"""
doit-fm AI engine.
Supports: NVIDIA NIM, OpenAI, Anthropic, Groq, xAI, Ollama.
All OpenAI-compatible providers use the same /v1/chat/completions path.
Anthropic uses /v1/messages with its own format.
"""
from __future__ import annotations
import asyncio, json, logging, re
from datetime import datetime
from typing import Optional

from doit_fm.config import AI_PROVIDERS, HTTP_TIMEOUT

logger = logging.getLogger("doit.ai")


def _system_prompt() -> str:
    now      = datetime.now().strftime("%A %B %d %Y, %I:%M %p")
    plat     = __import__("platform").system()
    from doit_fm.tools import tool_list_for_prompt
    tools = tool_list_for_prompt()
    return f"""You are Doit — a friendly AI assistant that controls a real computer for the user via Telegram.
Today: {now}. OS: {plat}.

PERSONALITY
• Warm, direct. Many users are non-technical — use plain language.
• Prefer doing over asking. Infer reasonable defaults.
• Only respond with "clarify" when you genuinely cannot proceed.
• Users WILL misspell and use casual language — always infer their real intent.
• Treat typos/mistakes generously: "delet that folder" = delete_item, "screenshoot" = screenshot, etc.

AVAILABLE TOOLS
{tools}

RESPONSE FORMAT — output ONLY valid JSON. No prose, no markdown, just JSON.

Single tool call:
{{"tool": "tool_name", "args": {{}}, "description": "Plain English what I'm doing", "confirm": false}}

Dangerous action (shell_run, delete_item) — MUST include confirm:true:
{{"tool": "tool_name", "args": {{}}, "description": "Plain English what I'm doing", "confirm": true}}

Need clarification:
{{"tool": "clarify", "args": {{"message": "Short friendly question"}}, "description": "clarify", "confirm": false}}

INTENT MAPPING RULES (always infer, never reject)
• "health" / "how's my computer" / "cpu" / "ram" → sys_health
• "run X" / "execute X" / "terminal X" / "bash X" → shell_run command=X, confirm=true
• "volume up/down/louder/quieter" → vol_set direction=up/down
• "set volume to N" / "volume N" → vol_set level=N
• "what's the volume" / "how loud" → vol_get
• "screenshot" / "screenshoot" / "capture screen" → screenshot
• "open X" / "launch X" / "start X" → app_open name=X
• "delete X" / "remove X" / "uninstall X" → delete_item target=X, kind inferred from context
• "export" / "history" / "what did you do" → export_log
• "show files" / "list X" / "ls X" → fs_list path=X
• "read X" / "show contents of X" → fs_read path=X
• "create file X" / "write to X" → fs_write
• "remember X" / "save that X" → memory_store
• "what do you remember" / "recall X" → memory_recall
• "search for X" / "google X" → web_search query=X
• "weather in X" / "temperature X" → weather city=X
• "calculate X" / "what is X" (math) → calc expression=X
• "remind me in X minutes" → reminder_set minutes=X
"""


class ConversationMemory:
    def __init__(self, max_turns: int = 10):
        self._turns = []
        self._max   = max_turns * 2

    def add(self, role: str, content: str):
        self._turns.append({"role": role, "content": content})
        if len(self._turns) > self._max:
            self._turns = self._turns[-self._max:]

    def messages(self, user_msg: str) -> list:
        return list(self._turns) + [{"role": "user", "content": user_msg}]

    def clear(self):
        self._turns = []


