"""
neurocli models — multi-model router (Gemini + Groq, all free tiers)
Uses google-genai SDK (not google-generativeai)
"""
import os
import json
import re
from .session import track

TOOL_SCHEMA = [
    {
        "name": "read_file",
        "description": "Read a file from the filesystem.",
        "parameters": {
            "type": "object",
            "properties": {
                "path": {"type": "string", "description": "File path to read"}
            },
            "required": ["path"]
        }
    },
    {
        "name": "write_file",
        "description": "Write or create a file with the given content.",
        "parameters": {
            "type": "object",
            "properties": {
                "path": {"type": "string", "description": "File path to write"},
                "content": {"type": "string", "description": "Full file content to write"}
            },
            "required": ["path", "content"]
        }
    },
    {
        "name": "list_files",
        "description": "List files in a directory.",
        "parameters": {
            "type": "object",
            "properties": {
                "directory": {"type": "string", "description": "Directory to list (default: .)"}
            },
            "required": []
        }
    },
    {
        "name": "search_files",
        "description": "Search for a pattern across all files.",
        "parameters": {
            "type": "object",
            "properties": {
                "pattern": {"type": "string", "description": "Text pattern to search for"},
                "directory": {"type": "string", "description": "Directory to search in"}
            },
            "required": ["pattern"]
        }
    },
    {
        "name": "run_command",
        "description": "Execute a shell command.",
        "parameters": {
            "type": "object",
            "properties": {
                "cmd": {"type": "string", "description": "Shell command to execute"}
            },
            "required": ["cmd"]
        }
    },
    {
        "name": "git_commit",
        "description": "Stage all changes and git commit with a message.",
        "parameters": {
            "type": "object",
            "properties": {
                "message": {"type": "string", "description": "Commit message following conventional commits format"}
            },
            "required": ["message"]
        }
    },
]


def _route_model(prompt: str, cfg: dict) -> str:
    """Intelligently route to the right model. Groq first if key exists."""
    default = cfg.get("default_model", "auto")
    if default != "auto":
        return default
    if cfg.get("groq_api_key"):
        return "groq"
    prompt_lower = prompt.lower()
    if any(k in prompt_lower for k in ["refactor", "architect", "design", "rewrite", "migrate"]):
        return "gemini-pro"
    return "gemini-flash"


def call_ai(messages: list, system: str, cfg: dict, stream_cb=None) -> dict:
    """
    Call the appropriate AI model. Returns:
    { "text": str, "tool_calls": list, "model_used": str }
    """
    model_choice = _route_model(messages[-1]["content"] if messages else "", cfg)
    track("ai_requests")
    track("models_used", 1, sub_key=model_choice)

    if model_choice == "groq" and cfg.get("groq_api_key"):
        return _call_groq(messages, system, cfg, stream_cb)
    else:
        return _call_gemini(messages, system, cfg, model_choice, stream_cb)


