#!/usr/bin/env python3
"""
backends.py - LLM Backend Abstraction (WO#GINNIE_02)
=====================================================

Abstracts different LLM backends:
- OpenAI (GPT-4, GPT-4o-mini)
- Anthropic (Claude)
- Local GGUF (llama.cpp)

All backends use JazzClient for HTTP calls (WO#GINNIE_07).
"""

import json
import os
from abc import ABC, abstractmethod
from typing import Any, Dict, List, Optional

from helix_cdc.ginnie.config import GinnieConfig, BackendConfig


class BaseBackend(ABC):
    """Abstract base class for LLM backends."""

    def __init__(self, config: BackendConfig):
        self.config = config
        self._jazz_client = None

    @property
    def jazz_client(self):
        """Lazy load Jazz client."""
        if self._jazz_client is None:
            try:
                from helix_cdc.jazz import JazzClient
                self._jazz_client = JazzClient(valley_threshold=40.0, max_wait=0.3)
            except ImportError:
                self._jazz_client = None
        return self._jazz_client

    @abstractmethod
    def complete(self, messages: List[Dict[str, str]], **kwargs) -> str:
        """
        Complete a chat conversation.

        Args:
            messages: List of message dicts with 'role' and 'content'
            **kwargs: Additional parameters (temperature, max_tokens, etc.)

        Returns:
            Generated text response
        """
        pass


class OpenAIBackend(BaseBackend):
    """OpenAI API backend."""

    ENDPOINT = "https://api.openai.com/v1/chat/completions"

    def complete(self, messages: List[Dict[str, str]], **kwargs) -> str:
        if not self.config.api_key:
            raise ValueError("OpenAI API key not configured")

        headers = {
            "Authorization": f"Bearer {self.config.api_key}",
            "Content-Type": "application/json",
        }

        payload = {
            "model": self.config.model or "gpt-4o-mini",
            "messages": messages,
            "temperature": kwargs.get("temperature", 0.7),
            "max_tokens": kwargs.get("max_tokens", 2048),
        }

        # Use JazzClient if available, else fall back to urllib
        if self.jazz_client:
            resp = self.jazz_client.post(self.ENDPOINT, headers=headers, json=payload)
            if resp.status_code != 200:
                raise Exception(f"OpenAI API error: {resp.status_code} - {resp.text()[:200]}")
            data = resp.json()
        else:
            import urllib.request
            import ssl
            ctx = ssl.create_default_context()
            req = urllib.request.Request(
                self.ENDPOINT,
                data=json.dumps(payload).encode(),
                headers=headers,
                method="POST",
            )
            with urllib.request.urlopen(req, timeout=60, context=ctx) as resp:
                data = json.load(resp)

        return data["choices"][0]["message"]["content"]


class AnthropicBackend(BaseBackend):
    """Anthropic API backend."""

    ENDPOINT = "https://api.anthropic.com/v1/messages"

    def complete(self, messages: List[Dict[str, str]], **kwargs) -> str:
        if not self.config.api_key:
            raise ValueError("Anthropic API key not configured")

        # Convert messages format (Anthropic uses different structure)
        system_content = ""
        user_messages = []
        for msg in messages:
            if msg["role"] == "system":
                system_content = msg["content"]
            else:
                user_messages.append(msg)

        headers = {
            "x-api-key": self.config.api_key,
            "Content-Type": "application/json",
            "anthropic-version": "2023-06-01",
        }

        payload = {
            "model": self.config.model or "claude-sonnet-4-20250514",
            "messages": user_messages,
            "max_tokens": kwargs.get("max_tokens", 2048),
        }

        if system_content:
            payload["system"] = system_content

        # Use JazzClient if available
        if self.jazz_client:
            resp = self.jazz_client.post(self.ENDPOINT, headers=headers, json=payload)
            if resp.status_code != 200:
                raise Exception(f"Anthropic API error: {resp.status_code} - {resp.text()[:200]}")
            data = resp.json()
        else:
            import urllib.request
            import ssl
            ctx = ssl.create_default_context()
            req = urllib.request.Request(
                self.ENDPOINT,
                data=json.dumps(payload).encode(),
                headers=headers,
                method="POST",
            )
            with urllib.request.urlopen(req, timeout=60, context=ctx) as resp:
                data = json.load(resp)

        return data["content"][0]["text"]


