"""
Brynq SDK — Multi-model AI orchestration in 3 lines of code.

    from brynq import Brynq

    b = Brynq()
    answer = b.ask("summarize this document")

That's it. Brynq auto-detects your models (Ollama, Claude, Gemini),
picks the best one for the task, and chains them if needed.

Features:
    b.ask(prompt)                     → auto-route to best model
    b.ask(prompt, model="claude")     → specific model
    b.chain("debate", prompt)         → multi-model debate
    b.chain("refine", prompt)         → draft → critique → revise
    b.chain("fan_out", prompt)        → parallel models, merged result
    b.models()                        → list available models
    b.fallback(prompt, ["claude", "gemini", "ollama:llama3"])
    b.cost_report()                   → token usage and cost summary
"""

from __future__ import annotations

import json
import logging
import os
import subprocess
import time
import urllib.error
import urllib.request
from dataclasses import dataclass, field
from typing import Any, Dict, List, Optional

log = logging.getLogger("brynq")

OLLAMA_URL = os.environ.get("OLLAMA_URL", "http://localhost:11434")


@dataclass
class BrynqResponse:
    """Response from a Brynq call."""
    content: str
    model: str
    strategy: str = "single"
    steps: int = 1
    elapsed_ms: int = 0
    tokens: int = 0
    cost: float = 0.0
    success: bool = True
    error: Optional[str] = None
    metadata: Dict[str, Any] = field(default_factory=dict)

    def __str__(self) -> str:
        return self.content

    def __repr__(self) -> str:
        return f"BrynqResponse(model={self.model!r}, strategy={self.strategy!r}, tokens={self.tokens})"


