from agicli.printplus import pr, inp
from agicli.welcome import get_detail, save_detail
from agicli.ai import ask
import os
import re
from agicli.tools import agi_tool, get_tools_prompt, handle_custom_tool, TOOL_REGISTRY
from agicli.agi_armamentarium import search_armamentarium
import agicli.agi_armamentarium_expanded
DEBUG = False

def get_currrent_date():
    from datetime import datetime
    return datetime.now().strftime("%B %d, %Y")

def get_current_time():
    from datetime import datetime
    return datetime.now().strftime("%I:%M %p")

@agi_tool
def current_weather():
    """Returns the current weather in the user's location based on the season."""
    from datetime import datetime
    month = datetime.now().month
    if month in [12, 1, 2]:
        return "It's winter season with cold temperatures."
    elif month in [3, 4, 5]:
        return "It's spring season with mild temperatures."
    elif month in [6, 7, 8]:
        return "It's summer season with warm temperatures."
    else:
        return "It's autumn season with cool temperatures."

@agi_tool
def write(filename, content, mode="rewrite"):
    """Writes a file with the specified filename and content. supported modes are 'rewrite' and 'append'."""
    SANDBOX_PATH = os.path.join(os.getcwd(), "agicli", "sandbox")
    try:
        if not os.path.exists(SANDBOX_PATH):
            os.makedirs(SANDBOX_PATH)
        
        filepath = os.path.join(SANDBOX_PATH, filename)
        
        if mode == "append":
            with open(filepath, "a") as f:
                f.write(content)
            return f"Successfully appended to file: {filename}"
        else:
            with open(filepath, "w") as f:
                f.write(content)
            return f"Successfully wrote file: {filename} ({len(content)} bytes)"
    except Exception as e:
        return f"Error writing file: {str(e)}"

@agi_tool
def filetree():
    """Lists the current files in the workspace (sandbox)."""
    SANDBOX_PATH = os.path.join(os.getcwd(), "agicli", "sandbox")
    try:
        if not os.path.exists(SANDBOX_PATH):
            return "Filetree: Sandbox is empty (no files yet)"
        
        file_list = []
        for root, dirs, files in os.walk(SANDBOX_PATH):
            rel_root = os.path.relpath(root, SANDBOX_PATH)
            for f in files:
                if rel_root == ".":
                    file_list.append(f)
                else:
                    file_list.append(os.path.join(rel_root, f))
        
        if file_list:
            tree_output = "\n".join(f"  - {f}" for f in sorted(file_list))
            return f"Filetree:\n{tree_output}"
        else:
            return "Filetree: Sandbox is empty (no files yet)"
    except Exception as e:
        return f"Error reading filetree: {str(e)}"

@agi_tool
def read(filename):
    """Reads the content of the specified file and returns it."""
    SANDBOX_PATH = os.path.join(os.getcwd(), "agicli", "sandbox")
    try:
        filepath = os.path.join(SANDBOX_PATH, filename)
        if not os.path.exists(filepath):
            return f"Error: File not found: {filename}"
        else:
            with open(filepath, "r") as f:
                file_content = f.read()
            return f"Content of {filename}:\n{file_content}"
    except Exception as e:
        return f"Error reading file: {str(e)}"

def parse_tools(response):
    """
    Extracts tool calls from a response string and returns a list of dictionaries with name and attrs.
    """
    if "<tools>" not in response:
        return []
    
    tools_found = []
    tool_sections = re.findall(r'<tools>([\s\S]*?)</tools>', response)
    
    for section in tool_sections:
        # More robust tag matching: skip over quoted strings when looking for closing > or />
        tag_pattern = r'<(\w+)\s*((?:[^"\'/>]|"[^"]*"|\'[^\']*\')*?)\s*/?>'
        # Wait, the regex I used before was better: r'<(\w+)\s*((?:[^"\'/>]|"[^"]*"|\'[^\']*\')*?)\s*/?>'
        # Let's use the one that works.
        tag_pattern = r'<(\w+)\s*((?:[^"\'/>]|"[^"]*"|\'[^\']*\')*?)\s*/?>'
        tags = re.findall(tag_pattern, section)
        
        for tag_name, attr_string in tags:
            # Robust attribute extraction for both single and double quotes
            attrs = {}
            attr_pattern = r'(\w+)=((?:"[\s\S]*?(?<!\\)")|(?:\'[\s\S]*?(?<!\\)\'))'
            for attr_match in re.finditer(attr_pattern, attr_string):
                key = attr_match.group(1)
                val = attr_match.group(2)[1:-1] # strip outer quotes
                
                # Unescape basic escape sequences and common XML entities
                val = val.replace("\\n", "\n").replace("\\t", "\t").replace('\\"', '"').replace("\\'", "'")
                val = val.replace("&quot;", '"').replace("&apos;", "'").replace("&lt;", "<").replace("&gt;", ">").replace("&amp;", "&")
                attrs[key] = val
            tools_found.append({"name": tag_name, "attrs": attrs})
            
    return tools_found

