"""Post-training evaluation integration.

This module provides a clean interface for running comprehensive evaluations
after training completes or at checkpoints during training.

Usage:
    from embergraphrouter.evaluation import run_post_training_eval

    # After training
    results = run_post_training_eval(
        model_path="model_path/my_model.pth",
        data_path="data/router_data.parquet",
        config=training_config,
    )
"""

from __future__ import annotations

import json
import logging
import subprocess
import sys
from dataclasses import dataclass, field
from pathlib import Path
from typing import Any

logger = logging.getLogger(__name__)


@dataclass
class PostTrainingEvalConfig:
    """Configuration for post-training evaluation."""

    # Whether to run evaluation
    enabled: bool = True

    # Run advanced analysis (error, cost, significance)
    run_advanced: bool = True

    # Print scorecard summary
    print_scorecard: bool = True

    # Save results to JSON
    save_results: bool = True

    # Output directory (defaults to model directory)
    output_dir: str | None = None


@dataclass
class EvalResults:
    """Results from post-training evaluation."""

    # Core metrics
    frac_oracle: float = 0.0
    top1_accuracy: float = 0.0
    regret: float = 0.0
    ece: float = 0.0

    # Difficulty detection
    difficulty_auc: float = 0.0
    mean_logit_corr: float = 0.0

    # Ranking
    spearman_corr: float = 0.0
    llm_discrimination: float = 0.0

    # Full results dict
    full_results: dict = field(default_factory=dict)

    # Paths
    results_path: str | None = None


def run_post_training_eval(
    model_path: str | Path,
    data_path: str | Path,
    config: dict[str, Any],
    eval_config: PostTrainingEvalConfig | None = None,
    llm_embedding_path: str | Path | None = None,
) -> EvalResults | None:
    """Run comprehensive post-training evaluation.

    This function calls the eval_difficulty_aware script with the trained model
    and returns structured results.

    Args:
        model_path: Path to trained model weights (.pth)
        data_path: Path to router data (parquet or csv)
        config: Training config dict (used to extract LLM info)
        eval_config: Evaluation configuration
        llm_embedding_path: Path to LLM embeddings (optional, uses config if not provided)

    Returns:
        EvalResults with metrics, or None if evaluation failed
    """
    if eval_config is None:
        eval_config = PostTrainingEvalConfig()

    if not eval_config.enabled:
        logger.info("Post-training evaluation disabled")
        return None

    model_path = Path(model_path)
    data_path = Path(data_path)

    if not model_path.exists():
        logger.warning("Model path does not exist: %s", model_path)
        return None

    if not data_path.exists():
        logger.warning("Data path does not exist: %s", data_path)
        return None

    # Determine output path
    if eval_config.output_dir:
        output_dir = Path(eval_config.output_dir)
    else:
        output_dir = model_path.parent

    output_json = output_dir / f"{model_path.stem}_eval_results.json"

    logger.info("=" * 70)
    logger.info("POST-TRAINING EVALUATION")
    logger.info("=" * 70)
    logger.info("Model: %s", model_path)
    logger.info("Data: %s", data_path)

    # Build the evaluation command
    script_dir = Path(__file__).parent.parent / "scripts"
    eval_script = script_dir / "eval_difficulty_aware.py"

    if not eval_script.exists():
        logger.error("Evaluation script not found: %s", eval_script)
        return None

    # Create a temporary checkpoint directory structure for the evaluator
    # The eval script expects a registry-style checkpoint with model.pth and config
    temp_checkpoint = _create_temp_checkpoint(model_path, config, llm_embedding_path)

    cmd = [
        sys.executable,
        str(eval_script),
        "--checkpoint", str(temp_checkpoint),
        "--data", str(data_path),
    ]

    if eval_config.run_advanced:
        cmd.append("--advanced")

    if eval_config.print_scorecard:
        cmd.append("--scorecard")

    if eval_config.save_results:
        cmd.extend(["--output", str(output_json)])

    try:
        logger.info("Running: %s", " ".join(cmd))
        result = subprocess.run(
            cmd,
            capture_output=True,
            text=True,
            timeout=600,  # 10 minute timeout
        )

        if result.returncode != 0:
            logger.error("Evaluation failed with code %d", result.returncode)
            logger.error("stderr: %s", result.stderr[-2000:] if result.stderr else "")
            # Still try to parse output
            if result.stdout:
                logger.info("stdout:\n%s", result.stdout[-3000:])
            return None

        # Log output
        if result.stdout:
            # Print key lines from output
            for line in result.stdout.split("\n"):
                if any(k in line.lower() for k in ["frac_oracle", "auc", "spearman", "ece", "key insight"]):
                    logger.info(line)

        # Parse results if saved
        eval_results = EvalResults()
        if eval_config.save_results and output_json.exists():
            with open(output_json) as f:
                full_results = json.load(f)

            eval_results.full_results = full_results
            eval_results.results_path = str(output_json)

            # Extract key metrics
            if "routing_metrics" in full_results:
                rm = full_results["routing_metrics"]
                eval_results.frac_oracle = rm.get("frac_of_oracle", 0)
                eval_results.top1_accuracy = rm.get("top1_acc", 0)
                eval_results.regret = rm.get("regret", 0)

            if "calibration_metrics" in full_results:
                cm = full_results["calibration_metrics"]
                eval_results.ece = cm.get("ece", 0)

            if "difficulty_detection" in full_results:
                dd = full_results["difficulty_detection"]
                eval_results.difficulty_auc = dd.get("trivial_impossible_auc", 0)
                eval_results.mean_logit_corr = dd.get("mean_logit_ncorrect_corr", 0)

            if "ranking_discrimination" in full_results:
                rd = full_results["ranking_discrimination"]
                eval_results.spearman_corr = rd.get("avg_spearman", 0)
                eval_results.llm_discrimination = rd.get("avg_llm_discrimination", 0)

            logger.info("Results saved to: %s", output_json)

        # Clean up temp checkpoint
        _cleanup_temp_checkpoint(temp_checkpoint)

        logger.info("=" * 70)
        logger.info("EVALUATION SUMMARY")
        logger.info("=" * 70)
        logger.info("  frac_oracle:     %.3f", eval_results.frac_oracle)
        logger.info("  top1_accuracy:   %.3f", eval_results.top1_accuracy)
        logger.info("  ECE:             %.3f", eval_results.ece)
        logger.info("  difficulty_AUC:  %.3f", eval_results.difficulty_auc)
        logger.info("  mean_logit_corr: %.3f", eval_results.mean_logit_corr)
        logger.info("  spearman_corr:   %.3f", eval_results.spearman_corr)
        logger.info("=" * 70)

        return eval_results

    except subprocess.TimeoutExpired:
        logger.error("Evaluation timed out after 600 seconds")
        _cleanup_temp_checkpoint(temp_checkpoint)
        return None
    except Exception as e:
        logger.error("Evaluation failed: %s", e)
        _cleanup_temp_checkpoint(temp_checkpoint)
        return None


