import json
import logging
import pickle
from dataclasses import dataclass
from pathlib import Path
from typing import Any, Dict, List, Optional

import joblib
from tensorflow.keras.models import load_model

logger = logging.getLogger(__name__)

COMMODITIES = {
    "rawit": "Cabai Rawit",
    "merah": "Cabai Merah Besar",
    "keriting": "Cabai Merah Keriting",
}
MARKETS = {
    "pasar_tanjung": "Pasar Tanjung",
    "pasar_kreongan": "Pasar Kreongan",
    "pasar_tegal_besar": "Pasar Tegal Besar",
    "pasar_mangli": "Pasar Mangli",
    "pasar_wirolegi": "Pasar Wirolegi",
}
DEFAULT_TIME_STEPS = {"rawit": 7, "merah": 7, "keriting": 14}


@dataclass
class ModelArtifact:
    commodity: str
    market: str
    model: Any
    scaler: Any
    config: Dict[str, Any]
    time_step: int
    model_time_step: int


class ModelRegistry:
    def __init__(self, models_dir: Path) -> None:
        self.models_dir = models_dir
        self.artifacts: Dict[str, Dict[str, ModelArtifact]] = {}
        self.skipped: List[Dict[str, str]] = []
        self.scan()

    def scan(self) -> None:
        self.artifacts.clear()
        self.skipped.clear()
        for commodity, commodity_name in COMMODITIES.items():
            commodity_dir = self.models_dir / commodity
            for market in MARKETS:
                market_dir = commodity_dir / market
                config_path = market_dir / "config.json"
                if not config_path.is_file():
                    self._skip(commodity, market, "config.json tidak ditemukan")
                    continue
                try:
                    config = json.loads(config_path.read_text(encoding="utf-8"))
                    files = config.get("files", {})
                    model_path = self._trusted_path(market_dir, files.get("model"))
                    scaler_path = self._trusted_path(market_dir, files.get("scaler"))
                    if not model_path.is_file() or not scaler_path.is_file():
                        raise FileNotFoundError("model atau scaler tidak lengkap")
                    model = load_model(model_path, compile=False)
                    try:
                        scaler = joblib.load(scaler_path)
                    except Exception:
                        with scaler_path.open("rb") as scaler_file:
                            scaler = pickle.load(scaler_file)
                    if not hasattr(scaler, "data_min_") or not hasattr(scaler, "data_max_"):
                        raise ValueError("scaler tidak memiliki data_min_ dan data_max_")
                    model_time_step = int(model.input_shape[1])
                    configured_step = int(config.get("model", {}).get("time_step", DEFAULT_TIME_STEPS[commodity]))
                    time_step = DEFAULT_TIME_STEPS[commodity]
                    if commodity == "keriting" and configured_step != time_step:
                        logger.warning("Menggunakan time_step API 14 untuk keriting; config model menyatakan %s", configured_step)
                    self.artifacts.setdefault(commodity, {})[market] = ModelArtifact(
                        commodity, market, model, scaler, config, time_step, model_time_step
                    )
                except Exception as error:
                    self._skip(commodity, market, f"{type(error).__name__}: {error}")

    @staticmethod
    def _trusted_path(directory: Path, filename: Optional[str]) -> Path:
        if not filename:
            raise FileNotFoundError("nama file artifact tidak ada di config.json")
        path = (directory / filename).resolve()
        if path.parent != directory.resolve():
            raise ValueError("path artifact keluar dari folder market")
        return path

    def _skip(self, commodity: str, market: str, reason: str) -> None:
        logger.warning("Melewati model %s/%s: %s", commodity, market, reason)
        self.skipped.append({"commodity": commodity, "market": market, "reason": reason})

    def get(self, commodity: str, market: str) -> Optional[ModelArtifact]:
        return self.artifacts.get(commodity, {}).get(market)

    def available_markets(self, commodity: str) -> List[str]:
        return sorted(self.artifacts.get(commodity, {}))