#!/usr/bin/env python3
# Copyright (c) 2026 PULSE Contributors. Proprietary. All rights reserved.
"""
Unified LLM Router — single dispatch point for all LLM calls.

Fast cascade (free Ollama -> cheap paid) and Premium cascade (quality first).
Supports Ollama, Anthropic (API key + CLI subscription fallback), and Google.
"""
from __future__ import annotations

import json
import os
import re
import subprocess
import time
from typing import Literal

from pulse.core.provider_state import (
    check_availability as _check_availability,
    reset_availability,
    _has_claude_cli,
)

FAST_CASCADE = [
    ("ollama", "qwen3-coder-next:cloud"),
    ("ollama", "qwen3.5:cloud"),
    ("ollama", "devstral-2:cloud"),
    ("anthropic", "claude-haiku-4-5-20251001"),
    ("google", "gemini-2.0-flash"),
]

PREMIUM_CASCADE = [
    ("ollama", "qwen3-coder:480b-cloud"),
    ("ollama", "qwen3-coder-next:cloud"),
    ("ollama", "qwen3.5:cloud"),
    ("ollama", "kimi-k2.5:cloud"),
    ("ollama", "minimax-m2.5:cloud"),
    ("ollama", "glm-5:cloud"),
    ("ollama", "devstral-2:cloud"),
    ("anthropic", "claude-sonnet-4-6"),
    ("google", "gemini-2.5-pro"),
]
_anthropic_client = None


def _call_anthropic(prompt: str, model: str, max_retries: int = 3, temperature: float = 0.1) -> str:
    """Call Anthropic API (with CLI subscription fallback)."""
    # Try API first if key exists
    if os.environ.get("ANTHROPIC_API_KEY"):
        return _call_anthropic_api(prompt, model, max_retries, temperature)
    # Fall back to CLI subscription
    if _has_claude_cli():
        return _call_anthropic_cli(prompt, model)
    raise RuntimeError("Anthropic: no API key and no CLI subscription")


def _call_anthropic_api(prompt: str, model: str, max_retries: int = 3, temperature: float = 0.1) -> str:
    """Call Anthropic HTTP API with retry on 429/529."""
    import anthropic
    global _anthropic_client
    if _anthropic_client is None:
        _anthropic_client = anthropic.Anthropic()
    for attempt in range(max_retries):
        try:
            resp = _anthropic_client.messages.create(
                model=model, max_tokens=2048, temperature=temperature,
                messages=[{"role": "user", "content": prompt}],
            )
            return resp.content[0].text
        except anthropic.RateLimitError:
            time.sleep(2 ** attempt + 1)
        except anthropic.APIStatusError as e:
            if e.status_code == 529:
                time.sleep(5 * (attempt + 1))
            else:
                raise
    raise RuntimeError(f"Anthropic API {model} failed after {max_retries} retries")


def _call_anthropic_cli(prompt: str, model: str = "") -> str:
    """Call Anthropic via Claude CLI subscription (stdin pipe, print mode)."""
    import shutil
    env = {k: v for k, v in os.environ.items() if k != "CLAUDECODE"}
    # Windows: subprocess can't find .CMD files without shell=True.
    # Resolve full path to avoid shell=True (injection risk).
    claude_bin = shutil.which("claude") or "claude"
    cmd = [claude_bin, "-p", "--output-format", "text"]
    if model:
        cmd.extend(["--model", model])
    result = subprocess.run(
        cmd, input=prompt, capture_output=True, text=True,
        encoding="utf-8", errors="replace", timeout=180, env=env,
    )
    if result.returncode != 0:
        raise RuntimeError(f"Claude CLI error: {result.stderr[:300]}")
    output = result.stdout.strip()
    if not output:
        raise RuntimeError("Claude CLI returned empty response")
    return output


def _call_google(prompt: str, model: str, temperature: float = 0.1) -> str:
    """Call Google Gemini API."""
    api_key = os.environ.get("GEMINI_API_KEY") or os.environ.get("GOOGLE_API_KEY") or ""
    try:
        from google import genai
        from google.genai import types
        client = genai.Client(api_key=api_key)
        resp = client.models.generate_content(
            model=model, contents=prompt,
            config=types.GenerateContentConfig(temperature=temperature),
        )
        return resp.text
    except ImportError:
        import google.generativeai as genai_old
        genai_old.configure(api_key=api_key)
        gm = genai_old.GenerativeModel(model)
        resp = gm.generate_content(prompt, generation_config={"temperature": temperature})
        return resp.text


