import subprocess
from pathlib import Path
import os
import json
from dotenv import load_dotenv
from shutil import copyfile
from openai import OpenAI

load_dotenv()

# OpenRouter configuration
OPENROUTER_API_KEY = os.getenv("OPENROUTER_API_KEY")
MODEL = "qwen/qwen3-coder:free"  # Qwen Distilled Coder model on OpenRouter

if not OPENROUTER_API_KEY:
    raise ValueError("OPENROUTER_API_KEY must be set in .env file")

# Create OpenAI client pointed to OpenRouter
client = OpenAI(
    base_url="https://openrouter.ai/api/v1",
    api_key=OPENROUTER_API_KEY,
    default_headers={"HTTP-Referer": "http://localhost:5000"}  # Required by OpenRouter
)

# Define tools in OpenAI format
TOOLS = [
    {
        "type": "function",
        "function": {
            "name": "view",
            "description": "View the contents of a file or directory",
            "parameters": {
                "type": "object",
                "properties": {
                    "path": {
                        "type": "string",
                        "description": "Path to the file or directory"
                    }
                },
                "required": ["path"]
            }
        }
    },
    {
        "type": "function",
        "function": {
            "name": "create",
            "description": "Create a new file with the given content",
            "parameters": {
                "type": "object",
                "properties": {
                    "path": {
                        "type": "string",
                        "description": "Path to the file to create"
                    },
                    "file_text": {
                        "type": "string",
                        "description": "Content to write to the file"
                    }
                },
                "required": ["path", "file_text"]
            }
        }
    },
    {
        "type": "function",
        "function": {
            "name": "str_replace",
            "description": "Replace a string in a file with another string",
            "parameters": {
                "type": "object",
                "properties": {
                    "path": {
                        "type": "string",
                        "description": "Path to the file"
                    },
                    "old_str": {
                        "type": "string",
                        "description": "String to replace"
                    },
                    "new_str": {
                        "type": "string",
                        "description": "String to replace with"
                    }
                },
                "required": ["path", "old_str", "new_str"]
            }
        }
    },
    {
        "type": "function",
        "function": {
            "name": "bash",
            "description": "Execute a bash command",
            "parameters": {
                "type": "object",
                "properties": {
                    "command": {
                        "type": "string",
                        "description": "Command to execute"
                    }
                },
                "required": ["command"]
            }
        }
    }
]


def execute_tool(tool_name: str, tool_input: dict) -> dict:
    """Execute a tool and return structured result with error handling."""
    try:
        # string replace tools
        if tool_name == "view":
            path = Path(str(tool_input.get("path")))
            if path.is_file():
                content = path.read_text()
                return {"content": content, "is_error": False}
            elif path.is_dir():
                content = "\n".join(sorted([f.name for f in path.iterdir()]))
                return {"content": content, "is_error": False}
            else:
                return {"content": f"Error: {path} does not exist", "is_error": True}
        elif tool_name == "create":
            path = Path(str(tool_input.get("path")))
            content = str(tool_input.get("file_text"))
            if not content:
                return {
                    "content": "Error: No content provided in file_text",
                    "is_error": True,
                }
            path.parent.mkdir(parents=True, exist_ok=True)
            path.write_text(content)
            return {"content": f"File {path} written successfully", "is_error": False}
        elif tool_name == "str_replace":
            path = Path(str(tool_input.get("path")))
            old_str = str(tool_input.get("old_str"))
            new_str = str(tool_input.get("new_str"))

            if not path.exists():
                return {
                    "content": f"Error: File {path} does not exist",
                    "is_error": True,
                }

            content = path.read_text()
            if old_str not in content:
                return {
                    "content": f"Error: String '{old_str}' not found in {path}",
                    "is_error": True,
                }

            new_content = content.replace(old_str, new_str, 1)
            path.write_text(new_content)
            return {
                "content": f"Replaced '{old_str}' with '{new_str}' in {path}",
                "is_error": False,
            }
        # bash tools
        elif tool_name == "bash":
            command = tool_input.get("command")
            print(command)
            if not command:
                return {
                    "content": "Error: No input in command",
                    "is_error": True,
                }
            result = subprocess.run(
                command,
                shell=True,
                capture_output=True,
                text=True,
                timeout=30,  # Add timeout for safety
            )

            # Return both stdout and stderr, mark as error if non-zero exit code
            output = f"stdout: {result.stdout}\nstderr: {result.stderr}"
            return {"content": output, "is_error": result.returncode != 0}
        else:
            return {
                "content": f"Error: Unknown tool '{tool_name}'",
                "is_error": True,
            }
    except Exception as e:
        return {
            "content": f"Error executing {tool_name}: {str(e)}",
            "is_error": True,
        }


if __name__ == "__main__":
    prompt_content = Path("prompt.md").read_text()

    system_prompt = prompt_content[
        prompt_content.find("<role>") + 6 : prompt_content.find("</role>")
    ].strip()

    instructions_content = prompt_content[
        prompt_content.find("<thinking_process>") :
    ].strip()

    # Initialize conversation history
    messages = [
        {"role": "system", "content": system_prompt},
        {"role": "user", "content": instructions_content}
    ]

    while True:
        user_input = input("💬 User: ")
        messages.append({"role": "user", "content": user_input})

        while True:
            try:
                # Call the OpenRouter API with Qwen model
                response = client.chat.completions.create(
                    model=MODEL,
                    messages=messages,
                    temperature=0.2,
                    max_tokens=4096,
                    tools=TOOLS,
                    tool_choice="auto"
                )
            except Exception as e:
                print(f"Error calling OpenRouter API: {str(e)}")
                break

            assistant_message = response.choices[0].message
            assistant_content = assistant_message.content or ""
            
            # Check if the model wants to use tools
            if assistant_message.tool_calls:
                tool_results = []
                
                # Print any text content from the assistant
                if assistant_content:
                    print(assistant_content)
                
                print(f"Executing {len(assistant_message.tool_calls)} tool(s)...")
                
                # Process each tool call
                for tool_call in assistant_message.tool_calls:
                    function_name = tool_call.function.name
                    function_args = json.loads(tool_call.function.arguments)  # Parse JSON arguments safely
                    
                    print(f"Executing tool: {function_name}")
                    
                    # Execute the tool
                    result = execute_tool(function_name, function_args)
                    print(result["content"])
                    
                    # Add the tool result to the conversation
                    tool_results.append({
                        "tool_call_id": tool_call.id,
                        "role": "tool",
                        "name": function_name,
                        "content": result["content"]
                    })
                
                # Add the assistant's response to the conversation
                messages.append({
                    "role": "assistant",
                    "content": assistant_content,
                    "tool_calls": assistant_message.tool_calls
                })
                
                # Add all tool results to the conversation
                messages.extend(tool_results)
                
                continue
            else:
                # No tool calls, just print the response
                print(assistant_content)
                
                # Add the assistant's response to the conversation
                messages.append({
                    "role": "assistant",
                    "content": assistant_content
                })
                
                break  # Break out of inner loop to restart conversation