import os
import json
import requests
from openai import OpenAI
from anthropic import Anthropic


class LLM:
    TOOLS = [
        {
            "type": "function",
            "function": {
                "name": "get_system_info",
                "description": "Get system information like OS, architecture, hostname",
                "parameters": {"type": "object", "properties": {}, "required": []},
            },
        },
        {
            "type": "function",
            "function": {
                "name": "get_cpu_usage",
                "description": "Get CPU usage percentage and core count",
                "parameters": {"type": "object", "properties": {}, "required": []},
            },
        },
        {
            "type": "function",
            "function": {
                "name": "get_memory_usage",
                "description": "Get RAM usage and total memory",
                "parameters": {"type": "object", "properties": {}, "required": []},
            },
        },
        {
            "type": "function",
            "function": {
                "name": "get_disk_usage",
                "description": "Get disk space information",
                "parameters": {"type": "object", "properties": {}, "required": []},
            },
        },
        {
            "type": "function",
            "function": {
                "name": "get_battery_status",
                "description": "Get battery status and charge level",
                "parameters": {"type": "object", "properties": {}, "required": []},
            },
        },
        {
            "type": "function",
            "function": {
                "name": "get_network_status",
                "description": "Get network statistics",
                "parameters": {"type": "object", "properties": {}, "required": []},
            },
        },
        {
            "type": "function",
            "function": {
                "name": "get_running_processes",
                "description": "Get list of running processes",
                "parameters": {
                    "type": "object",
                    "properties": {"limit": {"type": "integer", "default": 10}},
                    "required": [],
                },
            },
        },
        {
            "type": "function",
            "function": {
                "name": "take_screenshot",
                "description": "Take a screenshot of the screen",
                "parameters": {"type": "object", "properties": {}, "required": []},
            },
        },
        {
            "type": "function",
            "function": {
                "name": "open_application",
                "description": "Open an application or website",
                "parameters": {
                    "type": "object",
                    "properties": {
                        "app_name": {
                            "type": "string",
                            "description": "Name of app or URL to open",
                        }
                    },
                    "required": ["app_name"],
                },
            },
        },
        {
            "type": "function",
            "function": {
                "name": "shutdown_computer",
                "description": "Shutdown the computer",
                "parameters": {"type": "object", "properties": {}, "required": []},
            },
        },
        {
            "type": "function",
            "function": {
                "name": "restart_computer",
                "description": "Restart the computer",
                "parameters": {"type": "object", "properties": {}, "required": []},
            },
        },
        {
            "type": "function",
            "function": {
                "name": "lock_screen",
                "description": "Lock the computer screen",
                "parameters": {"type": "object", "properties": {}, "required": []},
            },
        },
        {
            "type": "function",
            "function": {
                "name": "scrape_website",
                "description": "Scrape content from a website",
                "parameters": {
                    "type": "object",
                    "properties": {
                        "url": {"type": "string", "description": "URL to scrape"},
                        "selector": {
                            "type": "string",
                            "description": "CSS selector to extract specific content",
                        },
                    },
                    "required": ["url"],
                },
            },
        },
        {
            "type": "function",
            "function": {
                "name": "browser_open",
                "description": "Open a website in browser using Selenium",
                "parameters": {
                    "type": "object",
                    "properties": {
                        "url": {"type": "string", "description": "URL to open"}
                    },
                    "required": ["url"],
                },
            },
        },
        {
            "type": "function",
            "function": {
                "name": "desktop_click",
                "description": "Click at coordinates on screen using PyAutoGUI",
                "parameters": {
                    "type": "object",
                    "properties": {"x": {"type": "integer"}, "y": {"type": "integer"}},
                    "required": ["x", "y"],
                },
            },
        },
        {
            "type": "function",
            "function": {
                "name": "desktop_type",
                "description": "Type text using PyAutoGUI",
                "parameters": {
                    "type": "object",
                    "properties": {"text": {"type": "string"}},
                    "required": ["text"],
                },
            },
        },
        {
            "type": "function",
            "function": {
                "name": "list_directory",
                "description": "List files and folders in a directory. Works on Linux, macOS, and Windows",
                "parameters": {
                    "type": "object",
                    "properties": {
                        "path": {
                            "type": "string",
                            "description": "Directory path to list (default: current directory)",
                        },
                        "show_hidden": {
                            "type": "boolean",
                            "description": "Show hidden files (default: False)",
                        },
                    },
                },
            },
        },
        {
            "type": "function",
            "function": {
                "name": "read_file",
                "description": "Read content of a text file. Works on Linux, macOS, and Windows",
                "parameters": {
                    "type": "object",
                    "properties": {
                        "path": {
                            "type": "string",
                            "description": "Full path to the file",
                        },
                        "max_lines": {
                            "type": "integer",
                            "description": "Maximum lines to read (default: 100)",
                        },
                    },
                    "required": ["path"],
                },
            },
        },
        {
            "type": "function",
            "function": {
                "name": "write_file",
                "description": "Write content to a file. Creates new file or overwrites. Works on Linux, macOS, and Windows",
                "parameters": {
                    "type": "object",
                    "properties": {
                        "path": {
                            "type": "string",
                            "description": "Full path for the file",
                        },
                        "content": {
                            "type": "string",
                            "description": "Content to write to file",
                        },
                    },
                    "required": ["path", "content"],
                },
            },
        },
        {
            "type": "function",
            "function": {
                "name": "create_folder",
                "description": "Create a new directory/folder. Works on Linux, macOS, and Windows",
                "parameters": {
                    "type": "object",
                    "properties": {
                        "path": {
                            "type": "string",
                            "description": "Path for new directory",
                        },
                    },
                    "required": ["path"],
                },
            },
        },
        {
            "type": "function",
            "function": {
                "name": "delete_file",
                "description": "Delete a file. Works on Linux, macOS, and Windows",
                "parameters": {
                    "type": "object",
                    "properties": {
                        "path": {
                            "type": "string",
                            "description": "Path to file to delete",
                        },
                    },
                    "required": ["path"],
                },
            },
        },
        {
            "type": "function",
            "function": {
                "name": "delete_folder",
                "description": "Delete an empty directory. Works on Linux, macOS, and Windows",
                "parameters": {
                    "type": "object",
                    "properties": {
                        "path": {
                            "type": "string",
                            "description": "Path to directory to delete",
                        },
                    },
                    "required": ["path"],
                },
            },
        },
        {
            "type": "function",
            "function": {
                "name": "get_file_info",
                "description": "Get information about a file or folder (size, modified date, type). Works on Linux, macOS, and Windows",
                "parameters": {
                    "type": "object",
                    "properties": {
                        "path": {
                            "type": "string",
                            "description": "Path to file or folder",
                        },
                    },
                    "required": ["path"],
                },
            },
        },
        {
            "type": "function",
            "function": {
                "name": "get_home_directory",
                "description": "Get the user's home directory path. Works on Linux, macOS, and Windows",
                "parameters": {"type": "object", "properties": {}, "required": []},
            },
        },
        {
            "type": "function",
            "function": {
                "name": "get_current_directory",
                "description": "Get the current working directory. Works on Linux, macOS, and Windows",
                "parameters": {"type": "object", "properties": {}, "required": []},
            },
        },
        {
            "type": "function",
            "function": {
                "name": "change_directory",
                "description": "Change the current working directory. Works on Linux, macOS, and Windows",
                "parameters": {
                    "type": "object",
                    "properties": {
                        "path": {
                            "type": "string",
                            "description": "Directory path to change to",
                        },
                    },
                    "required": ["path"],
                },
            },
        },
        {
            "type": "function",
            "function": {
                "name": "search_files",
                "description": "Search for files by name pattern. Works on Linux, macOS, and Windows",
                "parameters": {
                    "type": "object",
                    "properties": {
                        "directory": {
                            "type": "string",
                            "description": "Directory to search in",
                        },
                        "pattern": {
                            "type": "string",
                            "description": "File name pattern (e.g., '*.py', '*.txt')",
                        },
                    },
                    "required": ["directory", "pattern"],
                },
            },
        },
        {
            "type": "function",
            "function": {
                "name": "execute_command",
                "description": "Execute a shell command and return output. Works on Linux, macOS, and Windows",
                "parameters": {
                    "type": "object",
                    "properties": {
                        "command": {
                            "type": "string",
                            "description": "Shell command to execute",
                        },
                    },
                    "required": ["command"],
                },
            },
        },
        {
            "type": "function",
            "function": {
                "name": "copy_file",
                "description": "Copy a file from source to destination. Works on Linux, macOS, and Windows",
                "parameters": {
                    "type": "object",
                    "properties": {
                        "source": {"type": "string", "description": "Source file path"},
                        "destination": {
                            "type": "string",
                            "description": "Destination file path",
                        },
                    },
                    "required": ["source", "destination"],
                },
            },
        },
        {
            "type": "function",
            "function": {
                "name": "move_file",
                "description": "Move a file from source to destination. Works on Linux, macOS, and Windows",
                "parameters": {
                    "type": "object",
                    "properties": {
                        "source": {"type": "string", "description": "Source file path"},
                        "destination": {
                            "type": "string",
                            "description": "Destination file path",
                        },
                    },
                    "required": ["source", "destination"],
                },
            },
        },
        {
            "type": "function",
            "function": {
                "name": "run_python",
                "description": "Execute Python code and return output",
                "parameters": {
                    "type": "object",
                    "properties": {
                        "code": {
                            "type": "string",
                            "description": "Python code to execute",
                        },
                    },
                    "required": ["code"],
                },
            },
        },
        {
            "type": "function",
            "function": {
                "name": "append_file",
                "description": "Append content to a file. Works on Linux, macOS, and Windows",
                "parameters": {
                    "type": "object",
                    "properties": {
                        "path": {
                            "type": "string",
                            "description": "File path to append to",
                        },
                        "content": {
                            "type": "string",
                            "description": "Content to append",
                        },
                    },
                    "required": ["path", "content"],
                },
            },
        },
    ]

    def __init__(self, provider="ollama"):
        self.provider = provider
        if provider == "openai":
            self.client = OpenAI(api_key=os.getenv("OPENAI_API_KEY"))
            self.model = "gpt-3.5-turbo"
        elif provider == "anthropic":
            self.client = Anthropic(api_key=os.getenv("ANTHROPIC_API_KEY"))
            self.model = "claude-3-5-sonnet-20241014"
        elif provider == "ollama":
            self.base_url = "http://localhost:11434"
            self.model = "llama3.2"
        elif provider == "groq":
            self.client = OpenAI(
                api_key=os.getenv("GROQ_API_KEY"),
                base_url="https://api.groq.com/openai/v1",
            )
            self.model = "llama-3.1-8b-instant"
        elif provider == "gemini":
            self.base_url = "https://generativelanguage.googleapis.com/v1beta/models"
            self.model = "gemini-2.0-flash"
            self.api_key = os.getenv("GEMINI_API_KEY")
        elif provider == "deepseek":
            self.client = OpenAI(
                api_key=os.getenv("DEEPSEEK_API_KEY"),
                base_url="https://api.deepseek.com/v1",
            )
            self.model = "deepseek-chat"

    def ask(
        self, prompt, system_prompt="You are Ana, an advanced AI assistant.", tools=None
    ):
        if tools is None:
            tools = self.TOOLS

        if self.provider == "openai":
            response = self.client.chat.completions.create(
                model=self.model,
                messages=[
                    {"role": "system", "content": system_prompt},
                    {"role": "user", "content": prompt},
                ],
                tools=tools,
                tool_choice="auto",
            )
            return response.choices[0].message.content
        elif self.provider == "anthropic":
            response = self.client.messages.create(
                model=self.model,
                max_tokens=1024,
                system=system_prompt,
                messages=[{"role": "user", "content": prompt}],
            )
            return response.content[0].text
        elif self.provider == "ollama":
            response = requests.post(
                f"{self.base_url}/api/chat",
                json={
                    "model": self.model,
                    "messages": [
                        {"role": "system", "content": system_prompt},
                        {"role": "user", "content": prompt},
                    ],
                },
            )
            return response.json()["message"]["content"]
        elif self.provider == "groq":
            args = {
                "model": self.model,
                "messages": [
                    {"role": "system", "content": system_prompt},
                    {"role": "user", "content": prompt},
                ],
            }
            if tools:
                args["tools"] = tools
                args["tool_choice"] = "auto"
            response = self.client.chat.completions.create(**args)
            return response.choices[0].message.content
        elif self.provider == "gemini":
            url = f"{self.base_url}/{self.model}:generateContent?key={self.api_key}"
            data = {"contents": [{"parts": [{"text": f"{system_prompt}\n\n{prompt}"}]}]}
            response = requests.post(url, json=data)
            return response.json()["candidates"][0]["content"]["parts"][0]["text"]
        elif self.provider == "deepseek":
            args = {
                "model": self.model,
                "messages": [
                    {"role": "system", "content": system_prompt},
                    {"role": "user", "content": prompt},
                ],
            }
            if tools:
                args["tools"] = tools
                args["tool_choice"] = "auto"
            response = self.client.chat.completions.create(**args)
            return response.choices[0].message.content

    def ask_with_tools(self, prompt, system_prompt, tools, tool_callback):
        """Ask LLM with tools and execute them"""
        if self.provider in ["openai", "groq", "deepseek"]:
            return self._ask_openai_with_tools(
                prompt, system_prompt, tools, tool_callback
            )
        elif self.provider == "ollama":
            return self._ask_ollama_with_tools(
                prompt, system_prompt, tools, tool_callback
            )
        else:
            return self.ask(prompt, system_prompt)

    def _ask_ollama_with_tools(self, prompt, system_prompt, tools, tool_callback):
        """Ollama with tool calling"""
        response = requests.post(
            f"{self.base_url}/api/chat",
            json={
                "model": self.model,
                "messages": [
                    {
                        "role": "system",
                        "content": system_prompt
                        + "\n\nYou have access to these tools. Use them when needed.",
                    },
                    {"role": "user", "content": prompt},
                ],
                "tools": tools,
                "stream": False,
            },
        )

        result = response.json()

        if "tool_calls" in result.get("message", {}):
            tool_results = []
            for tool_call in result["message"]["tool_calls"]:
                tool_name = tool_call["function"]["name"]
                tool_args = tool_call["function"]["arguments"]

                try:
                    tool_result = tool_callback(tool_name, tool_args)
                    tool_results.append(f"Tool {tool_name}: {tool_result}")
                except Exception as e:
                    tool_results.append(f"Tool {tool_name} error: {str(e)}")

            final_response = requests.post(
                f"{self.base_url}/api/chat",
                json={
                    "model": self.model,
                    "messages": [
                        {"role": "system", "content": system_prompt},
                        {"role": "user", "content": prompt},
                        {"role": "assistant", "content": str(result["message"])},
                        {
                            "role": "user",
                            "content": "Tool results:\n" + "\n".join(tool_results),
                        },
                    ],
                    "stream": False,
                },
            )
            return final_response.json()["message"]["content"]

        return result["message"]["content"]

    def _ask_openai_with_tools(self, prompt, system_prompt, tools, tool_callback):
        """OpenAI/Groq with tool calling"""
        force_prompt = (
            system_prompt
            + "\n\nCRITICAL: You MUST call tools to answer questions. Do NOT write code. Call the appropriate tool now."
        )

        response = self.client.chat.completions.create(
            model=self.model,
            messages=[
                {"role": "system", "content": force_prompt},
                {"role": "user", "content": prompt},
            ],
            tools=tools,
            tool_choice="auto",
        )

        message = response.choices[0].message

        if message.tool_calls:
            tool_results = []
            for tool_call in message.tool_calls:
                tool_name = tool_call.function.name
                tool_args = json.loads(tool_call.function.arguments)

                try:
                    tool_result = tool_callback(tool_name, tool_args)
                    tool_results.append(f"Tool {tool_name} result: {tool_result}")
                except Exception as e:
                    tool_results.append(f"Tool {tool_name} error: {str(e)}")

            final_response = self.client.chat.completions.create(
                model=self.model,
                messages=[
                    {"role": "system", "content": system_prompt},
                    {"role": "user", "content": prompt},
                    message,
                    {
                        "role": "tool",
                        "content": "\n".join(tool_results),
                        "tool_call_id": message.tool_calls[0].id,
                    },
                ],
            )
            return final_response.choices[0].message.content

        return message.content
