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

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

Features:
  • Exponential-backoff retry on transient errors (5xx, network timeouts)
  • Thread-executor HTTP so the async event loop is never blocked
  • Robust JSON extraction: strips markdown fences, fixes trailing commas,
    normalises key aliases (action→tool, parameters→args, etc.)
  • Conversation rolling memory with configurable window size
  • API key validation with a real (tiny) request
"""
from __future__ import annotations

import asyncio
import datetime
import json
import math
import platform
import re
import time
import urllib.error
import urllib.parse
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."""


class ParseError(AIEngineError):
    """Raised when the AI response cannot be parsed into a valid action dict."""


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

class ConversationMemory:
    """Thread-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)})
            # Keep only last max_turns pairs (user + assistant = 2 per turn)
            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:
        """Synchronous clear for use outside async context."""
        self._turns.clear()


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

_SYSTEM_PROMPT_TEMPLATE = """\
You are the CallMe AI agent. You control a user's computer via Telegram.
The user sends you natural language commands and you MUST respond ONLY with valid JSON — never prose.

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

## Response schema (STRICT — output ONLY this JSON, no markdown, no explanation):
{{
  "tool": "<tool_name>",
  "args": {{<key>: <value>, ...}},
  "description": "<one sentence: what you are doing>",
  "confirm": <true|false>
}}

Rules:
- "tool" must be one of the available tools listed below.
- "args" must be a flat dict of keyword arguments matching the tool's parameters.
- "confirm" must be true for any dangerous/irreversible action (marked ⚠️ below).
- If you cannot understand the request, use tool "clarify" with args {{"message": "..."}}.
- NEVER output prose. NEVER wrap JSON in markdown fences. ONLY the JSON object.

## Intent mapping (common phrasings → correct tool):
- "how's my computer", "system stats", "health check", "check cpu", "memory usage"
  → system_health
- "screenshot", "screen capture", "take a photo of screen", "what's on screen"
  → screenshot
- "weather in X", "what's the weather", "temperature outside", "forecast for X"
  → weather, args: {{"city": "X"}}
- "remind me in X minutes to Y", "set reminder", "alert me in X seconds"
  → reminder, args: {{"message": "Y", "delay_minutes": X}}
- "run X", "execute X", "shell X", "terminal X", "bash X"
  → run_shell, args: {{"command": "X"}}
- "list files in X", "ls X", "show files", "what's in X folder"
  → list_directory, args: {{"path": "X"}}
- "read file X", "show me X", "contents of X", "cat X"
  → read_file, args: {{"path": "X"}}
- "write to file X", "save text to X", "create file X with content Y"
  → write_file, args: {{"path": "X", "content": "Y"}}
- "delete X", "remove X", "rm X"
  → delete_file, args: {{"path": "X"}}
- "copy X to Y", "duplicate X"
  → copy_file, args: {{"src": "X", "dst": "Y"}}
- "move X to Y", "rename X to Y"
  → move_file, args: {{"src": "X", "dst": "Y"}}
- "search for files named X", "find file X", "where is X"
  → search_files, args: {{"query": "X"}}
- "kill process X", "terminate X", "stop X process"
  → kill_process, args: {{"name": "X"}} or {{"pid": N}}
- "list processes", "what's running", "show processes", "ps"
  → list_processes
- "volume X", "set volume to X", "volume up/down"
  → set_volume, args: {{"level": X}}
- "mute", "silence", "mute audio"
  → mute_volume
- "what's the volume", "current volume"
  → get_volume
- "clipboard", "what's in clipboard", "paste content"
  → get_clipboard
- "copy X to clipboard", "set clipboard to X"
  → set_clipboard, args: {{"text": "X"}}
- "ping X", "is X reachable", "latency to X"
  → ping, args: {{"host": "X"}}
- "check internet", "am i online", "internet speed", "connectivity"
  → internet_speed