def _call_ollama(prompt: str, model: str, temperature: float = 0.1) -> str:
    """Call Ollama HTTP API."""
    import urllib.request

    endpoint = os.environ.get("OLLAMA_ENDPOINT", "http://localhost:11434")
    url = f"{endpoint}/api/generate"
    payload = json.dumps({
        "model": model,
        "prompt": prompt,
        "stream": False,
        "options": {"temperature": temperature, "num_predict": 2048},
    }).encode()

    req = urllib.request.Request(
        url, data=payload, headers={"Content-Type": "application/json"}
    )
    with urllib.request.urlopen(req, timeout=120) as resp:
        data = json.loads(resp.read().decode())
    return data.get("response", "")


_PROVIDER_FN = {
    "anthropic": _call_anthropic,
    "google": _call_google,
    "ollama": _call_ollama,
}

def dispatch(
    prompt: str,
    task_type: Literal["fast", "premium"] = "fast",
    max_retries: int = 3,
    temperature: float = 0.1,
) -> str:
    """Send a prompt through the LLM cascade and return the response text.

    Tries each provider in the cascade for the given task_type. Skips
    unavailable providers. Raises RuntimeError if ALL providers fail.

    Args:
        prompt: The full prompt to send.
        task_type: "fast" for cheap/free models, "premium" for quality.
        max_retries: Retries per provider (Anthropic only; others get 1 try).
        temperature: LLM temperature (0.1 for coding, 0.8 for creative/DMN).

    Returns:
        Raw response text from the first successful provider.
    """
    cascade = FAST_CASCADE if task_type == "fast" else PREMIUM_CASCADE
    avail = _check_availability()
    errors: list[str] = []

    for provider, model in cascade:
        if not avail.get(provider, False):
            continue

        try:
            fn = _PROVIDER_FN[provider]
            if provider == "anthropic":
                return fn(prompt, model, max_retries=max_retries, temperature=temperature)
            else:
                return fn(prompt, model, temperature=temperature)
        except Exception as e:
            errors.append(f"{provider}/{model}: {e}")

    raise RuntimeError(
        f"All providers failed for task_type={task_type}. "
        f"Errors: {'; '.join(errors) or 'No providers available'}"
    )


def dispatch_to(
    provider: str, model: str, prompt: str,
    temperature: float = 0.1, max_retries: int = 3,
) -> str:
    """Send a prompt to a specific provider/model (no cascade fallback)."""
    avail = _check_availability()
    if not avail.get(provider, False):
        raise RuntimeError(f"Provider '{provider}' not available")

    fn = _PROVIDER_FN.get(provider)
    if not fn:
        raise RuntimeError(f"Unknown provider '{provider}'")

    if provider == "anthropic":
        return fn(prompt, model, max_retries=max_retries, temperature=temperature)
    return fn(prompt, model, temperature=temperature)


def extract_json(raw: str) -> dict:
    """Parse JSON from LLM response, handling markdown fences and noise."""
    raw = raw.strip()

    # Strip markdown code fences
    if raw.startswith("```"):
        lines = raw.split("\n")
        lines = [l for l in lines if not l.strip().startswith("```")]
        raw = "\n".join(lines)

    # Try direct parse
    try:
        return json.loads(raw)
    except json.JSONDecodeError:
        pass

    # Try to find JSON object with expected keys
    match = re.search(
        r'\{[^{}]*"(?:lessons|failures|patterns)"[^{}]*\}', raw, re.DOTALL
    )
    if match:
        try:
            return json.loads(match.group())
        except json.JSONDecodeError:
            pass

    # Try outermost braces
    start = raw.find("{")
    end = raw.rfind("}")
    if start != -1 and end > start:
        try:
            return json.loads(raw[start : end + 1])
        except json.JSONDecodeError:
            pass

    return {"lessons": [], "failures": [], "patterns": []}


def dispatch_json(
    prompt: str,
    task_type: Literal["fast", "premium"] = "fast",
) -> dict:
    """Dispatch + parse JSON in one call."""
    return extract_json(dispatch(prompt, task_type=task_type))
