from __future__ import annotations

import json
import logging
import os
import shutil
import subprocess
import time
import urllib.request
from pathlib import Path
from typing import Any, Dict, List, Optional, Tuple

# Configure logging
logging.basicConfig(level=logging.INFO)
logger = logging.getLogger("brynq.api")

class Brynq:
    """
    The Brynq SDK - Intelligent Multi-Model Orchestration.
    
    Usage:
        from brynq import Brynq
        b = Brynq()
        res = b.ask("Explain quantum entanglement")
        print(res)
    """
    
    _BRYNQ_SYSTEM = (
        "You are running inside Brynq, an AI orchestration platform on the user's computer. "
        "You are one of several AI models working together. The user's data stays on their machine. "
        "Brynq can chain multiple models using these strategies: "
        "Sequential (A->B->C with context passing), "
        "Fan-Out (models run in parallel, results merged), "
        "Debate (two models argue, a third judges), "
        "Refine (one drafts, another critiques, first revises), "
        "or Single-Model (one model, simple tasks). "
        "Answer the user's question directly."
    )

    def __init__(self, ollama_url: str = "http://localhost:11434", cloud_url: str = "https://api.brynq.ai"):
        self.ollama_url = ollama_url
        self.cloud_url = cloud_url
        self.active_models: List[str] = []
        self.all_detected: List[str] = []
        self.detect_models()

    def detect_models(self):
        """Auto-detect available local and cloud models."""
        self.active_models = []
        
        # 1. Claude
        if shutil.which("claude") or shutil.which("claude.cmd"):
            self.active_models.append("claude")
            
        # 2. Gemini
        if shutil.which("gemini") or shutil.which("gemini.cmd"):
            self.active_models.append("gemini")
            
        # 3. ChatGPT
        if os.environ.get("OPENAI_API_KEY"):
            self.active_models.append("chatgpt")
            
        # 4. Ollama
        if self._ollama_is_running():
            for m in self._ollama_list_models():
                self.active_models.append(f"ollama:{m}")
                
        self.all_detected = list(self.active_models)
        logger.info(f"Detected models: {', '.join(self.active_models)}")

    def route(self, prompt: str) -> Dict[str, Any]:
        """
        Analyze prompt and return the recommended task type and strategy.
        """
        # Try cloud classification first
        cloud_res = self._classify_cloud(prompt)
        if cloud_res:
            return {
                "task_type": cloud_res[0],
                "strategy": cloud_res[1],
                "reason": cloud_res[2],
                "source": "cloud"
            }
            
        # Fallback to local classification
        local_res = self._classify_local(prompt)
        return {
            "task_type": local_res[0],
            "strategy": local_res[1],
            "reason": local_res[2],
            "source": "local"
        }

    def ask(self, prompt: str, model: Optional[str] = None, strategy: Optional[str] = None) -> str:
        """
        Ask a question. If strategy/model not provided, Brynq routes automatically.
        """
        if not self.active_models:
            return "[Error: No models detected. Install Ollama, Claude, or Gemini CLI.]"
            
        routing = self.route(prompt)
        task_type = routing["task_type"]
        actual_strategy = strategy or routing["strategy"]
        
        if model:
            target_models = [model]
        else:
            target_models = self._pick_models(task_type, actual_strategy)
            
        if not target_models:
            target_models = [self.active_models[0]]

        # For simple SDK, we'll just implement single model execution for now
        # and leave the complex chains for the .chain() method.
        if actual_strategy == "single" or len(target_models) == 1:
            return self._call_model(target_models[0], prompt)
        
        # If a complex strategy is requested in .ask(), we delegate to .chain()
        return self.chain(actual_strategy, prompt, models=target_models)

    def chain(self, strategy: str, prompt: str, models: Optional[List[str]] = None) -> str:
        """
        Run a multi-model orchestration chain.
        """
        if not models:
            routing = self.route(prompt)
            models = self._pick_models(routing["task_type"], strategy)
            
        if not models:
            return "[Error: No models available for chain]"

        if strategy == "debate":
            return self._run_debate(prompt, models)
        elif strategy == "refine":
            return self._run_refine(prompt, models)
        elif strategy == "fan_out":
            return self._run_fan_out(prompt, models)
        else:
            # Default to single
            return self._call_model(models[0], prompt)

    # ── Internal Helpers ───────────────────────────────────────────

    def _ollama_is_running(self) -> bool:
        try:
            req = urllib.request.Request(f"{self.ollama_url}/api/tags", method="GET")
            with urllib.request.urlopen(req, timeout=2.0):
                return True
        except Exception:
            return False

    def _ollama_list_models(self) -> List[str]:
        try:
            req = urllib.request.Request(f"{self.ollama_url}/api/tags", method="GET")
            with urllib.request.urlopen(req, timeout=3.0) as resp:
                data = json.loads(resp.read().decode("utf-8"))
                return [m["name"] for m in data.get("models", [])]
        except Exception:
            return []

    def _call_model(self, model: str, prompt: str) -> str:
        if model == "claude":
            return self._call_claude(prompt)
        elif model == "gemini":
            return self._call_gemini(prompt)
        elif model.startswith("ollama:"):
            return self._call_ollama(prompt, model.split(":", 1)[1])
        return f"[unknown model: {model}]"

    def _call_claude(self, prompt: str) -> str:
        full_prompt = f"[System: {self._BRYNQ_SYSTEM}]\n\nUser: {prompt}"
        
        # Try direct API if key exists
        api_key = os.environ.get("ANTHROPIC_API_KEY")
        if api_key:
            try:
                body = json.dumps({
                    "model": "claude-3-5-sonnet-20240620",
                    "max_tokens": 4096,
                    "messages": [{"role": "user", "content": full_prompt}],
                }).encode("utf-8")
                req = urllib.request.Request(
                    "https://api.anthropic.com/v1/messages",
                    data=body,
                    headers={
                        "Content-Type": "application/json",
                        "x-api-key": api_key,
                        "anthropic-version": "2023-06-01",
                    },
                )
                with urllib.request.urlopen(req, timeout=60) as resp:
                    result = json.loads(resp.read().decode("utf-8"))
                    return result.get("content", [{}])[0].get("text", "")
            except Exception:
                pass

        # Fallback to CLI
        try:
            r = subprocess.run(
                ["claude", "-p", full_prompt],
                capture_output=True, text=True, timeout=None,
                encoding="utf-8", errors="replace",
            )
            if r.returncode == 0:
                return r.stdout.strip()
        except Exception:
            pass
        return "[Claude call failed]"

    def _call_gemini(self, prompt: str) -> str:
        full_prompt = f"[System: {self._BRYNQ_SYSTEM}]\n\n{prompt}"
        cmd = "gemini.cmd" if os.name == "nt" else "gemini"
        try:
            r = subprocess.run(
                [cmd, "-p", full_prompt],
                capture_output=True, text=True, timeout=None,
                encoding="utf-8", errors="replace",
            )
            if r.returncode == 0:
                return r.stdout.strip()
        except Exception:
            pass
        return "[Gemini call failed]"

    def _call_ollama(self, prompt: str, model: str) -> str:
        try:
            body = {
                "model": model,
                "messages": [
                    {"role": "system", "content": self._BRYNQ_SYSTEM},
                    {"role": "user", "content": prompt},
                ],
                "stream": False,
            }
            data = json.dumps(body).encode("utf-8")
            req = urllib.request.Request(
                f"{self.ollama_url}/api/chat",
                data=data,
                headers={"Content-Type": "application/json"},
                method="POST",
            )
            with urllib.request.urlopen(req, timeout=120) as resp:
                result = json.loads(resp.read().decode("utf-8"))
                return result.get("message", {}).get("content", "")
        except Exception as e:
            return f"[Ollama error: {e}]"

    def _classify_cloud(self, text: str) -> Optional[Tuple[str, str, str]]:
        try:
            url = f"{self.cloud_url}/v1/plan"
            body = json.dumps({"input": text, "classify_only": True}).encode("utf-8")
            req = urllib.request.Request(
                url, data=body,
                headers={"Content-Type": "application/json"},
                method="POST",
            )
            with urllib.request.urlopen(req, timeout=3.0) as resp:
                data = json.loads(resp.read().decode("utf-8"))
                t = data.get("task_type", "").strip()
                s = data.get("strategy", "").strip()
                r = data.get("reason", "")
                if t and s:
                    return t, s, r
        except Exception:
            pass
        return None

    def _classify_local(self, text: str) -> Tuple[str, str, str]:
        t = text.lower()
        if "debate" in t or "vs" in t or "better" in t:
            return "compare", "debate", "debate keywords matched"
        if "write" in t or "draft" in t:
            return "creative", "refine", "creative keywords matched"
        if "analyze" in t or "research" in t:
            return "research", "fan_out", "research keywords matched"
        return "general", "single", "defaulting to single"

    def _pick_models(self, task_type: str, strategy: str) -> List[str]:
        if not self.active_models: return []
        
        if strategy == "debate":
            return self.active_models[:3]
        if strategy == "refine":
            return self.active_models[:2]
        if strategy == "fan_out":
            return self.active_models[:3]
        return self.active_models[:1]

    def _run_debate(self, prompt: str, models: List[str]) -> str:
        model_a = models[0]
        model_b = models[1] if len(models) > 1 else models[0]
        judge = models[2] if len(models) > 2 else models[0]
        
        pro = self._call_model(model_a, f"Present a case IN FAVOR of: {prompt}")
        con = self._call_model(model_b, f"Present a case AGAINST: {prompt}")
        
        verdict = self._call_model(judge, f"Topic: {prompt}\n\nPRO:\n{pro}\n\nCON:\n{con}\n\nEvaluate and give a verdict.")
        return verdict

    def _run_refine(self, prompt: str, models: List[str]) -> str:
        drafter = models[0]
        critic = models[1] if len(models) > 1 else models[0]
        
        draft = self._call_model(drafter, f"Draft a response for: {prompt}")
        critique = self._call_model(critic, f"Critically review this draft:\n{draft}")
        final = self._call_model(drafter, f"Revise this draft based on the critique.\nDraft:\n{draft}\nCritique:\n{critique}")
        return final

    def _run_fan_out(self, prompt: str, models: List[str]) -> str:
        # For simple SDK, we'll just run them sequentially for now or use threads
        results = []
        for m in models:
            results.append(self._call_model(m, prompt))
            
        synth = self._call_model(models[0], f"Synthesize these perspectives on '{prompt}':\n\n" + "\n\n".join(results))
        return synth
