"""
Metrics Storage for evaluation results.

Handles:
- Saving metrics.json per iteration
- Appending to iterations_history.parquet
- Loading baseline for comparison
"""

import json
from dataclasses import asdict
from datetime import datetime
from pathlib import Path
from typing import Dict, Any, Optional, List

import pandas as pd
from loguru import logger

from .scorer import EvalScore, PatternTypeMetrics


class MetricsStorage:
    """
    Stores and retrieves evaluation metrics.

    Output structure:
    data/output/torus/2025-12-10/256/
    ├── baseline/
    │   ├── patterns_*.parquet
    │   ├── features.parquet
    │   └── metrics.json
    ├── v_0001/
    │   ├── patterns_*.parquet
    │   ├── features.parquet
    │   └── metrics.json
    └── iterations_history.parquet
    """

    def __init__(self, output_base_path: Path):
        """
        Initialize metrics storage.

        Args:
            output_base_path: Base path for outputs (e.g., data/output/torus/2025-12-10/256)
        """
        self.output_base_path = Path(output_base_path)
        self.output_base_path.mkdir(parents=True, exist_ok=True)

    def save_metrics(
        self,
        iteration: str,
        score: EvalScore,
        build_command: Optional[str] = None,
        build_duration: Optional[float] = None,
    ) -> Path:
        """
        Save metrics.json for an iteration.

        Args:
            iteration: Iteration name ("baseline" or "v_XXXX")
            score: EvalScore with all metrics
            build_command: Optional build command that was run
            build_duration: Optional build duration in seconds

        Returns:
            Path to saved metrics.json
        """
        iteration_path = self.output_base_path / iteration
        iteration_path.mkdir(parents=True, exist_ok=True)

        metrics_file = iteration_path / "metrics.json"

        # Build metrics dict matching the schema from the plan
        metrics = {
            "iteration": iteration,
            "timestamp": datetime.utcnow().isoformat() + "Z",

            "summary": {
                "final_score": round(score.final_score, 4),
                "recall_score": round(score.recall_score, 4),
                "novelty_score": round(score.novelty_score, 4),
                "performance_score": round(score.performance_score, 4),
            },

            "counts": {
                "total_patterns": score.total_patterns,
                "synthetic_found": score.synthetic_found,
                "synthetic_expected": score.synthetic_expected,
                "novelty_valid": score.novelty_valid,
                "novelty_invalid": score.novelty_invalid,
            },

            "timing": {
                "total_seconds": round(score.total_seconds, 2),
                "detectors": {k: round(v, 2) for k, v in score.detector_timing.items()},
            },

            "per_pattern_type": {
                ptype: {
                    "detected": metrics.detected,
                    "recall": round(metrics.recall, 4),
                    "novelty_valid": metrics.novelty_valid,
                }
                for ptype, metrics in score.per_pattern_type.items()
            },
        }

        # Add build info if available
        if build_command:
            metrics["build"] = {
                "command": build_command,
                "duration_seconds": round(build_duration, 2) if build_duration else 0,
            }

        # Add comparison to baseline if available
        if score.baseline_comparison:
            metrics["comparison_to_baseline"] = {
                k: round(v, 4) for k, v in score.baseline_comparison.items()
            }

        # Write JSON
        with open(metrics_file, 'w') as f:
            json.dump(metrics, f, indent=2)

        logger.info(f"Saved metrics to {metrics_file}")
        return metrics_file

    def load_metrics(self, iteration: str) -> Optional[Dict[str, Any]]:
        """
        Load metrics.json for an iteration.

        Args:
            iteration: Iteration name

        Returns:
            Metrics dict or None if not found
        """
        metrics_file = self.output_base_path / iteration / "metrics.json"

        if not metrics_file.exists():
            return None

        with open(metrics_file, 'r') as f:
            return json.load(f)

    def load_baseline_score(self) -> Optional[EvalScore]:
        """
        Load baseline EvalScore for comparison.

        Returns:
            EvalScore or None if baseline doesn't exist
        """
        metrics = self.load_metrics("baseline")
        if metrics is None:
            return None

        # Reconstruct EvalScore from metrics
        return EvalScore(
            final_score=metrics["summary"]["final_score"],
            recall_score=metrics["summary"]["recall_score"],
            novelty_score=metrics["summary"]["novelty_score"],
            performance_score=metrics["summary"]["performance_score"],
            total_patterns=metrics["counts"]["total_patterns"],
            synthetic_found=metrics["counts"]["synthetic_found"],
            synthetic_expected=metrics["counts"]["synthetic_expected"],
            novelty_valid=metrics["counts"]["novelty_valid"],
            novelty_invalid=metrics["counts"]["novelty_invalid"],
            total_seconds=metrics["timing"]["total_seconds"],
            detector_timing=metrics["timing"].get("detectors", {}),
            per_pattern_type={
                ptype: PatternTypeMetrics(
                    detected=data["detected"],
                    recall=data["recall"],
                    novelty_valid=data["novelty_valid"],
                )
                for ptype, data in metrics.get("per_pattern_type", {}).items()
            },
        )

    def append_to_history(self, iteration: str, score: EvalScore) -> Path:
        """
        Append iteration results to history parquet.

        Args:
            iteration: Iteration name
            score: EvalScore with metrics

        Returns:
            Path to history parquet
        """
        history_file = self.output_base_path / "iterations_history.parquet"

        # Build row for history
        row = {
            "iteration": iteration,
            "timestamp": datetime.utcnow().isoformat(),
            "final_score": score.final_score,
            "recall_score": score.recall_score,
            "novelty_score": score.novelty_score,
            "performance_score": score.performance_score,
            "total_patterns": score.total_patterns,
            "synthetic_found": score.synthetic_found,
            "synthetic_expected": score.synthetic_expected,
            "novelty_valid": score.novelty_valid,
            "novelty_invalid": score.novelty_invalid,
            "total_seconds": score.total_seconds,
        }

        # Add baseline comparison if available
        if score.baseline_comparison:
            row["final_score_delta"] = score.baseline_comparison.get("final_score_delta", 0)
            row["recall_delta"] = score.baseline_comparison.get("recall_delta", 0)

        new_row = pd.DataFrame([row])

        # Load existing history or create new
        if history_file.exists():
            existing = pd.read_parquet(history_file, engine='pyarrow')
            # Remove existing row for this iteration (if re-running)
            existing = existing[existing['iteration'] != iteration]
            history = pd.concat([existing, new_row], ignore_index=True)
        else:
            history = new_row

        # Sort by timestamp
        history = history.sort_values('timestamp').reset_index(drop=True)

        # Save
        history.to_parquet(history_file, engine='pyarrow', index=False)

        logger.info(f"Appended to history: {history_file} ({len(history)} iterations)")
        return history_file

    def get_history(self) -> pd.DataFrame:
        """
        Get full iteration history.

        Returns:
            DataFrame with all iterations
        """
        history_file = self.output_base_path / "iterations_history.parquet"

        if not history_file.exists():
            return pd.DataFrame()

        return pd.read_parquet(history_file, engine='pyarrow')

    def list_iterations(self) -> List[str]:
        """
        List all iterations that have metrics.

        Returns:
            List of iteration names
        """
        iterations = []
        for path in self.output_base_path.iterdir():
            if path.is_dir() and (path / "metrics.json").exists():
                iterations.append(path.name)
        return sorted(iterations)

    def print_history_table(self) -> str:
        """
        Format history as a readable table.

        Returns:
            Formatted string table
        """
        history = self.get_history()

        if history.empty:
            return "No iterations found."

        # Select columns for display
        display_cols = [
            'iteration', 'final_score', 'recall_score', 'novelty_score',
            'performance_score', 'total_seconds'
        ]
        display_cols = [c for c in display_cols if c in history.columns]

        # Format
        lines = []
        lines.append("-" * 80)
        lines.append(f"{'Iteration':<12} {'Final':>8} {'Recall':>8} {'Novelty':>8} {'Perf':>8} {'Time':>8}")
        lines.append("-" * 80)

        for _, row in history.iterrows():
            lines.append(
                f"{row['iteration']:<12} "
                f"{row['final_score']:>8.4f} "
                f"{row['recall_score']:>8.4f} "
                f"{row['novelty_score']:>8.4f} "
                f"{row['performance_score']:>8.4f} "
                f"{row['total_seconds']:>7.1f}s"
            )

        lines.append("-" * 80)
        return "\n".join(lines)