class Brynq:
    """Multi-model AI orchestration.

    Auto-detects Ollama, Claude CLI, and Gemini CLI on your machine.
    Routes prompts to the best model. Chains models for complex tasks.

    Args:
        ollama_url: Ollama API base URL (default: http://localhost:11434)
        default_model: Preferred model (default: auto-detect best available)
        verbose: Print debug info (default: False)

    Example:
        >>> b = Brynq()
        >>> b.ask("What is the capital of France?")
        BrynqResponse(content="The capital of France is Paris.", model="ollama:llama3")
    """

    def __init__(
        self,
        ollama_url: str = "",
        default_model: str = "",
        verbose: bool = False,
    ):
        self._ollama_url = ollama_url or OLLAMA_URL
        self._default_model = default_model
        self._verbose = verbose
        self._models: List[Dict[str, str]] = []
        self._history: List[Dict[str, Any]] = []
        self._total_tokens = 0
        self._total_cost = 0.0
        self._total_calls = 0

        if verbose:
            logging.basicConfig(level=logging.DEBUG)

        # Auto-detect models on init
        self._detect_models()

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

    def ask(
        self,
        prompt: str,
        model: str = "",
        system: str = "",
        temperature: float = 0.7,
    ) -> BrynqResponse:
        """Send a prompt to an AI model. Auto-routes if no model specified.

        Args:
            prompt: Your question or instruction.
            model: Specific model (e.g., "claude", "gemini", "ollama:llama3").
                   If empty, Brynq picks the best available model.
            system: Optional system prompt.
            temperature: Creativity (0.0 = focused, 1.0 = creative).

        Returns:
            BrynqResponse with content, model used, timing, and token info.

        Example:
            >>> b.ask("Explain quantum computing simply")
            >>> b.ask("Write a poem", model="claude")
            >>> b.ask("Translate to Spanish", model="ollama:mistral")
        """
        target = model or self._pick_model(prompt)
        if not target:
            return BrynqResponse(
                content="No AI models available. Install Ollama (ollama.com) or set a Claude/Gemini API key.",
                model="none", success=False, error="no_models",
            )

        t0 = time.time()
        content, tokens = self._call_model(target, prompt, system, temperature)
        elapsed = int((time.time() - t0) * 1000)

        response = BrynqResponse(
            content=content,
            model=target,
            strategy="single",
            steps=1,
            elapsed_ms=elapsed,
            tokens=tokens,
            success=bool(content and not content.startswith("[Error")),
        )

        self._record(response)
        return response

    def chain(
        self,
        strategy: str,
        prompt: str,
        models: Optional[List[str]] = None,
        system: str = "",
    ) -> BrynqResponse:
        """Run a multi-model chain with a specific strategy.

        Strategies:
            "debate"     — Two models argue, a judge decides.
            "refine"     — Draft → critique → revise.
            "fan_out"    — Multiple models in parallel, results merged.
            "sequential" — Model A passes context to Model B to Model C.

        Args:
            strategy: Chain strategy name.
            prompt: The task/question.
            models: Specific models to use (default: auto-select).
            system: Optional system prompt.

        Returns:
            BrynqResponse with the final chain output.

        Example:
            >>> b.chain("debate", "Should startups use microservices?")
            >>> b.chain("refine", "Write a press release for our product launch")
            >>> b.chain("fan_out", "Research the pros and cons of Rust")
        """
        available = models or [m["id"] for m in self._models[:3]]
        if len(available) < 2:
            return self.ask(prompt, model=available[0] if available else "")

        t0 = time.time()

        if strategy == "debate":
            result = self._chain_debate(prompt, available, system)
        elif strategy == "refine":
            result = self._chain_refine(prompt, available, system)
        elif strategy == "fan_out":
            result = self._chain_fan_out(prompt, available, system)
        elif strategy == "sequential":
            result = self._chain_sequential(prompt, available, system)
        else:
            return self.ask(prompt)

        elapsed = int((time.time() - t0) * 1000)
        result.strategy = strategy
        result.elapsed_ms = elapsed
        self._record(result)
        return result

    def fallback(
        self,
        prompt: str,
        models: List[str],
        system: str = "",
    ) -> BrynqResponse:
        """Try models in order. If one fails, try the next.

        Args:
            prompt: Your question.
            models: Ordered list of models to try.
            system: Optional system prompt.

        Returns:
            Response from the first model that succeeds.

        Example:
            >>> b.fallback("Explain AI", ["claude", "gemini", "ollama:llama3"])
        """
        for model in models:
            response = self.ask(prompt, model=model, system=system)
            if response.success:
                return response
            if self._verbose:
                log.warning("Model %s failed, trying next...", model)

        return BrynqResponse(
            content="All models failed.",
            model="none", success=False, error="all_failed",
        )

    def models(self) -> List[Dict[str, str]]:
        """List all available AI models.

        Returns:
            List of dicts with id, name, provider, type.

        Example:
            >>> for m in b.models():
            ...     print(f"{m['id']} ({m['provider']})")
        """
        self._detect_models()
        return self._models

    def cost_report(self) -> Dict[str, Any]:
        """Get token usage and cost summary for this session.

        Returns:
            Dict with total_calls, total_tokens, total_cost, history.

        Example:
            >>> report = b.cost_report()
            >>> print(f"Spent {report['total_tokens']} tokens across {report['total_calls']} calls")
        """
        return {
            "total_calls": self._total_calls,
            "total_tokens": self._total_tokens,
            "total_cost": round(self._total_cost, 6),
            "history": self._history[-20:],
        }

    # ── Model Detection ─────────────────────────────────────────────

    def _detect_models(self) -> None:
        """Detect available models on this machine."""
        self._models = []

        # Ollama models
        try:
            req = urllib.request.Request(f"{self._ollama_url}/api/tags")
            with urllib.request.urlopen(req, timeout=3) as resp:
                data = json.loads(resp.read().decode())
                for m in data.get("models", []):
                    full_name = m.get("name", "")
                    short_name = full_name.split(":")[0]
                    self._models.append({
                        "id": f"ollama:{full_name}",
                        "name": short_name,
                        "provider": "ollama",
                        "type": "local",
                    })
        except Exception:
            pass

        # Claude CLI
        if self._cli_exists("claude") or self._cli_exists("claude.cmd"):
            self._models.append({
                "id": "claude",
                "name": "Claude",
                "provider": "anthropic",
                "type": "cloud",
            })

        # Gemini CLI
        if self._cli_exists("gemini") or self._cli_exists("gemini.cmd"):
            self._models.append({
                "id": "gemini",
                "name": "Gemini",
                "provider": "google",
                "type": "cloud",
            })

        # API keys
        if os.environ.get("ANTHROPIC_API_KEY"):
            if not any(m["id"] == "claude" for m in self._models):
                self._models.append({
                    "id": "claude",
                    "name": "Claude (API)",
                    "provider": "anthropic",
                    "type": "cloud",
                })

        if os.environ.get("GOOGLE_API_KEY") or os.environ.get("GEMINI_API_KEY"):
            if not any(m["id"] == "gemini" for m in self._models):
                self._models.append({
                    "id": "gemini",
                    "name": "Gemini (API)",
                    "provider": "google",
                    "type": "cloud",
                })

        if os.environ.get("OPENAI_API_KEY"):
            self._models.append({
                "id": "gpt",
                "name": "GPT-4o",
                "provider": "openai",
                "type": "cloud",
            })

        if self._verbose:
            log.info("Detected %d models: %s", len(self._models),
                     ", ".join(m["id"] for m in self._models))

    def _cli_exists(self, name: str) -> bool:
        """Check if a CLI tool exists on PATH."""
        import shutil
        return shutil.which(name) is not None

    def _pick_model(self, prompt: str = "") -> str:
        """Pick the best available model for a prompt."""
        if self._default_model:
            return self._default_model
        if not self._models:
            return ""
        # Prefer local models (free), then cloud
        local = [m for m in self._models if m["type"] == "local"]
        if local:
            return local[0]["id"]
        return self._models[0]["id"]

    # ── Model Calls ─────────────────────────────────────────────────

    def _call_model(
        self, model_id: str, prompt: str, system: str = "", temperature: float = 0.7
    ) -> tuple:
        """Call a specific model. Returns (content, token_count)."""
        try:
            if model_id.startswith("ollama:"):
                return self._call_ollama(model_id, prompt, system, temperature)
            elif model_id == "claude":
                return self._call_claude_cli(prompt)
            elif model_id == "gemini":
                return self._call_gemini_cli(prompt)
            elif model_id == "gpt":
                return self._call_openai_api(prompt, system, temperature)
            else:
                return f"[Error: Unknown model '{model_id}']", 0
        except Exception as e:
            return f"[Error: {e}]", 0

    def _call_ollama(
        self, model_id: str, prompt: str, system: str, temperature: float
    ) -> tuple:
        """Call Ollama API."""
        model_name = model_id.replace("ollama:", "")
        payload = {
            "model": model_name,
            "prompt": prompt,
            "stream": False,
            "options": {"temperature": temperature},
        }
        if system:
            payload["system"] = system

        body = json.dumps(payload).encode()
        req = urllib.request.Request(
            f"{self._ollama_url}/api/generate",
            data=body,
            headers={"Content-Type": "application/json"},
            method="POST",
        )
        with urllib.request.urlopen(req, timeout=120) as resp:
            data = json.loads(resp.read().decode())
            content = data.get("response", "")
            tokens = data.get("eval_count", 0) + data.get("prompt_eval_count", 0)
            return content.strip(), tokens

    def _call_claude_cli(self, prompt: str) -> tuple:
        """Call Claude via CLI (claude -p)."""
        cmd = "claude.cmd" if os.name == "nt" else "claude"
        result = subprocess.run(
            [cmd, "-p", prompt, "--output-format", "text"],
            capture_output=True, text=True, timeout=120,
        )
        content = result.stdout.strip() if result.returncode == 0 else result.stderr.strip()
        tokens = len(content.split()) * 2  # rough estimate
        return content, tokens

    def _call_gemini_cli(self, prompt: str) -> tuple:
        """Call Gemini via CLI."""
        cmd = "gemini.cmd" if os.name == "nt" else "gemini"
        result = subprocess.run(
            [cmd, "-p", prompt],
            capture_output=True, text=True, timeout=120,
        )
        content = result.stdout.strip() if result.returncode == 0 else result.stderr.strip()
        tokens = len(content.split()) * 2
        return content, tokens

    def _call_openai_api(self, prompt: str, system: str, temperature: float) -> tuple:
        """Call OpenAI API directly."""
        api_key = os.environ.get("OPENAI_API_KEY", "")
        if not api_key:
            return "[Error: OPENAI_API_KEY not set]", 0

        messages = []
        if system:
            messages.append({"role": "system", "content": system})
        messages.append({"role": "user", "content": prompt})

        payload = {
            "model": "gpt-4o",
            "messages": messages,
            "temperature": temperature,
        }
        body = json.dumps(payload).encode()
        req = urllib.request.Request(
            "https://api.openai.com/v1/chat/completions",
            data=body,
            headers={
                "Content-Type": "application/json",
                "Authorization": f"Bearer {api_key}",
            },
            method="POST",
        )
        with urllib.request.urlopen(req, timeout=120) as resp:
            data = json.loads(resp.read().decode())
            content = data["choices"][0]["message"]["content"]
            tokens = data.get("usage", {}).get("total_tokens", 0)
            return content.strip(), tokens

    # ── Chain Strategies ────────────────────────────────────────────

    def _chain_debate(self, prompt: str, models: List[str], system: str) -> BrynqResponse:
        """Two models argue, third judges."""
        advocate_a = models[0]
        advocate_b = models[1] if len(models) > 1 else models[0]
        judge = models[2] if len(models) > 2 else models[0]

        a_response, a_tokens = self._call_model(
            advocate_a, f"Argue IN FAVOR of this position:\n\n{prompt}", system)
        b_response, b_tokens = self._call_model(
            advocate_b, f"Argue AGAINST this position:\n\n{prompt}", system)
        judge_prompt = (
            f"Two advocates debated this question: {prompt}\n\n"
            f"Advocate A ({advocate_a}):\n{a_response}\n\n"
            f"Advocate B ({advocate_b}):\n{b_response}\n\n"
            f"As the judge, synthesize both arguments and give a final verdict."
        )
        verdict, j_tokens = self._call_model(judge, judge_prompt, system)

        return BrynqResponse(
            content=verdict,
            model=f"{advocate_a}+{advocate_b}+{judge}",
            steps=3,
            tokens=a_tokens + b_tokens + j_tokens,
        )

    def _chain_refine(self, prompt: str, models: List[str], system: str) -> BrynqResponse:
        """Draft → critique → revise."""
        drafter = models[0]
        critic = models[1] if len(models) > 1 else models[0]
        reviser = models[2] if len(models) > 2 else models[0]

        draft, d_tokens = self._call_model(drafter, prompt, system)
        critique, c_tokens = self._call_model(
            critic, f"Critique this draft and suggest improvements:\n\n{draft}")
        revised, r_tokens = self._call_model(
            reviser,
            f"Original request: {prompt}\n\nDraft:\n{draft}\n\nCritique:\n{critique}\n\n"
            f"Produce a final revised version incorporating the feedback.")

        return BrynqResponse(
            content=revised,
            model=f"{drafter}+{critic}+{reviser}",
            steps=3,
            tokens=d_tokens + c_tokens + r_tokens,
        )

    def _chain_fan_out(self, prompt: str, models: List[str], system: str) -> BrynqResponse:
        """Multiple models in parallel, results merged."""
        responses = []
        total_tokens = 0
        for model_id in models:
            content, tokens = self._call_model(model_id, prompt, system)
            responses.append(f"**{model_id}:**\n{content}")
            total_tokens += tokens

        merge_prompt = (
            f"Multiple AI models answered this question: {prompt}\n\n"
            + "\n\n---\n\n".join(responses)
            + "\n\nSynthesize all responses into one comprehensive answer."
        )
        merged, m_tokens = self._call_model(models[0], merge_prompt, system)

        return BrynqResponse(
            content=merged,
            model="+".join(models),
            steps=len(models) + 1,
            tokens=total_tokens + m_tokens,
        )

    def _chain_sequential(self, prompt: str, models: List[str], system: str) -> BrynqResponse:
        """Model A passes context to Model B to Model C."""
        context = prompt
        total_tokens = 0
        for i, model_id in enumerate(models):
            step_prompt = context if i == 0 else (
                f"Previous step output:\n{context}\n\n"
                f"Continue working on the original task: {prompt}"
            )
            context, tokens = self._call_model(model_id, step_prompt, system)
            total_tokens += tokens

        return BrynqResponse(
            content=context,
            model="→".join(models),
            steps=len(models),
            tokens=total_tokens,
        )

    # ── Internal ────────────────────────────────────────────────────

    def _record(self, response: BrynqResponse) -> None:
        """Record a call for cost reporting."""
        self._total_calls += 1
        self._total_tokens += response.tokens
        self._history.append({
            "model": response.model,
            "strategy": response.strategy,
            "tokens": response.tokens,
            "elapsed_ms": response.elapsed_ms,
            "success": response.success,
        })
