"""Temporal provenance and policy-scoped freshness helpers for weather data."""

from __future__ import annotations

from dataclasses import dataclass
from datetime import datetime
from typing import Iterable

from mountain_twin.analysis_contract import (
    FreshnessState,
    IntendedUseContext,
    TemporalProvenance,
    WeatherSourceType,
)


@dataclass(frozen=True)
class FreshnessPolicy:
    """Optional product policy; no thresholds are assumed by default."""

    policy_id: str
    fresh_max_age_seconds: float
    stale_after_seconds: float

    def __post_init__(self) -> None:
        if self.fresh_max_age_seconds < 0 or self.stale_after_seconds < self.fresh_max_age_seconds:
            raise ValueError("freshness thresholds must be ordered and non-negative")


def parse_time(value: str | None) -> datetime | None:
    if value is None:
        return None
    parsed = datetime.fromisoformat(value)
    if parsed.utcoffset() is None:
        raise ValueError("provenance timestamps must be timezone-aware")
    return parsed


def age_at_evaluation(
    fetched_at: str | None, evaluation_time: str | datetime | None
) -> float | None:
    fetched = parse_time(fetched_at)
    if fetched is None or evaluation_time is None:
        return None
    evaluated = parse_time(evaluation_time) if isinstance(evaluation_time, str) else evaluation_time
    if evaluated is None or evaluated.utcoffset() is None:
        raise ValueError("evaluation time must be timezone-aware")
    return (evaluated - fetched).total_seconds()


def classify_freshness(age_seconds: float | None, policy: FreshnessPolicy | None) -> FreshnessState:
    if age_seconds is None or age_seconds < 0 or policy is None:
        return FreshnessState.UNKNOWN
    if age_seconds <= policy.fresh_max_age_seconds:
        return FreshnessState.FRESH
    if age_seconds <= policy.stale_after_seconds:
        return FreshnessState.AGED
    return FreshnessState.STALE


def consistency_codes(
    *,
    source_type: WeatherSourceType,
    fetched_at: str | None,
    valid_at: str | None,
    model_run_at: str | None,
    evaluation_time: str | datetime | None = None,
    contributing_valid_times: Iterable[str] = (),
) -> tuple[str, ...]:
    codes: set[str] = set()
    fetched, valid, run = parse_time(fetched_at), parse_time(valid_at), parse_time(model_run_at)
    if fetched is None:
        codes.add("FETCH_TIME_UNKNOWN")
    if valid is None:
        codes.add("VALID_TIME_UNKNOWN")
    if run is None:
        codes.add("MODEL_RUN_TIME_NOT_EXPOSED")
    if run and valid and run > valid:
        codes.add("MODEL_RUN_AFTER_VALID_TIME")
    evaluated = parse_time(evaluation_time) if isinstance(evaluation_time, str) else evaluation_time
    if (
        evaluated
        and valid
        and source_type is WeatherSourceType.HISTORICAL_FORECAST
        and valid > evaluated
    ):
        codes.add("HISTORICAL_VALID_TIME_IN_FUTURE")
    times = tuple(contributing_valid_times)
    if len(times) > 1:
        codes.add("DERIVED_FROM_MULTIPLE_VALID_TIMES")
    if source_type is WeatherSourceType.OBSERVATION and run is not None:
        codes.add("OBSERVATION_RUN_TIME_NOT_APPLICABLE")
    return tuple(sorted(codes))


def make_temporal_provenance(
    *,
    source_type: WeatherSourceType,
    fetched_at: str | None,
    valid_at: str | None,
    contributing_valid_times: tuple[str, ...] = (),
    model_run_at: str | None = None,
    evaluation_time: str | datetime | None = None,
    freshness_policy: FreshnessPolicy | None = None,
    intended_use_context: IntendedUseContext = IntendedUseContext.PLANNING_CONTEXT,
) -> TemporalProvenance:
    age = age_at_evaluation(fetched_at, evaluation_time)
    return TemporalProvenance(
        fetched_at=fetched_at,
        valid_at=valid_at,
        contributing_valid_times=contributing_valid_times or ((valid_at,) if valid_at else ()),
        model_run_at=model_run_at,
        age_at_evaluation_seconds=age,
        freshness_state=classify_freshness(age, freshness_policy),
        freshness_policy_id=freshness_policy.policy_id if freshness_policy else None,
        intended_use_context=intended_use_context,
        consistency_reason_codes=consistency_codes(
            source_type=source_type,
            fetched_at=fetched_at,
            valid_at=valid_at,
            model_run_at=model_run_at,
            evaluation_time=evaluation_time,
            contributing_valid_times=contributing_valid_times,
        ),
    )


def propagate_temporal_provenance(
    provenances: Iterable[TemporalProvenance],
    *,
    evaluation_time: str | datetime | None = None,
    freshness_policy: FreshnessPolicy | None = None,
    intended_use_context: IntendedUseContext = IntendedUseContext.PLANNING_CONTEXT,
) -> TemporalProvenance:
    """Conservatively combine contributors; derived data cannot appear fresher."""
    items = tuple(provenances)
    if not items:
        return make_temporal_provenance(
            source_type=WeatherSourceType.FORECAST,
            fetched_at=None,
            valid_at=None,
            evaluation_time=evaluation_time,
            freshness_policy=freshness_policy,
            intended_use_context=intended_use_context,
        )
    fetched = [item.fetched_at for item in items if item.fetched_at]
    valid_times = tuple(sorted({time for item in items for time in item.contributing_valid_times}))
    ages = [
        item.age_at_evaluation_seconds
        for item in items
        if item.age_at_evaluation_seconds is not None
    ]
    age = max(ages) if ages else None
    state = classify_freshness(age, freshness_policy)
    codes = tuple(
        sorted(
            {code for item in items for code in item.consistency_reason_codes}
            | ({"DERIVED_FROM_MULTIPLE_VALID_TIMES"} if len(valid_times) > 1 else set())
        )
    )
    return TemporalProvenance(
        fetched_at=min(fetched) if fetched else None,
        valid_at=valid_times[0] if len(valid_times) == 1 else None,
        contributing_valid_times=valid_times,
        model_run_at=None,
        age_at_evaluation_seconds=age,
        freshness_state=state,
        freshness_policy_id=freshness_policy.policy_id if freshness_policy else None,
        intended_use_context=intended_use_context,
        consistency_reason_codes=codes,
    )