"""
Doit Agent - AI Intent Engine v2.1
- Ultra-detailed system prompt engineered for non-technical users
- Natural language examples for every tool category
- Multi-turn memory so "do that again" and "now for the other folder" work
- Dynamic tool discovery — prompt never goes stale
- Graceful fallback chain with informative user errors
"""
from __future__ import annotations

import asyncio, json, logging, re, time
from datetime import datetime
from typing import Optional

import aiohttp

from core.config import AI_PROVIDERS, HTTP_TIMEOUT, AI_RATE_LIMIT_PER_MIN

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


def _build_system_prompt(tool_list: list[dict]) -> str:
    """Dynamically build system prompt from live tool registry."""
    tool_lines = "\n".join(
        f"- {t['name']}: {t['description']}"
        + (" ⚠️ DANGEROUS — set confirm:true" if t.get("dangerous") else "")
        for t in tool_list
    )
    now = datetime.now().strftime("%A, %B %d %Y, %I:%M %p")
    platform_note = __import__('platform').system()

    return f"""You are Doit — a friendly, capable AI assistant that controls a real computer on behalf of the user.
Today is {now}. Running on: {platform_note}.

━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━
YOUR PERSONALITY
━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━
• Warm, clear and jargon-free. Many users are non-technical.
• When someone says "my computer" they mean the local machine.
• When someone says "my files" they probably mean ~/Documents or ~/Desktop.
• Prefer doing over asking. If you can infer a reasonable default, use it.
• Only use "clarify" when a task is truly ambiguous and you cannot infer what's wanted.
• Give short, human-readable descriptions (not technical jargon) in the "description" field.

━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━
AVAILABLE TOOLS
━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━
{tool_lines}

━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━
RESPONSE FORMAT — valid JSON only, never any prose outside JSON
━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━

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

Multi-step plan (use when 2+ tools needed):
{{"steps": [{{"tool": "...", "args": {{}}, "description": "..."}}, ...], "description": "Overall plain-English summary", "metadata": {{"confirm": false}}}}

Need clarification (use sparingly):
{{"tool": "clarify", "args": {{"message": "Friendly question to user"}}, "description": "Need more info", "metadata": {{"confirm": false}}}}

━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━
KEY RULES
━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━
1. ONLY use tools from the list. Never invent a tool name.
2. Respond ONLY with valid JSON. Zero prose outside the JSON structure.
3. ⚠️ Dangerous tools (delete, shutdown, kill, shell_exec, etc.) MUST have "confirm": true.
4. "description" must be plain English a non-technical user can understand.
5. Expand ~ to the user home directory in your mind — users say "my Desktop" meaning ~/Desktop.
6. If the user says "every X" or "daily" or "every morning" → use schedule_add.
7. If the user said something recently (it's in our conversation), you can refer to it — use context.
8. memory_store/recall = the user's personal long-term notes that survive restarts.
9. secret_get/set = storing passwords & API keys safely. Never repeat a stored secret back.
10. For "open an app", "launch X" → app_open. For typing/clicking → type_text / hotkey.
11. For "search for X online" → web_search. For "look up X" → web_search or web_scrape.
12. For "write code that does X" → repl_python with the code. For anything shell → shell_exec.
13. For "tell me about X" that requires live info → web_search first.
14. When user asks for weather, news, calculations, translation — just do it, no confirmation needed.
15. For unclear paths like "my project folder" → ask with clarify. For "Downloads" → ~/Downloads.

━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━
NATURAL LANGUAGE → TOOL MAPPING EXAMPLES
━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━
"show me what's in Downloads"          → fs_list {{"path": "~/Downloads"}}
"find my resume"                        → fs_search {{"path": "~", "pattern": "*resume*"}}
"find files with 'invoice' in the name" → fs_search {{"path": "~", "pattern": "*invoice*"}}
"search inside my notes for 'meeting'"  → fs_search {{"path": "~/Documents", "pattern": "*.txt", "content": "meeting"}}
"what's taking up space"               → sys_disk_usage + fs_list
"how's my computer doing"              → sys_monitor
"is the internet working"              → net_ping {{"host": "8.8.8.8"}}
"download this link"                   → net_download
"google for X" / "search for X"       → web_search
"check if my website is up"            → net_monitor
"take a picture of my screen"          → screenshot_take
"copy this to my clipboard"            → clipboard_write
"what's on my clipboard"               → clipboard_read
"open Chrome" / "launch Safari"        → app_open
"type this message for me"             → type_text
"turn volume up/down"                  → vol_set (higher or lower level)
"mute" / "unmute"                      → vol_mute
"say X out loud"                       → speech_say
"remind me in X minutes"               → reminder_set
"set an alarm for X"                   → reminder_set
"what time is it" / "what day is it"   → respond from current time in description
"what's the weather in X"              → weather {{"city": "X"}}
"translate X to Spanish"               → translate
"calculate X"                          → calc
"make a QR code for X"                 → qr_generate
"send a notification saying X"         → notify_send
"remember that X"                      → memory_store
"what did you remember" / "recall X"  → memory_recall
"save this as a note"                  → note_create
"show my notes"                        → note_list
"my to-do list" / "add to my tasks"   → todo_add / todo_list
"backup my Documents"                  → fs_backup
"clean up my Downloads"               → fs_organize + fs_clean
"compress/zip this folder"             → fs_compress
"extract/unzip this"                   → fs_extract
"git status" / "any new changes"       → git_status
"commit my work"                       → git_add + git_commit + git_push (multi-step)
"install X"                            → pkg_install
"run my tests"                         → test_run
"build my project"                     → build_run
"run: ls -la"                          → shell_exec (user explicitly wants shell)
"shutdown" / "restart" / "sleep"       → sys_shutdown / sys_restart / sys_sleep (all confirm:true)
"kill process X"                       → sys_kill (confirm:true)
"store my API key"                     → secret_set
"get my API key"                       → secret_get
"""


