"""
ai.py — AIEngine: sends user messages to the AI and parses JSON tool intents.

Supported provider families:
  • anthropic     — X-API-Key header, /messages endpoint
  • openai-compat — all others (OpenAI, Groq, Together, NVIDIA, Mistral,
                    Cohere, OpenRouter, Zhipu) → Bearer token, /chat/completions

Robustness features:
  • Exponential-backoff retry (transient 5xx / network errors)
  • Rate-limit (429) awareness
  • Nested JSON extraction (models that wrap in {response: {...}})
  • Unquoted JS-style keys, single quotes, trailing commas all repaired
  • Key alias normalisation: action/command/function → tool; parameters/arguments → args
  • danger tools always get confirm=True regardless of AI output
  • All errors fall back to {tool: clarify} — never raises to caller
"""
from __future__ import annotations

import asyncio
import datetime
import json
import platform
import re
import time
import urllib.error
import urllib.request
from typing import Any

from .config import (
    AI_PROVIDERS,
    CONVERSATION_MEMORY_SIZE,
    MAX_RETRIES,
    RETRY_BASE_DELAY,
    RETRY_MAX_DELAY,
)
from .logger import get_logger
from .tools import tool_list_for_prompt

log = get_logger("ai")


# ---------------------------------------------------------------------------
# Exceptions
# ---------------------------------------------------------------------------

class AIEngineError(Exception):
    """Base error for AIEngine."""

class APIKeyInvalid(AIEngineError):
    """Raised when the API key is definitively rejected (401/403)."""

class ProviderUnavailable(AIEngineError):
    """Raised when the provider is unreachable or returns a server error."""


# ---------------------------------------------------------------------------
# Conversation memory
# ---------------------------------------------------------------------------

class ConversationMemory:
    """Async-safe rolling window of the last N message turns."""

    def __init__(self, max_turns: int = CONVERSATION_MEMORY_SIZE) -> None:
        self._turns: list[dict] = []
        self._max = max(1, max_turns)
        self._lock = asyncio.Lock()

    async def add(self, role: str, content: str) -> None:
        async with self._lock:
            self._turns.append({"role": role, "content": str(content)})
            cap = self._max * 2
            if len(self._turns) > cap:
                self._turns = self._turns[-cap:]

    async def as_list(self) -> list[dict]:
        async with self._lock:
            return list(self._turns)

    async def clear(self) -> None:
        async with self._lock:
            self._turns.clear()

    def clear_sync(self) -> None:
        self._turns.clear()


# ---------------------------------------------------------------------------
# System prompt
# ---------------------------------------------------------------------------