def _create_temp_checkpoint(
    model_path: Path,
    config: dict[str, Any],
    llm_embedding_path: str | Path | None = None,
) -> Path:
    """Create a temporary checkpoint directory for the evaluator.

    The eval_difficulty_aware script expects a registry-style checkpoint with:
    - model.pth (or the actual model file)
    - config.yaml

    Args:
        model_path: Path to trained model weights
        config: Training config dict
        llm_embedding_path: Optional path to LLM embeddings

    Returns:
        Path to temporary checkpoint directory
    """
    import shutil
    import tempfile
    import yaml

    # Create temp directory
    temp_dir = Path(tempfile.mkdtemp(prefix="eval_checkpoint_"))

    # Copy or link model file
    model_dest = temp_dir / "model.pth"
    if model_path.exists():
        shutil.copy2(model_path, model_dest)

    # Write config
    config_dest = temp_dir / "config.yaml"
    with open(config_dest, "w") as f:
        yaml.dump(config, f)

    # Copy LLM embeddings if provided
    if llm_embedding_path:
        llm_path = Path(llm_embedding_path)
        if llm_path.exists():
            shutil.copy2(llm_path, temp_dir / "llm_embeddings.pkl")

    logger.debug("Created temp checkpoint at: %s", temp_dir)
    return temp_dir


def _cleanup_temp_checkpoint(temp_dir: Path) -> None:
    """Clean up temporary checkpoint directory."""
    import shutil
    try:
        if temp_dir.exists():
            shutil.rmtree(temp_dir)
            logger.debug("Cleaned up temp checkpoint: %s", temp_dir)
    except Exception as e:
        logger.warning("Failed to cleanup temp checkpoint: %s", e)


def run_inline_eval(
    model: Any,
    test_data: Any,
    config: dict[str, Any],
    epoch: int | None = None,
) -> dict[str, float]:
    """Run lightweight inline evaluation during training.

    This is a faster alternative to full post-training eval, suitable for
    periodic evaluation during training loops.

    Args:
        model: The GNN model (in eval mode)
        test_data: Test data tensor
        config: Training config
        epoch: Current epoch (for logging)

    Returns:
        Dict of evaluation metrics
    """
    import torch
    import numpy as np

    model.eval()
    metrics = {}

    with torch.no_grad():
        # Get logits
        # Note: This is a simplified version - full eval uses the complete pipeline
        pass

    prefix = f"[Epoch {epoch}] " if epoch is not None else ""
    logger.info("%sInline eval: %s", prefix, metrics)

    return metrics
