"""Calibration runner — orchestrates benchmark execution and model vector learning."""

from __future__ import annotations

import asyncio
import json
import logging
import random
import time
from dataclasses import dataclass, field
from pathlib import Path

import aiohttp
import numpy as np

from backend.bridge.base import BridgeRequest
from backend.bridge.factory import create_bridge
from backend.calibration.judge import judge_response
from backend.calibration.tasks import BenchmarkTask, get_all_tasks
from backend.calibration.vector_learner import learn_model_vector
from backend.config import EcateConfig, ModelConfig, ProviderConfig
from backend.db.connection import get_db
from backend.router.embedder import get_embedder

logger = logging.getLogger("ecate.calibration")

# Calibration constants
PASS_THRESHOLD = 3.5  # Composite score threshold for pass/fail (70% of 5.0)
MIN_SAMPLES_FOR_VECTOR = 10  # Minimum valid results to learn a vector


@dataclass
class TaskResult:
    """Result of running a single task against a model."""
    task_id: str
    model_id: str
    category: str
    response_text: str
    correctness: float
    completeness: float
    following: float
    composite_score: float
    passed: bool  # Binary pass/fail for vector learning
    latency_ms: float
    tokens_in: int
    tokens_out: int
    cost_usd: float
    judge_reasoning: str
    task_embedding: list[float] | None = None  # Cached embedding for vector learning


@dataclass
class CalibrationProgress:
    """Tracks calibration run progress."""
    total_tasks: int = 0
    completed_tasks: int = 0
    current_model: str = ""
    current_task: str = ""
    status: str = "idle"  # idle | running | completed | failed
    results: list[TaskResult] = field(default_factory=list)
    error: str | None = None
    run_id: int | None = None
    vectors_learned: dict[str, float] = field(default_factory=dict)  # model_id -> accuracy


# Global progress tracker
_progress = CalibrationProgress()


def get_progress() -> CalibrationProgress:
    return _progress


def _embed_tasks(tasks: list[BenchmarkTask]) -> dict[str, np.ndarray]:
    """Pre-embed all task prompts for vector learning."""
    embedder = get_embedder()
    prompts = [t.prompt for t in tasks]
    embeddings = embedder.embed_batch(prompts, show_progress=True)
    return {t.id: embeddings[i] for i, t in enumerate(tasks)}


def _filter_tasks(tasks: list[BenchmarkTask], config: EcateConfig) -> list[BenchmarkTask]:
    """Apply calibration config filters to task list."""
    cal = config.calibration
    filtered = tasks

    # Filter by categories (include list)
    if cal.categories:
        filtered = [t for t in filtered if t.category in cal.categories]

    # Filter by exclude_categories
    if cal.exclude_categories:
        filtered = [t for t in filtered if t.category not in cal.exclude_categories]

    # Filter by languages
    if cal.languages:
        filtered = [t for t in filtered if t.language in cal.languages]

    # Filter by difficulties
    if cal.difficulties:
        filtered = [t for t in filtered if t.difficulty in cal.difficulties]

    # Limit max tasks (random sample to maintain category distribution)
    if cal.max_tasks and len(filtered) > cal.max_tasks:
        random.shuffle(filtered)
        filtered = filtered[:cal.max_tasks]

    if not filtered:
        logger.warning("No tasks match the calibration filters")

    logger.info(f"Calibration task filter: {len(tasks)} → {len(filtered)} tasks")
    return filtered