def _call_gemini(messages: list, system: str, cfg: dict, model_choice: str, stream_cb=None) -> dict:
    """Call Gemini using google-genai SDK."""
    try:
        from google import genai
        from google.genai import types

        api_key = cfg.get("gemini_api_key") or os.environ.get("GEMINI_API_KEY", "")
        if not api_key:
            return {
                "text": "No Gemini API key.\nRun: neurocli config --gemini-key YOUR_KEY\nGet free key: https://aistudio.google.com/apikey",
                "tool_calls": [],
                "model_used": "none"
            }

        client = genai.Client(api_key=api_key)

        if model_choice == "gemini-pro":
            model_name = "gemini-2.5-pro-preview-03-25"
        elif model_choice == "gemini-flash":
            model_name = "gemini-2.0-flash"
        else:
            model_name = "gemini-2.0-flash-lite"

        tool_declarations = []
        for t in TOOL_SCHEMA:
            props = {}
            for pname, pinfo in t["parameters"].get("properties", {}).items():
                props[pname] = types.Schema(
                    type="STRING",
                    description=pinfo.get("description", "")
                )
            req = t["parameters"].get("required", [])
            tool_declarations.append(types.FunctionDeclaration(
                name=t["name"],
                description=t["description"],
                parameters=types.Schema(
                    type="OBJECT",
                    properties=props,
                    required=req
                )
            ))

        tools = [types.Tool(function_declarations=tool_declarations)]

        contents = []
        for msg in messages:
            role = "user" if msg["role"] == "user" else "model"
            contents.append(types.Content(
                role=role,
                parts=[types.Part(text=msg["content"])]
            ))

        config = types.GenerateContentConfig(
            system_instruction=system,
            tools=tools,
        )

        response = client.models.generate_content(
            model=model_name,
            contents=contents,
            config=config,
        )

        tool_calls = []
        text_parts = []

        candidate = response.candidates[0] if response.candidates else None
        if candidate and candidate.content and candidate.content.parts:
            for part in candidate.content.parts:
                if hasattr(part, "function_call") and part.function_call:
                    fc = part.function_call
                    args = dict(fc.args) if fc.args else {}
                    tool_calls.append({"name": fc.name, "args": args})
                elif hasattr(part, "text") and part.text:
                    text_parts.append(part.text)
                    if stream_cb:
                        stream_cb(part.text)

        return {
            "text": "\n".join(text_parts),
            "tool_calls": tool_calls,
            "model_used": model_name,
        }

    except Exception as e:
        return {
            "text": "Gemini error: " + str(e),
            "tool_calls": [],
            "model_used": "gemini-error"
        }


def _parse_tool_blocks(text: str) -> list:
    """Parse tool JSON blocks from LLM response."""
    tool_calls = []
    pattern = r"```tool\s*(\{.*?\})\s*```"
    matches = re.findall(pattern, text, re.DOTALL)
    for match in matches:
        try:
            obj = json.loads(match.strip())
            name = obj.get("name", "")
            args = obj.get("args", {})
            if name:
                tool_calls.append({"name": name, "args": args})
        except Exception:
            pass
    return tool_calls


def _call_groq(messages: list, system: str, cfg: dict, stream_cb=None) -> dict:
    """
    Groq with Llama — uses prompt-based tool calling (reliable).
    Native function calling is broken in Llama on Groq.
    """
    try:
        from groq import Groq

        client = Groq(api_key=cfg.get("groq_api_key") or os.environ.get("GROQ_API_KEY", ""))

        tools_desc = (
            "You have access to these tools. When you need to use a tool, "
            "output a JSON block exactly like this:\n\n"
            "```tool\n"
            '{"name": "tool_name", "args": {"arg1": "value1"}}\n'
            "```\n\n"
            "Available tools:\n"
            "- read_file(path): Read a file\n"
            "- write_file(path, content): Write or create a file\n"
            "- list_files(directory): List files in a directory\n"
            "- search_files(pattern, directory): Search for pattern in files\n"
            "- run_command(cmd): Execute a shell command\n"
            "- git_commit(message): Git add all and commit\n\n"
            "You can use multiple tools by writing multiple ```tool blocks.\n"
            "Always use tools when the user asks to create files, run commands, or edit code."
        )

        enhanced_system = system + "\n\n" + tools_desc

        groq_messages = [{"role": "system", "content": enhanced_system}] + messages

        response = client.chat.completions.create(
            model="llama-3.3-70b-versatile",
            messages=groq_messages,
            max_tokens=4096,
        )

        text = response.choices[0].message.content or ""

        if stream_cb and text:
            stream_cb(text)

        tool_calls = _parse_tool_blocks(text)
        clean_text = re.sub(r"```tool\s*\{.*?\}\s*```", "", text, flags=re.DOTALL).strip()

        return {
            "text": clean_text,
            "tool_calls": tool_calls,
            "model_used": "groq/llama-3.3-70b",
        }

    except Exception as e:
        return {
            "text": "Groq error: " + str(e),
            "tool_calls": [],
            "model_used": "groq-error"
        }