_SYSTEM_PROMPT_TEMPLATE = """\
You are the BySalim AI agent running on the user's computer. The user controls \
their computer by sending you natural language messages via Telegram. Your job is \
to translate their request into the exact JSON tool call that will execute it.

You MUST respond with ONLY a single valid JSON object — no prose, no markdown \
fences, no explanation before or after. Just raw JSON.

Current date/time: {datetime}
Operating system: {os_name}

## RESPONSE SCHEMA (output ONLY this JSON object, nothing else):
{{
  "tool": "<tool_name>",
  "args": {{<key>: <value>}},
  "description": "<one sentence describing what you will do>",
  "confirm": <true|false>
}}

## RULES:
- "tool" MUST be exactly one of the tool names listed in the TOOLS section.
- "args" MUST be a flat dict with keys matching the tool's parameters. Use {{}} if no args needed.
- "confirm" MUST be true for tools marked DANGEROUS below. Use false for all others.
- If the request is unclear or you cannot map it to a tool, use tool "clarify".
- If a location/city/host/path/etc is mentioned, extract it and put it in args.
- NEVER output any text outside the JSON object. No "Here is my response:" preamble.

## TOOL SELECTION GUIDE:

System monitoring:
- "how's my computer/cpu/memory/ram/health/system stats" → system_health
- "uptime / how long running / boot time" → uptime  
- "battery / charge level" → battery
- "network stats / bandwidth / interfaces" → network_stats
- "what's running / processes / ps / top" → list_processes
- "kill/stop/terminate [process name or pid]" → kill_process {{name or pid}}
- "disk space / disk usage [path]" → disk_usage {{path:"/"}}

Shell:
- "run / execute / bash / shell / terminal: [command]" → run_shell {{command:"[cmd]"}} DANGEROUS

Files:
- "list / ls / show files [in path]" → list_directory {{path:"~"}}
- "read / show / cat / open [file]" → read_file {{path:"file"}}
- "write / save / create file [file] with [content]" → write_file {{path,content}} DANGEROUS
- "delete / remove / rm [file/dir]" → delete_file {{path}} DANGEROUS
- "copy [src] to [dst]" → copy_file {{src,dst}}
- "move / rename [src] to [dst]" → move_file {{src,dst}} DANGEROUS
- "find / search files [name] [in path]" → search_files {{query,path:"~"}}
- "create directory / mkdir [path]" → create_directory {{path}}
- "count words / word count [text or file]" → word_count {{text or path}}

Screen & audio:
- "screenshot / take a picture / capture / what's on screen" → screenshot
- "what's the volume / current volume" → get_volume
- "set volume to [N] / volume [N]%" → set_volume {{level:N}}
- "mute / silence" → mute_volume
- "clipboard / what's copied / paste content" → get_clipboard
- "copy [text] to clipboard" → set_clipboard {{text:"..."}}

Internet & apps:
- "weather in [city] / temperature in [city]" → weather {{city:"[city]"}}
- "ping [host] / is [host] reachable" → ping {{host:"[host]"}}
- "check internet / am I online / connectivity" → internet_speed
- "open / browse / go to [url]" → open_url {{url:"https://..."}}
- "launch / open / start [app name]" → open_application {{name:"[app]"}}

Time & reminders:
- "time / date / what day / current time" → datetime
- "remind me in [N] minutes/hours to [msg]" → reminder {{message:"[msg]",delay_minutes:N}}

Development:
- "python version / python info" → python_info
- "environment variables / env vars" → env_vars
- "git status [path]" → git_status {{path:"."}}
- "git log / commits [path]" → git_log {{path:"."}}
- "install [package] / pip install [pkg]" → pip_install {{package:"[pkg]"}} DANGEROUS
- "installed packages / pip list" → pip_list

Math:
- "calculate / compute / what is [math expression]" → calculate {{expression:"[expr]"}}

## AVAILABLE TOOLS:
{tool_list}

## EXAMPLES:
User: "take a screenshot"
You output: {{"tool":"screenshot","args":{{}},"description":"Take a screenshot of the screen","confirm":false}}

User: "weather in London"
You output: {{"tool":"weather","args":{{"city":"London"}},"description":"Get current weather in London","confirm":false}}

User: "remind me in 10 minutes to call mom"
You output: {{"tool":"reminder","args":{{"message":"call mom","delay_minutes":10}},"description":"Set a 10-minute reminder to call mom","confirm":false}}

User: "rm /tmp/test.txt"
You output: {{"tool":"delete_file","args":{{"path":"/tmp/test.txt"}},"description":"Delete /tmp/test.txt","confirm":true}}

User: "how's my cpu"
You output: {{"tool":"system_health","args":{{}},"description":"Check system health stats","confirm":false}}

User: "ping google.com"
You output: {{"tool":"ping","args":{{"host":"google.com"}},"description":"Ping google.com","confirm":false}}

User: "list files in /home"
You output: {{"tool":"list_directory","args":{{"path":"/home"}},"description":"List files in /home","confirm":false}}
"""


def _build_system_prompt() -> str:
    return _SYSTEM_PROMPT_TEMPLATE.format(
        datetime=datetime.datetime.now().strftime("%A %B %d %Y %H:%M:%S"),
        os_name=f"{platform.system()} {platform.release()} {platform.machine()}",
        tool_list=tool_list_for_prompt(),
    )


# ---------------------------------------------------------------------------
# Retry helper
# ---------------------------------------------------------------------------

def _with_retry(fn, max_retries: int = MAX_RETRIES):
    last_exc: Exception | None = None
    for attempt in range(max_retries):
        try:
            return fn()
        except APIKeyInvalid:
            raise
        except (ProviderUnavailable, Exception) as exc:
            last_exc = exc
            delay = min(RETRY_BASE_DELAY * (2 ** attempt), RETRY_MAX_DELAY)
            log.warning("API attempt %d/%d failed: %s — retry in %.1fs",
                        attempt + 1, max_retries, exc, delay)
            time.sleep(delay)
    raise ProviderUnavailable(f"All {max_retries} attempts failed. Last: {last_exc}")