async def run_calibration(
    config: EcateConfig,
    session: aiohttp.ClientSession,
    model_ids: list[str] | None = None,
) -> CalibrationProgress:
    """Run calibration against models to learn routing vectors.

    1. Loads benchmark tasks and pre-embeds them
    2. Sends each task to each model (respects model.calibration_enabled)
    3. Evaluates responses with LLM-as-judge (binary pass/fail)
    4. Learns model vectors from pass/fail results
    5. Saves vectors to data/model_vectors.json
    """
    global _progress
    _progress = CalibrationProgress(status="running")

    try:
        # Load built-in + custom tasks
        tasks = get_all_tasks(config.calibration.custom_tasks_dir)
        tasks = _filter_tasks(tasks, config)
        
        # Pre-embed all tasks for vector learning
        logger.info("Pre-embedding tasks for vector learning...")
        task_embeddings = _embed_tasks(tasks)
        
        all_models = config.all_models()

        # Filter models: only those with calibration_enabled=True
        all_models = [
            (p, m) for p, m in all_models
            if m.calibration_enabled
        ]
        
        # Further filter if specific model_ids provided
        if model_ids:
            all_models = [
                (p, m) for p, m in all_models
                if f"{p.id}/{m.id}" in model_ids or m.id in model_ids
            ]

        if not all_models:
            _progress.status = "failed"
            _progress.error = "No models to calibrate (check calibration_enabled)"
            return _progress

        # Limit tasks per model based on calibration_samples
        _progress.total_tasks = sum(
            min(len(tasks), m.calibration_samples) for _, m in all_models
        )

        # Resolve judge model
        judge_resolved = config.get_model(config.calibration.judge_model)
        if not judge_resolved:
            _progress.status = "failed"
            _progress.error = f"Judge model '{config.calibration.judge_model}' not found in config"
            return _progress

        judge_provider, judge_model = judge_resolved

        # Create calibration run in DB
        db = await get_db()
        cursor = await db.execute(
            "INSERT INTO calibration_runs (started_at, judge_model, models_tested, tasks_count, status) "
            "VALUES (datetime('now'), ?, ?, ?, 'running')",
            (
                config.calibration.judge_model,
                json.dumps([f"{p.id}/{m.id}" for p, m in all_models]),
                len(tasks),
            ),
        )
        await db.commit()
        run_id = cursor.lastrowid
        _progress.run_id = run_id

        # Semaphore for rate limiting per provider
        semaphores: dict[str, asyncio.Semaphore] = {
            p.id: asyncio.Semaphore(config.calibration.concurrency_per_provider)
            for p, _ in all_models
        }

        total_cost = 0.0
        
        # Results grouped by model for vector learning
        model_results: dict[str, list[TaskResult]] = {}

        for provider, model in all_models:
            model_id = f"{provider.id}/{model.id}"
            _progress.current_model = model_id
            bridge = create_bridge(provider, session)
            
            # Sample tasks for this model
            n_samples = min(len(tasks), model.calibration_samples)
            model_tasks = random.sample(tasks, n_samples) if n_samples < len(tasks) else tasks
            model_results[model_id] = []

            for task in model_tasks:
                _progress.current_task = task.id

                async with semaphores[provider.id]:
                    try:
                        result = await _run_single_task(
                            bridge, provider, model, task,
                            task_embeddings.get(task.id),
                            session, judge_provider, judge_model, config,
                        )
                        _progress.results.append(result)
                        model_results[model_id].append(result)
                        total_cost += result.cost_usd

                        # Store result in DB
                        await db.execute(
                            """INSERT INTO calibration_results
                               (run_id, model_id, task_id, category, response_text,
                                correctness, completeness, following, composite_score,
                                latency_ttfb_ms, latency_total_ms, tokens_in, tokens_out,
                                cost_usd, judge_reasoning)
                               VALUES (?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?)""",
                            (
                                run_id, model_id, task.id, task.category,
                                result.response_text[:2000],
                                result.correctness, result.completeness,
                                result.following, result.composite_score,
                                result.latency_ms, result.latency_ms,
                                result.tokens_in, result.tokens_out,
                                result.cost_usd, result.judge_reasoning,
                            ),
                        )
                        await db.commit()

                    except Exception as e:
                        logger.error(f"Task {task.id} failed on {model_id}: {e}")

                _progress.completed_tasks += 1

        # Learn model vectors from results
        logger.info("Learning model vectors from calibration results...")
        vectors_path = Path(config.routing.vectors_path)
        await _learn_and_save_vectors(model_results, vectors_path, _progress)

        # Update run status
        await db.execute(
            "UPDATE calibration_runs SET completed_at = datetime('now'), total_cost_usd = ?, status = 'completed' WHERE id = ?",
            (total_cost, run_id),
        )
        await db.commit()

        _progress.status = "completed"
        logger.info(f"Calibration complete: {_progress.completed_tasks} tasks, ${total_cost:.2f}")

    except Exception as e:
        logger.error(f"Calibration failed: {e}")
        _progress.status = "failed"
        _progress.error = str(e)

    return _progress


