import os
import time

from litellm import completion

from .terminal import SPECIAL_KEYS

PROVIDERS = [
    ("ANTHROPIC_API_KEY", "claude-3-5-haiku-20241022"),
    ("OPENAI_API_KEY", "gpt-4o-mini"),
    ("GOOGLE_API_KEY", "gemini/gemini-2.5-flash"),
    ("GROQ_API_KEY", "groq/llama-3.3-70b-versatile"),
    ("OPENROUTER_API_KEY", "anthropic/claude-3.5-haiku"),
]

# timings:
#   claude-3-5-haiku: 1.06
#   gemini-flash-2.5: 0.62
#   gpt-4o-mini: 2.11
#   roq/llama-3.3-70b-versatile: 0.27


class AIAssist:
    def __init__(self):
        self.api_key, self.model = self._get_provider()
        self.special_keys_list = ", ".join(f"{{{k}}}" for k in SPECIAL_KEYS)

    @staticmethod
    def _get_provider():
        for env_key, model in PROVIDERS:
            if api_key := os.environ.get(env_key):
                return api_key, model

        raise ValueError(
            "No LLM API key found. Set one of: "
            + ", ".join(env_key for env_key, _ in PROVIDERS)
        )

    def get_command(
        self,
        screen: str,
        current_command: str,
        user_request: str,
        current_process: str | None,
    ) -> str:
        user_prompt = [
            "You help users in their terminal by typing the right commands/keystrokes.",
            "",
            "SCREEN OUTPUT:",
            screen,
            "",
        ]

        if current_process:
            user_prompt.extend(
                [
                    f"CURRENT APP: {current_process}",
                    "",
                    f"SPECIAL KEYS: {self.special_keys_list}",
                    "Use {{} for literal brace",
                    "",
                    "Use keyboard navigation to interact with this app.",
                    "",
                    "EXAMPLES:",
                    "  Exit vim: {esc}:wq{enter}",
                    "  Navigate: {left}{left}{down}",
                    "  Search in less: /pattern{enter}",
                ]
            )
        elif current_command:
            user_prompt.extend(
                [
                    f"PARTIAL COMMAND: {current_command}",
                    "",
                    f"Special keys available: {self.special_keys_list}",
                    "Use {{} for literal brace",
                    "",
                    "EXAMPLES:",
                    "  If typed 'ffmpeg' and user wants to convert output.mov:",
                    '    -i output.mov -filter:v "setpts=0.5*PTS" output.mp4',
                    "  If typed 'git comit':",
                    "    {left}{left}{left}{left}{left}m",
                ]
            )
        else:
            user_prompt.extend(
                [
                    "Provide the complete command to run.",
                    "",
                    "EXAMPLES:",
                    "  Fix 'gti status' error: git status",
                    "  List python files: find . -name '*.py'",
                    "  Convert video: ffmpeg -i input.mov output.mp4",
                ]
            )

        user_prompt.extend(
            [
                "",
                f"USER REQUEST: {user_request}",
                "",
                "Reply with ONLY the characters/keys to type (no explanation):",
            ]
        )

        start = time.time()

        response = completion(
            model=self.model,
            messages=[{"role": "user", "content": "\n".join(user_prompt)}],
            api_key=self.api_key,
            timeout=10.0,
        )

        res = response.choices[0].message.content
        print(f"{time.time()-start}: {res}")
        return res