# ---------------------------------------------------------------------------
# API key validation
# ---------------------------------------------------------------------------

def validate_api_key(provider: str, api_key: str, model: str) -> tuple[bool, str]:
    """Make a tiny real API call to validate the key. Returns (ok, message)."""
    if not api_key or not api_key.strip():
        return False, "API key cannot be empty"
    if provider not in AI_PROVIDERS:
        return False, f"Unknown provider: {provider}"

    info = AI_PROVIDERS[provider]
    headers: dict[str, str] = {"Content-Type": "application/json"}
    payload: dict[str, Any]

    if provider == "anthropic":
        headers["x-api-key"] = api_key.strip()
        headers["anthropic-version"] = "2023-06-01"
        url = f"{info['base_url']}/messages"
        payload = {"model": model, "max_tokens": 5,
                   "messages": [{"role": "user", "content": "hi"}]}
    else:
        headers["Authorization"] = f"Bearer {api_key.strip()}"
        if provider == "openrouter":
            headers["HTTP-Referer"] = "https://github.com/bysalim-agent/bysalim"
            headers["X-Title"] = "BySalim"
        url = f"{info['base_url']}/chat/completions"
        payload = {"model": model, "max_tokens": 5,
                   "messages": [{"role": "user", "content": "hi"}]}

    try:
        body = json.dumps(payload).encode("utf-8")
        req = urllib.request.Request(url, data=body, headers=headers, method="POST")
        with urllib.request.urlopen(req, timeout=20) as resp:
            resp.read()
        return True, "API key is valid ✓"
    except urllib.error.HTTPError as exc:
        raw = exc.read()
        try:
            err = json.loads(raw)
            msg = err.get("error", {}).get("message") or err.get("detail") or str(err)
        except Exception:
            msg = raw.decode(errors="replace")[:300]
        if exc.code in (401, 403):
            return False, f"Invalid API key: {msg}"
        if 400 <= exc.code < 500:
            return True, f"API key accepted (HTTP {exc.code})"
        return False, f"Provider error HTTP {exc.code}: {msg}"
    except urllib.error.URLError as exc:
        return False, f"Network error reaching {provider}: {exc.reason}"
    except Exception as exc:
        return False, f"Could not reach provider: {exc}"


# ---------------------------------------------------------------------------
# AIEngine
# ---------------------------------------------------------------------------