async def run_single_model_calibration(
    config: EcateConfig,
    session: aiohttp.ClientSession,
    model_id: str,
    n_samples: int = 200,
) -> tuple[np.ndarray, float] | None:
    """Calibrate a single model and learn its vector.
    
    Use this to add a new model without full recalibration.
    
    Args:
        config: Ecate configuration
        session: HTTP session
        model_id: Model to calibrate (e.g., "anthropic/claude-sonnet")
        n_samples: Number of tasks to evaluate
    
    Returns:
        (model_vector, accuracy) or None if failed
    """
    # Resolve model
    resolved = config.get_model(model_id)
    if not resolved:
        logger.error(f"Model {model_id} not found")
        return None
    
    provider, model = resolved
    full_id = f"{provider.id}/{model.id}"
    
    # Load tasks
    tasks = get_all_tasks(config.calibration.custom_tasks_dir)
    tasks = _filter_tasks(tasks, config)
    
    if len(tasks) < n_samples:
        logger.warning(f"Only {len(tasks)} tasks available, using all")
        n_samples = len(tasks)
    
    # Sample and embed tasks
    sample_tasks = random.sample(tasks, n_samples)
    task_embeddings = _embed_tasks(sample_tasks)
    
    # Resolve judge
    judge_resolved = config.get_model(config.calibration.judge_model)
    if not judge_resolved:
        logger.error(f"Judge model not found")
        return None
    
    judge_provider, judge_model = judge_resolved
    bridge = create_bridge(provider, session)
    
    # Run evaluation
    results: list[TaskResult] = []
    for i, task in enumerate(sample_tasks):
        try:
            result = await _run_single_task(
                bridge, provider, model, task,
                task_embeddings.get(task.id),
                session, judge_provider, judge_model, config,
            )
            results.append(result)
            if (i + 1) % 20 == 0:
                logger.info(f"Progress: {i + 1}/{n_samples}")
        except Exception as e:
            logger.error(f"Task {task.id} failed: {e}")
    
    if len(results) < 10:
        logger.error(f"Not enough results ({len(results)}) to learn vector")
        return None
    
    # Learn vector
    valid_results = [r for r in results if r.task_embedding is not None]
    embeddings = np.array([r.task_embedding for r in valid_results])
    pass_fail = np.array([1.0 if r.passed else 0.0 for r in valid_results])
    
    vector, accuracy = learn_model_vector(embeddings, pass_fail)
    logger.info(f"Learned vector for {full_id}: accuracy={accuracy:.3f}")
    
    # Save to vectors file
    vectors_path = Path(config.routing.vectors_path)
    _save_single_vector(vectors_path, full_id, vector)
    
    return vector, accuracy


def _save_single_vector(vectors_path: Path, model_id: str, vector: np.ndarray) -> None:
    """Save a single model vector to the vectors file."""
    embedder = get_embedder()
    
    # Load existing
    existing = {}
    if vectors_path.exists():
        with open(vectors_path) as f:
            data = json.load(f)
            existing = data.get("vectors", {})
    
    # Add/update
    existing[model_id] = vector.tolist()
    
    # Save
    vectors_path.parent.mkdir(parents=True, exist_ok=True)
    data = {
        "version": 1,
        "embed_dim": embedder.dim,
        "embedder_model": embedder._model_name,
        "vectors": existing,
    }
    with open(vectors_path, "w") as f:
        json.dump(data, f, indent=2)
    
    logger.info(f"Saved vector for {model_id} to {vectors_path}")