class AIEngine:
    def __init__(self, config: dict):
        self.provider   = config.get("ai_provider", "nvidia")
        self.model      = config.get("ai_model", "meta/llama-4-maverick-17b-128e-instruct")
        self.api_key    = config.get("ai_api_key", "")
        self.base_url   = config.get("ai_base_url", "")
        self._memory    = ConversationMemory()
        self._available = True

    def clear_memory(self):
        self._memory.clear()

    def _is_anthropic(self) -> bool:
        return self.provider == "anthropic"

    def _headers(self) -> dict:
        if self._is_anthropic():
            return {"x-api-key": self.api_key,
                    "anthropic-version": "2023-06-01",
                    "Content-Type": "application/json"}
        return {"Authorization": f"Bearer {self.api_key}",
                "Content-Type": "application/json"}

    def _endpoint(self) -> str:
        if self._is_anthropic():
            return f"{self.base_url}/messages"
        return f"{self.base_url}/chat/completions"

    def _payload(self, messages: list) -> dict:
        sys_prompt = _system_prompt()
        if self._is_anthropic():
            return {"model": self.model, "max_tokens": 1024,
                    "system": sys_prompt, "messages": messages}
        full = [{"role": "system", "content": sys_prompt}] + messages
        p = {"model": self.model, "messages": full,
             "max_tokens": 1024, "temperature": 0.1}
        # json_object mode — supported by OpenAI, NVIDIA, and most compatible providers
        if self.provider in ("openai", "nvidia", "groq", "xai"):
            p["response_format"] = {"type": "json_object"}
        return p

    def _extract(self, data: dict) -> str:
        if self._is_anthropic():
            for b in data.get("content", []):
                if isinstance(b, dict) and b.get("type") == "text":
                    return b["text"]
            raise ValueError("No text in Anthropic response")
        choices = data.get("choices", [])
        if not choices:
            raise ValueError(f"Empty choices. Response: {str(data)[:300]}")
        return choices[0]["message"]["content"] or ""

    def _parse(self, raw: str) -> dict:
        text = raw.strip()
        text = re.sub(r"^```(?:json)?\s*\n?", "", text, flags=re.MULTILINE)
        text = re.sub(r"\n?```\s*$", "", text, flags=re.MULTILINE)
        text = text.strip()
        m = re.search(r"\{.*\}", text, re.DOTALL)
        if m:
            text = m.group()
        try:
            a = json.loads(text)
            if not isinstance(a, dict):
                raise ValueError("not a dict")
            # Normalise key names
            for alt in ("action", "command", "function"):
                if "tool" not in a and alt in a:
                    a["tool"] = a.pop(alt)
            if "tool" not in a:
                raise ValueError("no tool key")
            a.setdefault("args", {})
            a.setdefault("description", a.get("tool", ""))
            a.setdefault("confirm", False)
            return a
        except Exception as e:
            logger.warning("AI parse failed: %s | raw: %.200s", e, raw)
            return {
                "tool": "clarify",
                "args": {"message": "I didn't quite catch that — could you rephrase?"},
                "description": "clarify",
                "confirm": False,
            }

    def _http_call(self, endpoint: str, headers: dict, payload: dict) -> dict:
        """Blocking HTTP POST — runs in executor thread."""
        import urllib.request, urllib.error
        data = json.dumps(payload).encode()
        req  = urllib.request.Request(endpoint, data=data, headers=headers)
        try:
            with urllib.request.urlopen(req, timeout=HTTP_TIMEOUT) as r:
                return json.loads(r.read())
        except urllib.error.HTTPError as e:
            body = {}
            try: body = json.loads(e.read())
            except Exception: pass
            code = e.code
            desc = (body.get("error", {}).get("message")
                    or body.get("error")
                    or body.get("message")
                    or str(e))
            if code == 401: raise _InvalidKey(str(desc))
            if code == 429: raise _Quota(str(desc))
            raise _AIError(f"HTTP {code}: {desc}")
        except Exception as e:
            raise _AIError(str(e))

    async def parse_intent(self, user_message: str) -> dict:
        messages = self._memory.messages(user_message)
        loop     = asyncio.get_event_loop()
        try:
            data   = await loop.run_in_executor(
                None, lambda: self._http_call(self._endpoint(), self._headers(), self._payload(messages)))
            text   = self._extract(data)
            action = self._parse(text)
            self._available = True
            self._memory.add("user", user_message)
            self._memory.add("assistant", json.dumps(action))
            return action
        except _InvalidKey as e:
            self._available = False
            return _err(f"Invalid API key for {self.provider}. Run `doitfm init --reconfigure` to update it.")
        except _Quota:
            return _err("Rate limit hit — please wait a moment and try again.")
        except Exception as e:
            logger.error("AI call failed: %s", e)
            self._available = False
            return _err(f"AI unreachable ({self.provider}). Check internet/API key. Error: {str(e)[:100]}")


class _AIError(Exception): pass
class _InvalidKey(_AIError): pass
class _Quota(_AIError): pass

def _err(msg: str) -> dict:
    return {"tool": "error", "args": {"message": msg}, "description": "error", "confirm": False}


async def validate_api_key(provider: str, model: str, api_key: str, base_url: str) -> tuple[bool, str]:
    """Quick live test. Returns (ok, message)."""
    engine = AIEngine({"ai_provider": provider, "ai_model": model,
                       "ai_api_key": api_key, "ai_base_url": base_url})
    loop = asyncio.get_event_loop()
    try:
        payload = engine._payload([{
            "role": "user",
            "content": 'Reply ONLY with valid JSON: {"tool":"ok","args":{},"description":"ok","confirm":false}',
        }])
        data   = await loop.run_in_executor(
            None, lambda: engine._http_call(engine._endpoint(), engine._headers(), payload))
        text   = engine._extract(data)
        action = engine._parse(text)
        if action.get("tool") in ("ok", "clarify", "error", "fs_list", "sys_health"):
            return True, f"✓ Connected: {provider} / {model}"
        return True, f"✓ Connected (response: {text[:60]})"
    except _InvalidKey:
        return False, "Invalid API key — check the key and try again"
    except _Quota:
        return False, "Quota exceeded on this key"
    except Exception as e:
        return False, str(e)
