from __future__ import annotations

import logging
from datetime import datetime, timezone
from collections.abc import Sequence

from fastapi import APIRouter, FastAPI, HTTPException, Request, status
from pydantic import BaseModel, Field

from app.ml.evaluation import Evaluator
from app.ml.feature_store import FeatureVector
from app.ml.predictor import Predictor
from app.ml.registry.model_registry import ModelRegistry
from app.ml.retraining import retrain_model

logger = logging.getLogger(__name__)

router = APIRouter(prefix="/ml", tags=["ml"])


class HealthResponse(BaseModel):
    status: str
    updated_at: datetime = Field(default_factory=lambda: datetime.now(timezone.utc))


class PredictRequest(BaseModel):
    model_id: str = Field(..., alias="modelId")
    sensor_id: str
    values: dict[str, float]


class PredictResponse(BaseModel):
    model_id: str
    sensor_id: str
    predictions: dict[str, float]
    confidence: float
    model_type: str


class BatchRequest(BaseModel):
    requests: Sequence[PredictRequest]


class BatchResponse(BaseModel):
    predictions: Sequence[PredictResponse]


class ModelsResponse(BaseModel):
    models: list[str]


class TrainingSample(BaseModel):
    sensor_id: str = Field(min_length=1)
    values: dict[str, float]
    label: str | None = None


class RetrainRequest(BaseModel):
    model_id: str = Field(..., alias="modelId", min_length=1, max_length=128)
    samples: list[TrainingSample] = Field(min_length=1)


class RetrainResponse(BaseModel):
    model_id: str
    supported_sensors: list[str]
    trained_features: int
    model_type: str
    replaced: bool


class EvaluateRequest(BaseModel):
    model_id: str = Field(..., alias="modelId", min_length=1, max_length=128)
    samples: list[TrainingSample] = Field(min_length=1)


class MetricResponse(BaseModel):
    name: str
    value: float
    threshold: float | None = None


class EvaluateResponse(BaseModel):
    model_id: str
    sample_size: int
    metrics: list[MetricResponse]


@router.get("/health", response_model=HealthResponse, status_code=200)
def health() -> HealthResponse:
    return HealthResponse(status="ok")


@router.get("/models", response_model=ModelsResponse, status_code=200)
def list_models(request: Request) -> ModelsResponse:
    registry = _require_registry(request)
    models = [artifact.artifact_id for artifact in registry.list_models()]
    return ModelsResponse(models=models)


@router.post("/retrain", response_model=RetrainResponse, status_code=200)
def retrain(payload: RetrainRequest, request: Request) -> RetrainResponse:
    registry = _require_registry(request)
    vectors = [
        FeatureVector(
            sensor_id=sample.sensor_id,
            values=sample.values,
            label=sample.label,
        )
        for sample in payload.samples
    ]
    try:
        result = retrain_model(registry, payload.model_id, vectors)
    except ValueError as exc:
        raise HTTPException(
            status_code=status.HTTP_422_UNPROCESSABLE_CONTENT,
            detail=str(exc),
        ) from exc
    return RetrainResponse(
        model_id=result.artifact.artifact_id,
        supported_sensors=list(result.artifact.supported_sensors),
        trained_features=sum(
            len(feature_models)
            for feature_models in result.artifact.feature_models.values()
        ),
        model_type=result.artifact.model_type,
        replaced=result.replaced,
    )


@router.post("/evaluate", response_model=EvaluateResponse, status_code=200)
def evaluate(payload: EvaluateRequest, request: Request) -> EvaluateResponse:
    registry = _require_registry(request)
    vectors = [
        FeatureVector(
            sensor_id=sample.sensor_id,
            values=sample.values,
            label=sample.label,
        )
        for sample in payload.samples
    ]
    try:
        report = Evaluator(registry=registry).evaluate(payload.model_id, vectors)
    except ValueError as exc:
        try:
            registry.load_artifact(payload.model_id)
        except KeyError:
            raise HTTPException(
                status_code=status.HTTP_404_NOT_FOUND,
                detail=str(exc),
            ) from exc
        raise HTTPException(
            status_code=status.HTTP_422_UNPROCESSABLE_CONTENT,
            detail=str(exc),
        ) from exc
    return EvaluateResponse(
        model_id=report.artifact_id,
        sample_size=report.sample_size,
        metrics=[
            MetricResponse(name=metric.name, value=metric.value, threshold=metric.threshold)
            for metric in report.metrics
        ],
    )


@router.post("/predict", response_model=PredictResponse, status_code=200)
def predict(payload: PredictRequest, request: Request) -> PredictResponse:
    registry = _require_registry(request)
    predictor = Predictor(registry=registry)
    vector = FeatureVector(sensor_id=payload.sensor_id, values=payload.values)
    try:
        prediction = predictor.predict(payload.model_id, vector)
    except KeyError as exc:
        raise HTTPException(status_code=status.HTTP_404_NOT_FOUND, detail=str(exc)) from exc
    except ValueError as exc:
        raise HTTPException(
            status_code=status.HTTP_422_UNPROCESSABLE_CONTENT,
            detail=str(exc),
        ) from exc
    return PredictResponse(
        model_id=payload.model_id,
        sensor_id=payload.sensor_id,
        predictions=prediction.predictions,
        confidence=prediction.confidence,
        model_type=prediction.model_type,
    )


@router.post("/batch", response_model=BatchResponse, status_code=200)
def predict_batch(payload: BatchRequest, request: Request) -> BatchResponse:
    registry = _require_registry(request)
    predictor = Predictor(registry=registry)
    responses: list[PredictResponse] = []
    for item in payload.requests:
        vector = FeatureVector(sensor_id=item.sensor_id, values=item.values)
        try:
            prediction = predictor.predict(item.model_id, vector)
        except KeyError as exc:
            raise HTTPException(status_code=status.HTTP_404_NOT_FOUND, detail=str(exc)) from exc
        except ValueError as exc:
            raise HTTPException(
                status_code=status.HTTP_422_UNPROCESSABLE_CONTENT,
                detail=str(exc),
            ) from exc
        responses.append(
            PredictResponse(
                model_id=item.model_id,
                sensor_id=item.sensor_id,
                predictions=prediction.predictions,
                confidence=prediction.confidence,
                model_type=prediction.model_type,
            )
        )
    return BatchResponse(predictions=responses)


def _require_registry(request: Request) -> ModelRegistry:
    registry = getattr(request.app.state, "registry", None)
    if not isinstance(registry, ModelRegistry):
        raise HTTPException(
            status_code=status.HTTP_503_SERVICE_UNAVAILABLE,
            detail="ML registry nicht initialisiert.",
        )
    return registry


def init_ml_routes(app: FastAPI, model_store: str = ".model_store") -> None:
    app.state.model_store = model_store
    app.include_router(router)
    logger.info("ML routes registered")