async def _run_single_task(
    bridge,
    provider: ProviderConfig,
    model: ModelConfig,
    task: BenchmarkTask,
    task_embedding: np.ndarray | None,
    session: aiohttp.ClientSession,
    judge_provider: ProviderConfig,
    judge_model: ModelConfig,
    config: EcateConfig,
) -> TaskResult:
    """Run a single benchmark task against a model and evaluate the response."""
    model_id = f"{provider.id}/{model.id}"

    # Send task to model
    start = time.monotonic()
    request = BridgeRequest(
        messages=[{"role": "user", "content": task.prompt}],
        model=model.id,
        temperature=0,
        max_tokens=1024,
    )

    response = await bridge.complete(request, model)
    latency = (time.monotonic() - start) * 1000

    # Calculate cost
    cost = (
        response.tokens_in * model.input_cost_per_mtok / 1_000_000
        + response.tokens_out * model.output_cost_per_mtok / 1_000_000
    )

    # Judge the response
    score = await judge_response(
        session=session,
        judge_base_url=judge_provider.base_url,
        judge_api_key=judge_provider.api_key,
        judge_model=judge_model.id,
        judge_provider_type=judge_provider.type,
        task_prompt=task.prompt,
        model_response=response.content,
        expected_behavior=task.expected_behavior,
    )

    # Add judge cost estimate
    judge_cost = (judge_model.input_cost_per_mtok * 500 + judge_model.output_cost_per_mtok * 200) / 1_000_000
    cost += judge_cost
    
    # Binary pass/fail
    passed = score.composite >= PASS_THRESHOLD

    return TaskResult(
        task_id=task.id,
        model_id=model_id,
        category=task.category,
        response_text=response.content,
        correctness=score.correctness,
        completeness=score.completeness,
        following=score.following,
        composite_score=score.composite,
        passed=passed,
        latency_ms=latency,
        tokens_in=response.tokens_in,
        tokens_out=response.tokens_out,
        cost_usd=cost,
        judge_reasoning=score.reasoning,
        task_embedding=task_embedding.tolist() if task_embedding is not None else None,
    )


async def _learn_and_save_vectors(
    model_results: dict[str, list[TaskResult]],
    vectors_path: Path,
    progress: CalibrationProgress,
) -> None:
    """Learn model vectors from calibration results and save to file."""
    
    # Load existing vectors (to preserve models not in this run)
    existing_vectors: dict[str, list[float]] = {}
    if vectors_path.exists():
        try:
            with open(vectors_path) as f:
                data = json.load(f)
                existing_vectors = data.get("vectors", {})
        except Exception as e:
            logger.warning(f"Could not load existing vectors: {e}")
    
    embedder = get_embedder()
    new_vectors: dict[str, list[float]] = {}
    
    for model_id, results in model_results.items():
        if not results:
            continue
        
        # Filter results with embeddings
        valid_results = [r for r in results if r.task_embedding is not None]
        if len(valid_results) < MIN_SAMPLES_FOR_VECTOR:
            logger.warning(f"Not enough valid results for {model_id} ({len(valid_results)}), need {MIN_SAMPLES_FOR_VECTOR}")
            continue
        
        # Build embedding matrix and pass/fail vector
        embeddings = np.array([r.task_embedding for r in valid_results])
        pass_fail = np.array([1.0 if r.passed else 0.0 for r in valid_results])
        
        # Learn vector
        try:
            vector, accuracy = learn_model_vector(embeddings, pass_fail)
            new_vectors[model_id] = vector.tolist()
            progress.vectors_learned[model_id] = accuracy
            logger.info(f"Learned vector for {model_id}: accuracy={accuracy:.3f}")
        except Exception as e:
            logger.error(f"Failed to learn vector for {model_id}: {e}")
    
    # Merge with existing vectors
    all_vectors = {**existing_vectors, **new_vectors}
    
    # Save
    vectors_path.parent.mkdir(parents=True, exist_ok=True)
    data = {
        "version": 1,
        "embed_dim": embedder.dim,
        "embedder_model": embedder._model_name,
        "vectors": all_vectors,
    }
    with open(vectors_path, "w") as f:
        json.dump(data, f, indent=2)
    
    logger.info(f"Saved {len(all_vectors)} model vectors to {vectors_path}")
