from __future__ import annotations

import logging
import math
from collections.abc import Sequence
from dataclasses import dataclass

from app.ml.feature_store import FeatureVector
from app.ml.predictor import Predictor
from app.ml.registry.model_registry import ModelRegistry
from app.ml.training import TrainingPipeline

logger = logging.getLogger(__name__)


@dataclass
class Metric:
    name: str
    value: float
    threshold: float | None = None


@dataclass
class EvalReport:
    artifact_id: str
    sample_size: int
    metrics: list[Metric]


class Evaluator:
    def __init__(
        self,
        pipeline: TrainingPipeline | None = None,
        registry: ModelRegistry | None = None,
    ) -> None:
        if isinstance(pipeline, ModelRegistry) and registry is None:
            registry = pipeline
            pipeline = None
        if pipeline is None and registry is None:
            raise ValueError("Evaluator erfordert TrainingPipeline oder ModelRegistry.")
        self._pipeline = pipeline
        self._registry = registry
        self._predictor = Predictor(pipeline=pipeline, registry=registry)

    def evaluate(self, artifact_id: str, samples: Sequence[FeatureVector]) -> EvalReport:
        try:
            if self._registry is not None:
                self._registry.load_artifact(artifact_id)
            elif self._pipeline is not None:
                self._pipeline.export(artifact_id)
        except KeyError as exc:
            raise ValueError("Kein trainiertes Modell für Evaluation vorhanden.") from exc

        absolute_errors: list[float] = []
        squared_errors: list[float] = []
        for sample in samples:
            try:
                prediction = self._predictor.predict(artifact_id, sample)
            except ValueError:
                continue
            for feature_name, predicted in prediction.predictions.items():
                actual = float(sample.values[feature_name])
                error = predicted - actual
                absolute_errors.append(abs(error))
                squared_errors.append(error**2)

        sample_size = len(absolute_errors)
        mae = sum(absolute_errors) / sample_size if sample_size else 0.0
        rmse = math.sqrt(sum(squared_errors) / sample_size) if sample_size else 0.0
        expected_values = sum(len(sample.values) for sample in samples)
        coverage = sample_size / expected_values if expected_values else 0.0

        report = EvalReport(
            artifact_id=artifact_id,
            sample_size=sample_size,
            metrics=[
                Metric(name="mae", value=mae),
                Metric(name="rmse", value=rmse),
                Metric(name="coverage", value=coverage, threshold=0.8),
            ],
        )
        logger.info(
            "Evaluation %s -> mae=%.4f, rmse=%.4f, coverage=%.2f",
            artifact_id,
            mae,
            rmse,
            coverage,
        )
        return report