from __future__ import annotations

import os
import json
import time
import logging
from typing import Any, Dict, Optional

import requests

from .base import BaseLLMProvider

logger = logging.getLogger(__name__)

OLLAMA_BASE_URL = os.getenv("OLLAMA_BASE_URL", "http://127.0.0.1:11434")
FALLBACK_MODEL = os.getenv("OLLAMA_FALLBACK_MODEL", "llama3.2:1b-instruct-q4_K_M")
P95_SLO_MS = float(os.getenv("LLM_P95_SLO_MS", "800")) / 1000.0  # Convert to seconds


class LocalOllamaProvider(BaseLLMProvider):
    def __init__(self, model: str = None):
        # Use env var or default to 3B model
        if model is None:
            model = os.getenv("OLLAMA_MODEL", "llama3.2:3b-instruct-q4_K_M")
        super().__init__(model=model)

    def _generate_once(
        self, model: str, prompt: str, schema: Optional[Dict] = None, **opts
    ) -> Dict[str, Any]:
        """Single generation attempt with a specific model"""
        url = f"{OLLAMA_BASE_URL}/api/generate"
        body: Dict[str, Any] = {
            "model": model,
            "prompt": prompt,
            "stream": False,
            "options": {
                "temperature": 0,
                "num_predict": opts.get("num_predict", 64),
                "keep_alive": "30m",
            },
        }

        # Prefer structured outputs if supported; else JSON mode
        if schema:
            body["format"] = schema
        else:
            body["format"] = "json"

        # Use tuple timeout: (connect, read)
        r = requests.post(url, json=body, timeout=(10, 60))

        if r.status_code == 400 and schema:
            # Schema not supported, fallback to JSON mode
            body["format"] = "json"
            r = requests.post(url, json=body, timeout=(10, 60))

        r.raise_for_status()
        data = r.json()
        raw = (data.get("response") or "").strip()

        try:
            value = json.loads(raw) if raw else {}
        except json.JSONDecodeError:
            # Return as heuristic result, not error
            value = {"value": raw}

        return {
            "value": value,
            "provider": "ollama",
            "model": model,
            "tokens": data.get("eval_count", 0),
        }

    def generate_json(self, prompt: str, schema: Optional[Dict] = None, **opts):
        """Generate with automatic fallback on error or slow response"""
        model = opts.get("model") or self.model
        start_time = time.monotonic()
        metrics = {"model": model, "fallback_used": False}

        try:
            # Try primary model
            result = self._generate_once(model, prompt, schema, **opts)
            elapsed = time.monotonic() - start_time

            # Check if response was too slow
            if elapsed > P95_SLO_MS and model != FALLBACK_MODEL:
                metrics["slow_response"] = True
                logger.warning(
                    f"Slow response from {model}: {elapsed:.2f}s > {P95_SLO_MS:.2f}s SLO"
                )

            metrics["latency_ms"] = int(elapsed * 1000)
            result["metrics"] = metrics
            return result

        except Exception as e:
            elapsed = time.monotonic() - start_time
            logger.warning(f"Primary model {model} failed after {elapsed:.2f}s: {e}")

            # Try fallback if not already using it
            if model != FALLBACK_MODEL:
                try:
                    logger.info(f"Attempting fallback to {FALLBACK_MODEL}")
                    metrics["fallback_used"] = True
                    metrics["fallback_reason"] = str(e)[:100]

                    result = self._generate_once(FALLBACK_MODEL, prompt, schema, **opts)
                    elapsed = time.monotonic() - start_time
                    metrics["latency_ms"] = int(elapsed * 1000)
                    result["metrics"] = metrics
                    return result

                except Exception as fallback_error:
                    logger.error(f"Fallback model also failed: {fallback_error}")
                    raise
            raise
