"""
Multi-provider LLM adapter for CrashSense.

Supports (in priority order):
  1. Gemini  (default — uses google-genai SDK)
  2. OpenAI  (GPT-4o-mini or user choice)
  3. Claude  (Anthropic)
  4. Ollama  (local, offline)

The active provider is read from config; the user can override at runtime.
"""

from __future__ import annotations

import os
import re
import time
from typing import Any, Dict, List, Optional, Tuple

import requests
from rich.console import Console

from ..config import get_api_key, load_config

console = Console()

# ── API endpoints ─────────────────────────────────────────────────────────
OPENAI_CHAT_URL = "https://api.openai.com/v1/chat/completions"
ANTHROPIC_MSG_URL = "https://api.anthropic.com/v1/messages"


class LLMAdapter:
    """Unified LLM interface with caching, rate-limiting, and retry."""

    def __init__(
        self,
        provider: str = "gemini",
        model: Optional[str] = None,
        cfg: Optional[Dict[str, Any]] = None,
    ) -> None:
        self.cfg = cfg or load_config()
        self.provider = provider or self.cfg.get("provider", "gemini")

        # Resolve model per provider
        self.model = model or self._default_model()

        # API keys from env
        self.gemini_key = get_api_key("gemini")
        self.openai_key = get_api_key("openai")
        self.claude_key = get_api_key("claude")

        # Cache & rate limiting
        self._cache: Dict[tuple, Dict] = {}
        try:
            self._min_interval = float(os.environ.get("CRASHSENSE_MIN_CALL_INTERVAL", "0.5"))
        except ValueError:
            self._min_interval = 0.5
        self._last_call_ts = 0.0

        # Ollama host
        self.ollama_host = self.cfg.get("ollama", {}).get("host", "http://localhost:11434")

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

    def analyze(self, prompt: str, system: str = "") -> Dict[str, str]:
        """
        Send *prompt* to the active LLM and return structured result.

        Returns dict with keys: explanation, root_cause, suggested_fix.
        """
        if not system:
            system = (
                "You are CrashSense, an expert AI agent that diagnoses failed or reverted "
                "blockchain transactions. You explain errors in clear, human-readable language "
                "and suggest concrete fixes. Be concise but thorough."
            )

        cache_key = (self.provider, self.model, system[:256], prompt[:4096])
        if cache_key in self._cache:
            return self._cache[cache_key]

        dispatch = {
            "gemini": self._call_gemini,
            "openai": self._call_openai,
            "claude": self._call_claude,
            "ollama": self._call_ollama,
        }
        fn = dispatch.get(self.provider, self._heuristic_answer)
        result = fn(prompt, system)
        self._cache[cache_key] = result
        return result

    def retrieve(self, query: str, docs: List[str], top_k: int = 3) -> List[Tuple[str, float]]:
        """Simple RAG retriever: keyword-overlap (embedding optional)."""
        docs = [d for d in docs if isinstance(d, str) and d.strip()]
        if not docs:
            return []

        qk = {w for w in re.findall(r"\b\w+\b", query.lower()) if len(w) > 2}
        scored: List[Tuple[str, float]] = []
        for doc in docs:
            tk = {w for w in re.findall(r"\b\w+\b", doc.lower()) if len(w) > 2}
            overlap = len(qk & tk)
            scored.append((doc, float(overlap)))
        return sorted(scored, key=lambda x: x[1], reverse=True)[:top_k]

    def list_providers(self) -> List[Dict[str, str]]:
        """Return available providers and their status."""
        providers = []
        for name, check in [
            ("gemini", bool(self.gemini_key)),
            ("openai", bool(self.openai_key)),
            ("claude", bool(self.claude_key)),
            ("ollama", self._ollama_reachable()),
        ]:
            providers.append({
                "name": name,
                "available": check,
                "active": name == self.provider,
            })
        return providers

    # ── Provider implementations ──────────────────────────────────────────

    def _call_gemini(self, prompt: str, system: str) -> Dict[str, str]:
        """Call Google Gemini via the google-genai SDK."""
        if not self.gemini_key:
            return self._error_result("GEMINI_API_KEY not set. Add it to your .env file.")
        try:
            from google import genai  # type: ignore

            self._rate_limit()
            client = genai.Client(api_key=self.gemini_key)
            response = client.models.generate_content(
                model=self.model,
                contents=f"{system}\n\n{prompt}",
            )
            text = response.text.strip() if response.text else ""
            if not text:
                return self._error_result("Gemini returned empty response.")
            return self._parse_llm_text(text)
        except ImportError:
            return self._error_result(
                "google-genai package not installed. Run: pip install google-genai"
            )
        except Exception as exc:
            return self._error_result(f"Gemini error: {exc}")

    def _call_openai(self, prompt: str, system: str) -> Dict[str, str]:
        """Call OpenAI Chat Completions API."""
        if not self.openai_key:
            return self._error_result("OPENAI_API_KEY not set.")
        headers = {
            "Authorization": f"Bearer {self.openai_key}",
            "Content-Type": "application/json",
        }
        body = {
            "model": self.model,
            "messages": [
                {"role": "system", "content": system},
                {"role": "user", "content": prompt},
            ],
            "max_tokens": 1024,
            "temperature": 0.2,
        }
        try:
            resp = self._post_with_retries(OPENAI_CHAT_URL, headers=headers, json=body)
            resp.raise_for_status()
            text = resp.json()["choices"][0]["message"]["content"].strip()
            return self._parse_llm_text(text)
        except Exception as exc:
            return self._error_result(f"OpenAI error: {exc}")

    def _call_claude(self, prompt: str, system: str) -> Dict[str, str]:
        """Call Anthropic Claude Messages API."""
        if not self.claude_key:
            return self._error_result("ANTHROPIC_API_KEY not set.")
        headers = {
            "x-api-key": self.claude_key,
            "Content-Type": "application/json",
            "anthropic-version": "2023-06-01",
        }
        body = {
            "model": self.model,
            "max_tokens": 1024,
            "system": system,
            "messages": [{"role": "user", "content": prompt}],
        }
        try:
            resp = self._post_with_retries(ANTHROPIC_MSG_URL, headers=headers, json=body)
            resp.raise_for_status()
            data = resp.json()
            text = data.get("content", [{}])[0].get("text", "").strip()
            return self._parse_llm_text(text)
        except Exception as exc:
            return self._error_result(f"Claude error: {exc}")

    def _call_ollama(self, prompt: str, system: str) -> Dict[str, str]:
        """Call local Ollama HTTP API."""
        url = f"{self.ollama_host}/api/generate"
        full_prompt = f"{system}\n\n{prompt}" if system else prompt
        try:
            self._rate_limit()
            resp = self._post_with_retries(
                url,
                headers=None,
                json={"model": self.model, "prompt": full_prompt, "stream": False},
                timeout=180,
            )
            if resp.status_code == 200:
                text = resp.json().get("response", "").strip()
                return self._parse_llm_text(text) if text else self._error_result("Ollama returned empty response.")
            err = resp.json().get("error", resp.text)
            if "not found" in (err or "").lower():
                return self._error_result(f"Model '{self.model}' not found. Run: ollama pull {self.model}")
            return self._error_result(f"Ollama error: {err}")
        except requests.exceptions.ConnectionError:
            return self._error_result("Ollama not reachable. Start it with: ollama serve")
        except Exception as exc:
            return self._error_result(f"Ollama error: {exc}")

    def _heuristic_answer(self, prompt: str, system: str = "") -> Dict[str, str]:
        """Offline heuristic fallback when no LLM is available."""
        lower = prompt.lower()
        explanations = []

        if "out of gas" in lower or "gas_utilisation_pct" in lower:
            explanations.append("The transaction ran out of gas before completing execution.")
        if "revert" in lower or "execution reverted" in lower:
            explanations.append("The transaction was reverted by a require/revert statement in the smart contract.")
        if "insufficient" in lower and "balance" in lower:
            explanations.append("The sender did not have enough token balance to complete the operation.")
        if "allowance" in lower:
            explanations.append("The contract does not have sufficient token allowance (approval) to spend on your behalf.")
        if "overflow" in lower or "underflow" in lower:
            explanations.append("An arithmetic overflow or underflow occurred in the contract.")
        if "slippage" in lower or "expired" in lower or "deadline" in lower:
            explanations.append("The swap deadline expired or slippage tolerance was exceeded.")
        if "nonce" in lower:
            explanations.append("Transaction nonce mismatch — a pending transaction may be blocking this one.")

        explanation = " ".join(explanations) if explanations else "Unable to determine the root cause without an LLM."
        return {
            "explanation": f"[Heuristic] {explanation}",
            "root_cause": "See explanation",
            "suggested_fix": "Connect an LLM provider (gemini/openai/claude/ollama) for detailed diagnosis.",
        }

    # ── Internal helpers ──────────────────────────────────────────────────

    def _default_model(self) -> str:
        """Get model from config or use provider defaults."""
        # Check if user configured a model in their config
        provider_cfg = self.cfg.get(self.provider, {})
        if isinstance(provider_cfg, dict) and "model" in provider_cfg:
            return str(provider_cfg["model"])
        
        # Fallback to provider defaults (these should be set by user)
        # These are generic fallbacks only
        defaults = {
            "gemini": "gemini-2.0-flash-exp",  # Generic Gemini model
            "openai": "gpt-4o-mini",  # Affordable general-purpose
            "claude": "claude-3-5-sonnet-20241022",  # Latest Sonnet
            "ollama": "llama3.2:1b",  # Local model
        }
        return defaults.get(self.provider, "gemini-2.0-flash-exp")

    def _rate_limit(self) -> None:
        elapsed = time.time() - self._last_call_ts
        if elapsed < self._min_interval:
            time.sleep(self._min_interval - elapsed)
        self._last_call_ts = time.time()

    def _post_with_retries(
        self, url: str, *, headers=None, json=None, timeout: int = 60, max_retries: int = 3
    ) -> requests.Response:
        backoff = 1.0
        last_exc: Optional[Exception] = None
        for attempt in range(max_retries):
            self._rate_limit()
            try:
                resp = requests.post(url, headers=headers, json=json, timeout=timeout)
                if resp.status_code in (429, 500, 502, 503, 504) and attempt < max_retries - 1:
                    time.sleep(backoff)
                    backoff *= 2
                    continue
                return resp
            except (requests.exceptions.Timeout, requests.exceptions.ConnectionError) as exc:
                last_exc = exc
                if attempt < max_retries - 1:
                    time.sleep(backoff)
                    backoff *= 2
                    continue
                raise
        raise last_exc or RuntimeError("All retries exhausted")

    def _ollama_reachable(self) -> bool:
        try:
            resp = requests.get(f"{self.ollama_host}/api/tags", timeout=3)
            return resp.ok
        except Exception:
            return False

    @staticmethod
    def _parse_llm_text(text: str) -> Dict[str, str]:
        """
        Parse free-form LLM output into structured fields.
        Looks for markdown-style headers; falls back to full text.
        """
        result = {"explanation": text, "root_cause": "", "suggested_fix": ""}

        # Try to extract structured sections
        sections = {
            "root_cause": re.compile(r"(?:^|\n)#+\s*Root\s*Cause[:\s]*\n(.*?)(?=\n#+|\Z)", re.I | re.S),
            "explanation": re.compile(r"(?:^|\n)#+\s*Explanation[:\s]*\n(.*?)(?=\n#+|\Z)", re.I | re.S),
            "suggested_fix": re.compile(
                r"(?:^|\n)#+\s*(?:Suggested?\s*Fix(?:es)?|Remediation|Solution|How\s*to\s*Fix)[:\s]*\n(.*?)(?=\n#+|\Z)",
                re.I | re.S,
            ),
        }
        for key, pat in sections.items():
            m = pat.search(text)
            if m:
                result[key] = m.group(1).strip()

        if not result["root_cause"]:
            result["root_cause"] = "See explanation"

        return result

    @staticmethod
    def _error_result(msg: str) -> Dict[str, str]:
        return {"explanation": msg, "root_cause": "error", "suggested_fix": ""}