def launch_agent():
    if len(str(get_detail("gemini_api_key"))) == 0:
        pr("Welcome to AGI CLI! Let's set up your Gemini API key.")
        api_key = inp("Please enter your Gemini API key: ")
        save_detail("gemini_api_key", api_key)
        pr("API key saved successfully.", color="green")

    pr("AGI Agent is online.", color="green")

    user_input = inp("What would you like the agent to do? ", color="red", margin=0)
    context = []
    context.append(user_input)

    system_prompt = f"""
You are Alpha, an AGI Agent developed by AG Corp, today is {get_currrent_date()} at {get_current_time()}.
The user is called {get_detail("user_name")}. They have tasked you with the following prompt: {user_input}
As an AI Agent you are not limited to a single response, you can think step by step and take actions to complete the task.
In any part of your response, include one or more tool calls using XML format wrapped in a <tools> block like this:
<tools>
  <tool_name arg1="value1" arg2="value2" />
</tools>

You now have access to the following CORE tools:
{get_tools_prompt()}
<preamble text="I will..." /> - you may use preamble text to say the next tool you will use
<done result="Provide a summary of your work and result here." /> - indicates you have completed the user's task, once you provide this tool the system will terminate you.

IMPORTANT: You have an "Armamentarium" of hundreds of other tools (calculators, stock prices, news, translations, etc.).
If you need a tool that is not listed above, you MUST call <search_armamentarium query="search term" /> to find and register it for use in subsequent turns.

Use these tools as needed to complete the user's task.
When using tool attributes, ensure you escape double quotes with a backslash (\") or use single quotes if the content contains double quotes.
You MUST call the <done /> tool within a <tools> block when you have completed the user's task.
If you call a tool, ALWAYS wait for the tool call result before calling the <done /> tool.
Never call the <done> tool your first turn when you call multiple tools, wait for the system to get back to you with the tool call results before calling the <done> tool.
The armamentarium has many tools like design guidelines, user timezone, random number generator, etc.
"""
    not_complete = True

    context.append(system_prompt)
    i = 0
    max_iterations = 100  # Prevent infinite loops

    # SANDBOX_PATH = os.path.join(os.getcwd(), "agicli", "sandbox") # This is now defined within each tool function
    
    while not_complete and i < max_iterations:
        i += 1

        if DEBUG:
            pr(f"\n--- Agent Iteration {i} ---\n", color="magenta")
            pr(f"Current Context:\n" + "\n".join(context), color="yellow")
        
        response = ask("\n".join(context))
        pr(response, color="cyan")
        
        if response.startswith("Error:"):
            pr("API error encountered. Stopping agent.", color="red")
            break
        
        context.append(f"Assistant: {response}")
        
        tools = parse_tools(response)
        if tools:
            results = []
            for tool in tools:
                tag_name = tool["name"]
                attrs = tool["attrs"]
                
                if tag_name == "done":
                    pr("Task completed by the agent.", color="green")
                    summary = attrs.get("result", "Task completed.")
                    pr(f"Summary of work: {summary}", color="blue")
                    results.append(f"Summary of work: {summary}")
                    not_complete = False
                    break
                
                elif tag_name == "preamble":
                    preamble_text = attrs.get("text", "")
                    if preamble_text:
                        pr(f"🔮 {preamble_text}", color="magenta")
                        results.append(f"Preamble acknowledged: {preamble_text}")
                    else:
                        pr("🔮 Agent is thinking...", color="magenta")
                        results.append("Preamble acknowledged.")
                
                else:
                    # Check dynamic registry
                    result = handle_custom_tool(tag_name, attrs)
                    if result is not None:
                        # Print relevant feedback for core tools
                        if tag_name == "current_weather":
                            pr(f"Current weather: {result}", color="blue")
                        elif tag_name == "write":
                            pr(f"📝 {result}", color="green")
                        elif tag_name == "filetree":
                            pr(f"📂 {result}", color="green")
                        elif tag_name == "read":
                            pr(f"📖 Read file: {attrs.get('filename')}", color="green")
                        else:
                            pr(f"🛠 Executed tool: {tag_name}", color="yellow")
                            
                        results.append(result)
                    else:
                        results.append(f"Error: Tool '{tag_name}' not found.")

            if results:
                context.append(f"Tool results:\n" + "\n".join(results))
            
            if not not_complete:
                break
    
    if i >= max_iterations:
        pr("Max iterations reached. Stopping agent.", color="red")