class AIEngine:
    """
    Wraps a single AI provider and parses user messages into tool call JSON.
    All HTTP calls run in asyncio's thread executor — never blocks the event loop.
    """

    def __init__(self, config: dict) -> None:
        self._provider: str = config["ai_provider"]
        self._model: str = config["ai_model"]
        self._api_key: str = config["ai_api_key"]
        info = AI_PROVIDERS.get(self._provider, {})
        self._base_url: str = config.get("ai_base_url") or info.get("base_url", "")
        self.memory = ConversationMemory()

    # ------------------------------------------------------------------
    # HTTP
    # ------------------------------------------------------------------

    def _build_request(self, messages: list[dict]) -> tuple[str, dict, bytes]:
        headers: dict[str, str] = {"Content-Type": "application/json"}
        payload: dict[str, Any]

        if self._provider == "anthropic":
            headers["x-api-key"] = self._api_key
            headers["anthropic-version"] = "2023-06-01"
            url = f"{self._base_url}/messages"
            user_msgs = [m for m in messages if m["role"] != "system"]
            system = next((m["content"] for m in messages if m["role"] == "system"), "")
            payload = {
                "model": self._model, "max_tokens": 512, "temperature": 0.0,
                "system": system, "messages": user_msgs,
            }
        else:
            headers["Authorization"] = f"Bearer {self._api_key}"
            if self._provider == "openrouter":
                headers["HTTP-Referer"] = "https://github.com/bysalim-agent/bysalim"
                headers["X-Title"] = "BySalim"
            url = f"{self._base_url}/chat/completions"
            payload = {
                "model": self._model, "max_tokens": 512, "temperature": 0.0,
                "messages": messages,
            }

        return url, headers, json.dumps(payload, ensure_ascii=False).encode("utf-8")

    def _extract_text(self, data: dict) -> str:
        if self._provider == "anthropic":
            for block in data.get("content", []):
                if isinstance(block, dict) and block.get("type") == "text":
                    return block.get("text", "")
            return ""
        else:
            choices = data.get("choices", [])
            if choices:
                return choices[0].get("message", {}).get("content", "") or ""
            return ""

    def _call_api_sync(self, messages: list[dict]) -> str:
        url, headers, body = self._build_request(messages)

        def _do() -> str:
            req = urllib.request.Request(url, data=body, headers=headers, method="POST")
            try:
                with urllib.request.urlopen(req, timeout=45) as resp:
                    raw = resp.read()
            except urllib.error.HTTPError as exc:
                raw_err = exc.read()
                try:
                    err = json.loads(raw_err)
                    msg = err.get("error", {}).get("message") or str(err)
                except Exception:
                    msg = raw_err.decode(errors="replace")[:500]
                if exc.code in (401, 403):
                    raise APIKeyInvalid(f"Key rejected by {self._provider}: {msg}")
                if exc.code == 429:
                    raise ProviderUnavailable(f"Rate limited: {msg}")
                if exc.code >= 500:
                    raise ProviderUnavailable(f"Server error {exc.code}: {msg}")
                raise ProviderUnavailable(f"HTTP {exc.code}: {msg}")
            except urllib.error.URLError as exc:
                raise ProviderUnavailable(f"Network error: {exc.reason}")

            try:
                data = json.loads(raw)
            except json.JSONDecodeError as exc:
                raise ProviderUnavailable(f"Invalid JSON from provider: {exc}")

            if "error" in data:
                msg = data["error"].get("message", str(data["error"]))
                raise ProviderUnavailable(f"Provider error: {msg}")

            return self._extract_text(data)

        return _with_retry(_do)

    # ------------------------------------------------------------------
    # JSON parsing — extremely robust
    # ------------------------------------------------------------------

    @staticmethod
    def _parse_json_response(text: str) -> dict:
        """
        Extract and normalise a JSON action object from any model output.

        Handles:
          - Markdown code fences
          - Prose before/after the JSON object
          - Models that wrap in {response: {...}} or {data: {...}}
          - Single-quoted strings
          - Unquoted JS-style object keys {key: "val"}
          - Trailing commas
          - Missing required fields
          - Key aliases: action/command/function→tool, parameters/arguments→args
        """
        _fallback: dict = {
            "tool": "clarify",
            "args": {"message": "I couldn't understand that. Please try rephrasing."},
            "description": "Clarify",
            "confirm": False,
        }

        if not text or not text.strip():
            return {**_fallback, "args": {"message": "Empty response from AI."}}

        # Strip markdown fences
        clean = re.sub(r"```(?:json)?\s*", "", text, flags=re.IGNORECASE).strip()
        clean = clean.replace("```", "").strip()

        def _try_parse(s: str) -> dict | None:
            try:
                obj = json.loads(s)
                return obj if isinstance(obj, dict) else None
            except json.JSONDecodeError:
                return None

        def _repair_and_parse(s: str) -> dict | None:
            fixed = s
            # Quote unquoted JS-style keys: {key: val} → {"key": val}
            fixed = re.sub(r'(?<=[{,])\s*([A-Za-z_][A-Za-z0-9_]*)\s*:', r' "\1":', fixed)
            # Replace single quotes with double (simple non-escaped case)
            fixed = re.sub(r"(?<![\\])'", '"', fixed)
            # Remove trailing commas
            fixed = re.sub(r",\s*([}\]])", r"\1", fixed)
            # Remove JS // comments
            fixed = re.sub(r"//[^\n]*", "", fixed)
            return _try_parse(fixed)

        def _find_outer_object(s: str) -> str | None:
            start = s.find("{")
            if start == -1:
                return None
            depth, in_str, esc = 0, False, False
            for i, ch in enumerate(s[start:], start=start):
                if esc:
                    esc = False; continue
                if ch == "\\" and in_str:
                    esc = True; continue
                if ch == '"' and not esc:
                    in_str = not in_str; continue
                if in_str:
                    continue
                if ch == "{":
                    depth += 1
                elif ch == "}":
                    depth -= 1
                    if depth == 0:
                        return s[start:i + 1]
            return s[start:]  # Unterminated — return what we have

        # Strategy 1: direct parse
        obj = _try_parse(clean) or _repair_and_parse(clean)

        # Strategy 2: extract outermost { } and parse
        if obj is None:
            outer = _find_outer_object(clean)
            if outer:
                obj = _try_parse(outer) or _repair_and_parse(outer)

        # Strategy 3: find any complete { } blob
        if obj is None:
            for m in re.finditer(r"\{[^{}]+\}", clean):
                obj = _try_parse(m.group()) or _repair_and_parse(m.group())
                if obj is not None:
                    break

        if obj is None or not isinstance(obj, dict):
            return {**_fallback, "args": {"message": "AI returned invalid JSON. Please try rephrasing."}}

        # Strategy 4: unwrap nested wrapper {response: {...}, data: {...}, etc.}
        _tool_keys = {"tool", "action", "command", "function"}
        if not (_tool_keys & set(obj.keys())):
            for v in obj.values():
                if isinstance(v, dict) and (_tool_keys & set(v.keys())):
                    obj = v
                    break

        # Normalise key aliases
        if "tool" not in obj:
            for alias in ("action", "command", "function", "type"):
                if alias in obj:
                    obj["tool"] = obj.pop(alias)
                    break
        if "args" not in obj:
            for alias in ("parameters", "arguments", "params", "input", "kwargs", "options"):
                if alias in obj:
                    obj["args"] = obj.pop(alias)
                    break

        # Fill missing fields
        obj.setdefault("tool", "clarify")
        obj.setdefault("args", {})
        obj.setdefault("description", "")
        obj.setdefault("confirm", False)

        # Type safety
        if not isinstance(obj["tool"], str) or not obj["tool"].strip():
            obj["tool"] = "clarify"
        else:
            obj["tool"] = obj["tool"].strip()

        if not isinstance(obj["args"], dict):
            if isinstance(obj["args"], str):
                try:
                    parsed = json.loads(obj["args"])
                    obj["args"] = parsed if isinstance(parsed, dict) else {}
                except Exception:
                    obj["args"] = {}
            else:
                obj["args"] = {}

        if isinstance(obj["confirm"], str):
            obj["confirm"] = obj["confirm"].lower() in ("true", "1", "yes")
        elif not isinstance(obj["confirm"], bool):
            obj["confirm"] = bool(obj["confirm"])

        return obj

    # ------------------------------------------------------------------
    # Public API
    # ------------------------------------------------------------------

    async def parse_intent(self, user_message: str) -> dict:
        """
        Send user_message to the AI with conversation history.
        Returns {tool, args, description, confirm}.
        Never raises — always returns a usable dict.
        """
        if not user_message or not user_message.strip():
            return {
                "tool": "clarify",
                "args": {"message": "Your message was empty. Please try again."},
                "description": "Empty message",
                "confirm": False,
            }

        system_prompt = _build_system_prompt()
        history = await self.memory.as_list()
        messages = [{"role": "system", "content": system_prompt}]
        messages.extend(history)
        messages.append({"role": "user", "content": user_message.strip()})

        loop = asyncio.get_event_loop()
        raw_text = ""
        try:
            raw_text = await loop.run_in_executor(None, self._call_api_sync, messages)
        except APIKeyInvalid as exc:
            log.error("API key error: %s", exc)
            return {"tool": "clarify",
                    "args": {"message": f"⚠️ API key error: {exc}"},
                    "description": "API key rejected", "confirm": False}
        except ProviderUnavailable as exc:
            log.warning("Provider unavailable: %s", exc)
            return {"tool": "clarify",
                    "args": {"message": f"⚠️ AI provider unavailable: {exc}"},
                    "description": "Provider unavailable", "confirm": False}
        except Exception as exc:
            log.exception("Unexpected error in parse_intent")
            return {"tool": "clarify",
                    "args": {"message": f"⚠️ Unexpected error: {exc}"},
                    "description": "Unexpected error", "confirm": False}

        action = self._parse_json_response(raw_text)

        # Safety: ALWAYS enforce confirm=True for dangerous tools
        # regardless of what the AI decided
        from .tools import TOOLS
        if TOOLS.get(action["tool"], {}).get("danger", False):
            action["confirm"] = True

        # Store conversation
        await self.memory.add("user", user_message)
        await self.memory.add("assistant", json.dumps(action))

        log.debug("Intent: tool=%s confirm=%s desc=%s",
                  action["tool"], action["confirm"], action.get("description", "")[:60])
        return action