class LocalGGUFBackend(BaseBackend):
    """Local GGUF model backend via llama.cpp with FibPi3D timing."""

    LLAMA_CLI = "/home/voidstr3m33/llama.cpp/build/bin/llama-cli"

    # Prompt format templates
    PROMPT_FORMATS = {
        "mistral": {
            "system_prefix": "[INST] <<SYS>>\n",
            "system_suffix": "\n<</SYS>>\n",
            "user_prefix": "",
            "user_suffix": " [/INST]",
            "assistant_prefix": "",
            "assistant_suffix": "\n[INST] ",
        },
        "chatml": {
            # Qwen2.5, Phi-3, and other ChatML models
            "system_prefix": "<|im_start|>system\n",
            "system_suffix": "<|im_end|>\n",
            "user_prefix": "<|im_start|>user\n",
            "user_suffix": "<|im_end|>\n",
            "assistant_prefix": "<|im_start|>assistant\n",
            "assistant_suffix": "<|im_end|>\n",
        },
        "llama2": {
            "system_prefix": "[INST] <<SYS>>\n",
            "system_suffix": "\n<</SYS>>\n\n",
            "user_prefix": "",
            "user_suffix": " [/INST] ",
            "assistant_prefix": "",
            "assistant_suffix": " </s><s>[INST] ",
        },
    }

    def __init__(self, config):
        super().__init__(config)
        self._combustion_sampler = None
        # Detect prompt format from config or default to mistral
        self._prompt_format = getattr(config, 'prompt_format', None) or "mistral"

    @property
    def combustion_sampler(self):
        """Lazy load combustion sampler for valley timing."""
        if self._combustion_sampler is None:
            try:
                from helix_cdc.jazz import CombustionSampler
                self._combustion_sampler = CombustionSampler(valley_threshold=40.0)
            except ImportError:
                self._combustion_sampler = None
        return self._combustion_sampler

    def _wait_for_valley(self) -> tuple:
        """Wait for CPU valley before firing inference."""
        try:
            from helix_cdc.jazz import fire_on_valley, get_cpu_load
            import time
            start = time.perf_counter()
            cpu = get_cpu_load()

            # Wait for valley (max 0.5s)
            while cpu > 40.0 and (time.perf_counter() - start) < 0.5:
                time.sleep(0.02)
                cpu = get_cpu_load()

            wait_ms = (time.perf_counter() - start) * 1000
            valley_found = cpu <= 40.0
            return valley_found, wait_ms, cpu
        except ImportError:
            return True, 0, 0  # No jazz available, proceed anyway

    def complete(self, messages: List[Dict[str, str]], **kwargs) -> str:
        import subprocess
        import tempfile

        if not self.config.model_path or not self.config.model_path.exists():
            raise ValueError(f"Local model not found: {self.config.model_path}")

        # FibPi3D: Wait for CPU valley before inference
        valley_found, wait_ms, cpu_at_fire = self._wait_for_valley()

        # Build prompt from messages using detected prompt format
        prompt = self._build_prompt(messages)

    def _build_prompt(self, messages: List[Dict[str, str]]) -> str:
        """
        Build a prompt string from messages using the configured prompt format.

        Supports: mistral, chatml (Qwen2.5, Phi-3), llama2
        """
        fmt = self.PROMPT_FORMATS.get(self._prompt_format, self.PROMPT_FORMATS["mistral"])

        prompt_parts = []
        for msg in messages:
            role = msg["role"]
            content = msg["content"]
            if role == "system":
                prompt_parts.append(f"{fmt['system_prefix']}{content}{fmt['system_suffix']}")
            elif role == "user":
                prompt_parts.append(f"{fmt['user_prefix']}{content}{fmt['user_suffix']}")
            elif role == "assistant":
                prompt_parts.append(f"{fmt['assistant_prefix']}{content}{fmt['assistant_suffix']}")

        # For ChatML, add the assistant prompt to trigger generation
        if self._prompt_format == "chatml":
            prompt_parts.append("<|im_start|>assistant\n")

        return "".join(prompt_parts)

    def complete(self, messages: List[Dict[str, str]], **kwargs) -> str:
        import subprocess
        import tempfile

        if not self.config.model_path or not self.config.model_path.exists():
            raise ValueError(f"Local model not found: {self.config.model_path}")

        # FibPi3D: Wait for CPU valley before inference
        valley_found, wait_ms, cpu_at_fire = self._wait_for_valley()

        # Build prompt from messages using detected prompt format
        prompt = self._build_prompt(messages)

        # Write prompt to temp file
        with tempfile.NamedTemporaryFile(mode='w', suffix='.txt', delete=False) as f:
            f.write(prompt)
            prompt_file = f.name

        try:
            max_tokens = kwargs.get("max_tokens", 512)
            cmd = [
                self.LLAMA_CLI,
                "-m", str(self.config.model_path),
                "-f", prompt_file,
                "-n", str(max_tokens),
                "--temp", str(kwargs.get("temperature", 0.7)),
                "-c", "4096",
                "--no-display-prompt",
            ]

            result = subprocess.run(
                cmd,
                capture_output=True,
                text=True,
                timeout=120,
            )

            if result.returncode != 0:
                raise Exception(f"llama-cli error: {result.stderr[:500]}")

            return result.stdout.strip()

        finally:
            import os
            os.unlink(prompt_file)


class MockBackend(BaseBackend):
    """Mock backend for testing."""

    def complete(self, messages: List[Dict[str, str]], **kwargs) -> str:
        last_user_msg = ""
        for msg in reversed(messages):
            if msg["role"] == "user":
                last_user_msg = msg["content"]
                break

        # Return a simple mock response
        return json.dumps({
            "steps": [
                {"action": "analyze", "target": "mock", "details": f"Mock analysis of: {last_user_msg[:50]}..."}
            ]
        })


def get_backend(name: str, config: GinnieConfig) -> BaseBackend:
    """
    Get a backend instance by name.

    Args:
        name: Backend name (openai, anthropic, local, mock)
        config: Ginnie configuration

    Returns:
        Backend instance
    """
    backend_config = config.get_backend(name)

    if not backend_config:
        # Fall back to mock
        return MockBackend(BackendConfig(name="mock", type="mock"))

    backend_type = backend_config.type.lower()

    if backend_type == "openai":
        return OpenAIBackend(backend_config)
    elif backend_type == "anthropic":
        return AnthropicBackend(backend_config)
    elif backend_type == "local_gguf":
        return LocalGGUFBackend(backend_config)
    else:
        return MockBackend(backend_config)
