"""Experiment tracker -- logs algorithm configs, data configs, and results."""

from __future__ import annotations

import json
import time
from dataclasses import dataclass, field
from pathlib import Path
from typing import Protocol


class _ExperimentStore(Protocol):
    def save_experiment(self, experiment: Experiment) -> None: ...


@dataclass
class Experiment:
    """A tracked experiment with config, results, and lifecycle state."""

    experiment_id: str
    name: str
    algorithm: str
    algorithm_config: dict = field(default_factory=dict)
    data_config: dict = field(default_factory=dict)
    results: dict = field(default_factory=dict)
    start_time: float = field(default_factory=time.time)
    end_time: float = 0.0
    status: str = "pending"  # pending, running, completed, failed
    tags: list[str] = field(default_factory=list)

    @property
    def duration_ms(self) -> float:
        """Duration in milliseconds, or 0 if not yet ended."""
        if self.end_time <= 0:
            return 0.0
        return (self.end_time - self.start_time) * 1000

    def to_dict(self) -> dict:
        """Serialize to dictionary."""
        return {
            "experiment_id": self.experiment_id,
            "name": self.name,
            "algorithm": self.algorithm,
            "algorithm_config": self.algorithm_config,
            "data_config": self.data_config,
            "results": self.results,
            "start_time": self.start_time,
            "end_time": self.end_time,
            "duration_ms": round(self.duration_ms, 2),
            "status": self.status,
            "tags": self.tags,
        }


class ExperimentTracker:
    """Tracks experiments with their configs and results."""

    def __init__(
        self,
        storage_path: str | Path | None = None,
        db: _ExperimentStore | None = None,
    ) -> None:
        self._experiments: dict[str, Experiment] = {}
        self._storage_path = Path(storage_path) if storage_path else None
        self._db = db

    def create(
        self,
        experiment_id: str,
        name: str,
        algorithm: str,
        algorithm_config: dict | None = None,
        data_config: dict | None = None,
        tags: list[str] | None = None,
    ) -> Experiment:
        """Create a new experiment in pending state."""
        exp = Experiment(
            experiment_id=experiment_id,
            name=name,
            algorithm=algorithm,
            algorithm_config=algorithm_config or {},
            data_config=data_config or {},
            tags=tags or [],
        )
        self._experiments[experiment_id] = exp
        return exp

    def start(self, experiment_id: str) -> Experiment:
        """Transition experiment to running state."""
        exp = self._experiments[experiment_id]
        exp.status = "running"
        exp.start_time = time.time()
        return exp

    def complete(self, experiment_id: str, results: dict) -> Experiment:
        """Mark experiment as completed with results."""
        exp = self._experiments[experiment_id]
        exp.status = "completed"
        exp.end_time = time.time()
        exp.results = results
        if self._storage_path:
            self._persist()
        if self._db is not None:
            self._db.save_experiment(exp)
        return exp

    def fail(self, experiment_id: str, error: str) -> Experiment:
        """Mark experiment as failed with error message."""
        exp = self._experiments[experiment_id]
        exp.status = "failed"
        exp.end_time = time.time()
        exp.results = {"error": error}
        return exp

    def get(self, experiment_id: str) -> Experiment | None:
        """Retrieve an experiment by ID."""
        return self._experiments.get(experiment_id)

    def list_experiments(
        self,
        status: str | None = None,
        algorithm: str | None = None,
        tag: str | None = None,
    ) -> list[Experiment]:
        """List experiments with optional filters."""
        exps = list(self._experiments.values())
        if status:
            exps = [e for e in exps if e.status == status]
        if algorithm:
            exps = [e for e in exps if e.algorithm == algorithm]
        if tag:
            exps = [e for e in exps if tag in e.tags]
        return sorted(exps, key=lambda e: e.start_time)

    def summary(self) -> dict:
        """Aggregate summary of all tracked experiments."""
        exps = list(self._experiments.values())
        if not exps:
            return {"total": 0, "by_status": {}, "by_algorithm": {}}
        by_status: dict[str, int] = {}
        by_algorithm: dict[str, int] = {}
        for e in exps:
            by_status[e.status] = by_status.get(e.status, 0) + 1
            by_algorithm[e.algorithm] = by_algorithm.get(e.algorithm, 0) + 1
        return {
            "total": len(exps),
            "by_status": by_status,
            "by_algorithm": by_algorithm,
        }

    def _persist(self) -> None:
        """Write all experiments to storage."""
        if self._storage_path:
            self._storage_path.parent.mkdir(parents=True, exist_ok=True)
            data = {eid: e.to_dict() for eid, e in self._experiments.items()}
            self._storage_path.write_text(json.dumps(data, indent=2))

    def load(self, path: str | Path) -> None:
        """Load experiments from a JSON file."""
        p = Path(path)
        if p.exists():
            data = json.loads(p.read_text())
            for eid, d in data.items():
                tags = d.pop("tags", [])
                d.pop("duration_ms", None)
                self._experiments[eid] = Experiment(**d, tags=tags)