class ConversationMemory:
    """Per-session rolling context for multi-turn natural language."""

    def __init__(self, max_turns: int = 12):
        self.max_turns = max_turns
        self._turns: list[dict] = []

    def add(self, role: str, content: str):
        self._turns.append({"role": role, "content": content})
        # Keep last N turns to avoid token bloat
        if len(self._turns) > self.max_turns * 2:
            self._turns = self._turns[-(self.max_turns * 2):]

    def get_messages(self, user_message: str) -> list[dict]:
        """Build messages array from history + new message."""
        return list(self._turns) + [{"role": "user", "content": user_message}]

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

    @property
    def turn_count(self) -> int:
        return len(self._turns) // 2


class RateLimiter:
    def __init__(self, max_per_minute: int):
        self.max_per_minute = max_per_minute
        self._calls: list[float] = []

    async def acquire(self):
        now = time.time()
        self._calls = [t for t in self._calls if now - t < 60]
        if len(self._calls) >= self.max_per_minute:
            wait = 60 - (now - self._calls[0])
            logger.debug("Rate limit: waiting %.1fs", wait)
            await asyncio.sleep(wait)
        self._calls.append(time.time())


class AIEngine:
    """Multi-provider AI engine — friendly, capable, context-aware."""

    def __init__(self, config: dict):
        self.provider = config.get("ai_provider", "openai")
        self.model = config.get("ai_model", "gpt-4o-mini")
        self.api_key = config.get("ai_api_key", "")
        self.base_url = config.get("ai_base_url", "")
        self.fallback_provider = config.get("ai_fallback_provider")
        self.fallback_model = config.get("ai_fallback_model")
        self.fallback_key = config.get("ai_fallback_key")
        self._rate_limiter = RateLimiter(AI_RATE_LIMIT_PER_MIN)
        self._available = True
        self._memory = ConversationMemory()
        self._tool_list: list[dict] = []
        self._system_prompt: str = ""

    def refresh_tools(self):
        """Sync tool list from the live registry and rebuild system prompt."""
        try:
            from tools.registry import get_registry
            self._tool_list = get_registry().list_with_descriptions()
            self._system_prompt = _build_system_prompt(self._tool_list)
            logger.debug("Refreshed tool list: %d tools", len(self._tool_list))
        except Exception as e:
            logger.warning("Could not refresh tool list: %s", e)

    def clear_memory(self):
        """Clear conversation context — fresh start."""
        self._memory.clear()
        logger.info("Conversation memory cleared")

    @property
    def memory_turns(self) -> int:
        return self._memory.turn_count

    # ── HTTP helpers ─────────────────────────────────────────────────────────

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

    def _endpoint(self, provider: str, base_url: str) -> str:
        return f"{base_url}/messages" if provider == "anthropic" else f"{base_url}/chat/completions"

    def _payload(self, provider: str, model: str, messages: list[dict]) -> dict:
        system = self._system_prompt or _build_system_prompt(self._tool_list)
        if provider == "anthropic":
            return {"model": model, "max_tokens": 2048, "system": system, "messages": messages}
        full = [{"role": "system", "content": system}] + messages
        p = {"model": model, "messages": full, "max_tokens": 2048, "temperature": 0.1}
        if provider == "openai":
            p["response_format"] = {"type": "json_object"}
        return p

    def _extract_text(self, provider: str, data: dict) -> str:
        if provider == "anthropic":
            for block in data.get("content", []):
                if isinstance(block, dict) and block.get("type") == "text":
                    return block["text"]
            raise ValueError(f"Unexpected Anthropic response structure")
        choices = data.get("choices", [])
        if not choices:
            raise ValueError("Empty choices in provider response")
        return choices[0].get("message", {}).get("content", "")

    async def _call(self, provider: str, model: str, api_key: str,
                    base_url: str, messages: list[dict]) -> dict:
        await self._rate_limiter.acquire()
        async with aiohttp.ClientSession(
                timeout=aiohttp.ClientTimeout(total=HTTP_TIMEOUT)) as sess:
            async with sess.post(
                    self._endpoint(provider, base_url),
                    headers=self._headers(provider, api_key),
                    json=self._payload(provider, model, messages)) as resp:
                if resp.status == 429: raise QuotaExceededError("API quota exceeded")
                if resp.status == 401: raise InvalidKeyError("Invalid API key")
                if resp.status >= 500: raise ProviderError(f"Provider HTTP {resp.status}")
                resp.raise_for_status()
                data = await resp.json()
        return self._parse(self._extract_text(provider, data))

    def _parse(self, text: str) -> dict:
        """Parse AI response into action dict, tolerating markdown fences."""
        text = text.strip()
        # Strip markdown code fences
        text = re.sub(r'^```(?:json)?\s*\n?', '', text, flags=re.MULTILINE)
        text = re.sub(r'\n?```\s*$', '', text, flags=re.MULTILINE)
        # Extract outermost JSON object
        m = re.search(r'\{.*\}', text, re.DOTALL)
        if m:
            text = m.group()
        try:
            action = json.loads(text)
            # Ensure required fields
            if "tool" not in action and "steps" not in action:
                raise ValueError("No 'tool' or 'steps' key in response")
            return action
        except (json.JSONDecodeError, ValueError) as e:
            logger.warning("AI response parse failed: %s | raw: %.200s", e, text)
            return {
                "tool": "clarify",
                "args": {"message": "I'm not sure I understood that correctly. Could you rephrase?"},
                "description": "Clarification needed",
                "metadata": {"confirm": False},
            }

    # ── Public API ────────────────────────────────────────────────────────────

    async def parse_intent(self, user_message: str) -> dict:
        """Parse natural language into an action plan, with full conversation context."""
        if not self._tool_list:
            self.refresh_tools()
        if not self._system_prompt:
            self._system_prompt = _build_system_prompt(self._tool_list)

        messages = self._memory.get_messages(user_message)

        # Try primary provider
        try:
            action = await self._call(
                self.provider, self.model, self.api_key, self.base_url, messages)
            self._available = True
            self._memory.add("user", user_message)
            self._memory.add("assistant", json.dumps(action))
            return action
        except (InvalidKeyError, QuotaExceededError) as e:
            logger.error("Primary AI hard error: %s", e)
            self._available = False
        except Exception as e:
            logger.warning("Primary AI call failed: %s", e)
            self._available = False

        # Try fallback provider
        if self.fallback_provider and self.fallback_key:
            fb_info = AI_PROVIDERS.get(self.fallback_provider, {})
            fb_url = fb_info.get("base_url", "")
            try:
                logger.info("Trying fallback AI: %s", self.fallback_provider)
                action = await self._call(
                    self.fallback_provider,
                    self.fallback_model or "",
                    self.fallback_key, fb_url, messages)
                self._memory.add("user", user_message)
                self._memory.add("assistant", json.dumps(action))
                return action
            except Exception as e2:
                logger.error("Fallback AI also failed: %s", e2)

        return {
            "tool": "error",
            "args": {"message": (
                "I'm having trouble connecting to the AI right now. "
                "Please check your internet connection and try again in a moment. "
                "You can also use /run for direct shell commands."
            )},
            "description": "AI unavailable",
            "metadata": {"confirm": False},
        }

    async def test_connection(self) -> tuple[bool, str]:
        """Test AI provider connectivity. Returns (success, message)."""
        if not self._system_prompt:
            self._system_prompt = "Respond only with this exact JSON: {\"tool\":\"status\",\"args\":{},\"description\":\"ok\",\"metadata\":{\"confirm\":false}}"
        try:
            result = await self._call(
                self.provider, self.model, self.api_key, self.base_url,
                [{"role": "user", "content":
                  '{"tool":"status","args":{},"description":"test","metadata":{"confirm":false}}'}])
            if "tool" in result:
                return True, f"✅ Connected to {self.provider} / {self.model}"
            return False, "Unexpected response format"
        except InvalidKeyError: return False, "❌ Invalid API key"
        except QuotaExceededError: return False, "❌ API quota exceeded"
        except Exception as e: return False, f"❌ {e}"


# ── Custom exceptions ─────────────────────────────────────────────────────────

class AIError(Exception): pass
class InvalidKeyError(AIError): pass
class QuotaExceededError(AIError): pass
class ProviderError(AIError): pass


async def validate_api_key(provider: str, model: str, api_key: str, base_url: str) -> tuple[bool, str]:
    """Validate an API key with a real test call (3 retries)."""
    engine = AIEngine({
        "ai_provider": provider, "ai_model": model,
        "ai_api_key": api_key, "ai_base_url": base_url,
    })
    for attempt in range(3):
        success, msg = await engine.test_connection()
        if success: return True, msg
        if "Invalid API key" in msg or "quota" in msg.lower(): return False, msg
        if attempt < 2: await asyncio.sleep(2 ** attempt)
    return False, f"Failed after 3 attempts. Last error: {msg}"