- "open X", "open website X", "browse to X", "go to X"
  → open_url, args: {{"url": "X"}} (add https:// if missing)
- "launch X", "open app X", "start X application"
  → open_application, args: {{"name": "X"}}
- "what time", "current date", "date today", "what day is it"
  → datetime
- "uptime", "how long running", "when did it boot"
  → uptime
- "battery", "battery level", "how much battery"
  → battery
- "network stats", "data usage", "bandwidth"
  → network_stats
- "disk space", "disk usage", "free space"
  → disk_usage, args: {{"path": "/"}}
- "python version", "which python", "packages installed"
  → python_info
- "env vars", "environment variables", "show env"
  → env_vars
- "git status", "changes in repo"
  → git_status, args: {{"path": "."}}
- "git log", "recent commits"
  → git_log, args: {{"path": "."}}
- "install package X", "pip install X"
  → pip_install, args: {{"package": "X"}}
- "list packages", "installed packages", "pip list"
  → pip_list
- "calculate X", "what is X", "X * Y", "X + Y", math expressions
  → calculate, args: {{"expression": "X"}}
- "word count of X", "count words in X"
  → word_count, args: {{"text": "X"}} or {{"path": "X"}}

## Available tools:
{tool_list}
"""


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):
    """
    Call *fn()* up to *max_retries* times with exponential back-off.
    Re-raises APIKeyInvalid immediately (no retry).
    Re-raises ProviderUnavailable after all attempts are exhausted.
    """
    last_exc: Exception | None = None
    for attempt in range(max_retries):
        try:
            return fn()
        except APIKeyInvalid:
            raise  # Never retry auth errors
        except ProviderUnavailable as exc:
            last_exc = exc
            delay = min(RETRY_BASE_DELAY * (2 ** attempt), RETRY_MAX_DELAY)
            log.warning("Provider unavailable (attempt %d/%d): %s — retrying in %.1fs",
                        attempt + 1, max_retries, exc, delay)
            time.sleep(delay)
        except Exception as exc:
            last_exc = exc
            delay = min(RETRY_BASE_DELAY * (2 ** attempt), RETRY_MAX_DELAY)
            log.warning("Unexpected error (attempt %d/%d): %s — retrying in %.1fs",
                        attempt + 1, max_retries, exc, delay)
            time.sleep(delay)
    raise ProviderUnavailable(f"All {max_retries} attempts failed. Last error: {last_exc}")


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

def validate_api_key(provider: str, api_key: str, model: str) -> tuple[bool, str]:
    """
    Make a minimal real API call to validate the key.
    Returns (ok: bool, message: str).
    """
    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]
    base_url = info["base_url"]
    payload: dict[str, Any] = {}
    headers: dict[str, str] = {"Content-Type": "application/json"}

    if provider == "anthropic":
        headers["x-api-key"] = api_key.strip()
        headers["anthropic-version"] = "2023-06-01"
        url = f"{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/callme-agent/callme"
            headers["X-Title"] = "CallMe"
        url = f"{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:
        body_bytes = exc.read()
        try:
            err = json.loads(body_bytes)
            msg = (
                err.get("error", {}).get("message")
                or err.get("detail")
                or str(err)
            )
        except Exception:
            msg = body_bytes.decode(errors="replace")[:300]
        if exc.code in (401, 403):
            return False, f"Invalid API key: {msg}"
        if 400 <= exc.code < 500:
            # Key was accepted but request had some other issue — still valid
            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 so the event loop stays free.
    """

    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 transport
    # ------------------------------------------------------------------

    def _build_request(self, messages: list[dict]) -> tuple[str, dict, bytes]:
        """Return (url, headers, body_bytes) for this provider."""
        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_messages = [m for m in messages if m["role"] != "system"]
            system_content = next(
                (m["content"] for m in messages if m["role"] == "system"), ""
            )
            payload = {
                "model": self._model,
                "max_tokens": 512,
                "system": system_content,
                "messages": user_messages,
                "temperature": 0.0,
            }
        else:
            headers["Authorization"] = f"Bearer {self._api_key}"
            if self._provider == "openrouter":
                headers["HTTP-Referer"] = "https://github.com/callme-agent/callme"
                headers["X-Title"] = "CallMe"
            url = f"{self._base_url}/chat/completions"
            payload = {
                "model": self._model,
                "max_tokens": 512,
                "messages": messages,
                "temperature": 0.0,  # deterministic JSON output
            }

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

    def _extract_text(self, data: dict) -> str:
        """Pull the assistant reply text out of the API response dict."""
        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:
        """Synchronous API call — runs in executor thread. Returns raw model text."""
        url, headers, body = self._build_request(messages)

        def _do_request() -> 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:
                body_bytes = exc.read()
                try:
                    err = json.loads(body_bytes)
                    msg = err.get("error", {}).get("message") or str(err)
                except Exception:
                    msg = body_bytes.decode(errors="replace")[:500]
                if exc.code in (401, 403):
                    raise APIKeyInvalid(f"API key rejected by {self._provider}: {msg}")
                if exc.code == 429:
                    raise ProviderUnavailable(f"Rate limited by {self._provider}: {msg}")
                if exc.code >= 500:
                    raise ProviderUnavailable(f"Server error {exc.code} from {self._provider}: {msg}")
                raise ProviderUnavailable(f"HTTP {exc.code} from {self._provider}: {msg}")
            except urllib.error.URLError as exc:
                raise ProviderUnavailable(f"Network error reaching {self._provider}: {exc.reason}")

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

            # Check for embedded errors (some providers return 200 with error body)
            if "error" in data:
                err_msg = data["error"].get("message", str(data["error"]))
                raise ProviderUnavailable(f"{self._provider} returned error: {err_msg}")

            return self._extract_text(data)

        return _with_retry(_do_request)

    # ------------------------------------------------------------------
    # JSON parsing
    # ------------------------------------------------------------------

    @staticmethod
    def _parse_json_response(text: str) -> dict:
        """
        Robustly extract a JSON action object from the model's reply.

        Handles:
          - Markdown fences (```json ... ```)
          - Extra prose before/after the JSON
          - Single quotes instead of double quotes
          - Trailing commas before } or ]
          - Key aliases: action/command/function → tool, parameters/arguments → args
          - Missing required keys (filled with safe defaults)
        """
        _fallback = {
            "tool": "clarify",
            "args": {"message": "I couldn't understand the response. Please rephrase."},
            "description": "Clarify",
            "confirm": False,
        }

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

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

        # Find the outermost JSON object { ... }
        brace_start = clean.find("{")
        if brace_start == -1:
            # No JSON object — try to use the raw text as clarify message
            return {**_fallback, "args": {"message": text.strip()[:300]}}

        # Locate matching closing brace, respecting nested structures and strings
        depth = 0
        in_string = False
        escape_next = False
        brace_end = -1
        for i, ch in enumerate(clean[brace_start:], start=brace_start):
            if escape_next:
                escape_next = False
                continue
            if ch == "\\" and in_string:
                escape_next = True
                continue
            if ch == '"' and not escape_next:
                in_string = not in_string
                continue
            if in_string:
                continue
            if ch == "{":
                depth += 1
            elif ch == "}":
                depth -= 1
                if depth == 0:
                    brace_end = i + 1
                    break

        candidate = clean[brace_start: (brace_end if brace_end > 0 else len(clean))]

        # Attempt 1: direct parse
        obj: dict | None = None
        try:
            obj = json.loads(candidate)
        except json.JSONDecodeError:
            pass

        # Attempt 2: fix common issues
        if obj is None:
            fixed = candidate
            # Replace Python-style single-quoted strings with double-quoted
            fixed = re.sub(r"(?<![\\])\'", '"', fixed)
            # Remove trailing commas before } or ]
            fixed = re.sub(r",\s*([}\]])", r"\1", fixed)
            # Remove comments (// ...)
            fixed = re.sub(r"//[^\n]*", "", fixed)
            # Quote unquoted object keys: {key: "val"} → {"key": "val"}
            fixed = re.sub(r'(?<=[{,])\s*([A-Za-z_][A-Za-z0-9_]*)\s*:', r' "\1":', fixed)
            try:
                obj = json.loads(fixed)
            except json.JSONDecodeError:
                pass

        # Attempt 3: try to extract just the innermost complete JSON
        if obj is None:
            # Find all {...} substrings and try each
            for match in re.finditer(r"\{[^{}]*\}", candidate):
                try:
                    obj = json.loads(match.group())
                    break
                except json.JSONDecodeError:
                    pass

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

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

        # Ensure required keys with safe defaults
        obj.setdefault("tool", "clarify")
        obj.setdefault("args", {})
        obj.setdefault("description", "")
        obj.setdefault("confirm", False)

        # Type coercions
        if not isinstance(obj["args"], dict):
            # Try to recover: if args is a string, try parsing it
            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"])

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

        return obj

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

    async def parse_intent(self, user_message: str) -> dict:
        """
        Send *user_message* to the AI (with conversation history) and return
        a parsed action dict: {tool, args, description, confirm}.

        Never raises — falls back to clarify on any error.
        """
        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)

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

        log.debug("Intent parsed: tool=%s confirm=%s", action.get("tool"), action.get("confirm"))
        return action
