"""
Typed contracts for WastePredictor inputs and outputs.
Import these in any consumer: FL clients, API routes, tests, event handlers.
"""
from __future__ import annotations

from dataclasses import dataclass, field
from typing import Dict, List, Optional


# ── Input ──────────────────────────────────────────────────────────────────

@dataclass(frozen=True)
class PredictionInput:
    """
    Validated input for a single prediction.

    Example
    -------
    >>> inp = PredictionInput(
    ...     production_volume=50_000,
    ...     rain_sum=200.0,
    ...     temperature_mean=28.0,
    ...     humidity_mean=85.0,
    ...     wind_speed_mean=15.0,
    ...     month=6,
    ... )
    """
    production_volume: float
    rain_sum: float
    temperature_mean: float
    humidity_mean: float
    wind_speed_mean: float
    month: int = 6
    year: int = 2026

    def __post_init__(self) -> None:
        if not (1 <= self.month <= 12):
            raise ValueError(f"month must be 1-12, got {self.month}")
        if self.production_volume < 0:
            raise ValueError("production_volume must be >= 0")
        if not (0 <= self.humidity_mean <= 100):
            raise ValueError("humidity_mean must be 0-100")

    def to_dict(self) -> Dict[str, float]:
        return {
            "production_volume": self.production_volume,
            "rain_sum": self.rain_sum,
            "temperature_mean": self.temperature_mean,
            "humidity_mean": self.humidity_mean,
            "wind_speed_mean": self.wind_speed_mean,
            "month": float(self.month),
            "year": float(self.year),
        }


# ── Output ─────────────────────────────────────────────────────────────────

OUTPUT_COLUMNS: List[str] = [
    "Total_Waste_kg",
    "Solid_Waste_Limestone_kg",
    "Solid_Waste_Gypsum_kg",
    "Solid_Waste_Industrial_Salt_kg",
    "Liquid_Waste_Bittern_Liters",
    "Potential_Epsom_Salt_kg",
    "Potential_Potash_kg",
    "Potential_Magnesium_Oil_Liters",
]

INPUT_COLUMNS: List[str] = [
    "production_volume",
    "rain_sum",
    "temperature_mean",
    "humidity_mean",
    "wind_speed_mean",
    "Month",  # capital-M, as expected by AdvancedFeatureEngineer
]


@dataclass
class PredictionResult:
    """
    Strongly-typed wrapper for model output.

    Access fields directly or call `.to_dict()` for JSON serialisation.
    """
    total_waste_kg: float
    solid_waste_limestone_kg: float
    solid_waste_gypsum_kg: float
    solid_waste_industrial_salt_kg: float
    liquid_waste_bittern_liters: float
    potential_epsom_salt_kg: float
    potential_potash_kg: float
    potential_magnesium_oil_liters: float
    model_version: str = "unknown"
    confidence: Optional[float] = None

    @classmethod
    def from_dict(cls, d: Dict[str, float], model_version: str = "unknown") -> "PredictionResult":
        return cls(
            total_waste_kg=d["Total_Waste_kg"],
            solid_waste_limestone_kg=d["Solid_Waste_Limestone_kg"],
            solid_waste_gypsum_kg=d["Solid_Waste_Gypsum_kg"],
            solid_waste_industrial_salt_kg=d["Solid_Waste_Industrial_Salt_kg"],
            liquid_waste_bittern_liters=d["Liquid_Waste_Bittern_Liters"],
            potential_epsom_salt_kg=d["Potential_Epsom_Salt_kg"],
            potential_potash_kg=d["Potential_Potash_kg"],
            potential_magnesium_oil_liters=d["Potential_Magnesium_Oil_Liters"],
            model_version=model_version,
        )

    def to_dict(self) -> Dict[str, float]:
        return {
            "Total_Waste_kg": self.total_waste_kg,
            "Solid_Waste_Limestone_kg": self.solid_waste_limestone_kg,
            "Solid_Waste_Gypsum_kg": self.solid_waste_gypsum_kg,
            "Solid_Waste_Industrial_Salt_kg": self.solid_waste_industrial_salt_kg,
            "Liquid_Waste_Bittern_Liters": self.liquid_waste_bittern_liters,
            "Potential_Epsom_Salt_kg": self.potential_epsom_salt_kg,
            "Potential_Potash_kg": self.potential_potash_kg,
            "Potential_Magnesium_Oil_Liters": self.potential_magnesium_oil_liters,
        }
