import asyncio
import json
import time
from concurrent.futures import ThreadPoolExecutor
from typing import Any, Dict, Optional, Tuple
from huggingface_hub import hf_hub_download
import joblib
import pandas as pd


class ModelService:
    def __init__(self, model_repo: str, model_file: str, feature_file: str, metadata_file: str):
        self.model_repo = model_repo
        self.model_file = model_file
        self.feature_file = feature_file
        self.metadata_file = metadata_file

        self.pipeline = None
        self.expected_features = None
        self.load_time = None
        self.model_metadata: Optional[Dict[str, Any]] = None

        self.executor = ThreadPoolExecutor(max_workers=4)

        self._load_model()
        self._load_model_metadata()
        
    def _download_from_hf(self, filename: str) -> str:
        path = hf_hub_download(repo_id=self.model_repo, filename=filename, cache_dir="models")
        return path
    
    def _load_model(self):
        try:
            start_time = time.time()
            model_path = self._download_from_hf(self.model_file)
            feature_path = self._download_from_hf(self.feature_file)
    
            self.pipeline = joblib.load(model_path)
            self.expected_features = joblib.load(feature_path)
    
            self.load_time = round(time.time() - start_time, 3)
            print("Model Loaded Successfully!")
        except Exception as e:
            self.pipeline = None
            self.load_time = None
            print(f"❌ Model failed to load: {e}")

    def _load_model_metadata(self):
        try:
            metadata_path = self._download_from_hf(self.metadata_file)
    
            with open(metadata_path, "r") as f:
                self.model_metadata = json.load(f)
                print("Metadata loaded!")
        except Exception as e:
            print("❌ Failed to load metadata:", e)
            self.model_metadata = None

    def get_model_info(self) -> Dict[str, Any]:
        # Safety check
        if self.model_metadata is None:
            self._load_model_metadata()

        return {
            "path": self.model_path,
            "load_time_sec": self.load_time,
            "type": self.model_metadata["model_info"]["model_type"],
            "pipeline": self.model_metadata["model_info"]["model_name"],
        }

    def is_model_loaded(self) -> bool:
        return self.pipeline is not None

    async def preprocess_data(self, data_dict: Dict[str, Any]) -> pd.DataFrame:
        loop = asyncio.get_event_loop()
        return await loop.run_in_executor(
            self.executor, self._preprocess_sync, data_dict
        )

    def _preprocess_sync(self, data_dict: Dict[str, Any]) -> pd.DataFrame:
        df = pd.DataFrame([data_dict])

        # Reindex ke expected_features
        if self.expected_features:
            missing_model_features = [
                c for c in self.expected_features if c not in df.columns
            ]
            if missing_model_features:
                raise ValueError(f"Missing model features: {missing_model_features}")
            df = df.reindex(columns=self.expected_features, fill_value=0)

        return df

    async def predict(
        self, df: pd.DataFrame
    ) -> Tuple[int, Optional[float], Optional[float]]:
        if not self.is_model_loaded():
            raise RuntimeError("Model not loaded")

        loop = asyncio.get_event_loop()
        return await loop.run_in_executor(self.executor, self._predict_sync, df)

    def _predict_sync(
        self, df: pd.DataFrame
    ) -> Tuple[int, Optional[float], Optional[float], Optional[str]]:
        # Prediksi label
        pred = self.pipeline.predict(df)[0]

        # Prediksi probabilitas
        if hasattr(self.pipeline, "predict_proba"):
            proba = self.pipeline.predict_proba(df)[0]
            proba_no, proba_yes = float(proba[0]), float(proba[1])
        else:
            proba_yes = None
            proba_no = None

        # Tentukan priority/confidence berdasarkan probabilitas yes
        if proba_yes is not None:
            if proba_yes >= 0.7:
                priority = "HIGH"
            elif proba_yes >= 0.5:
                priority = "MEDIUM"
            elif proba_yes >= 0.3:
                priority = "LOW"
            else:
                priority = "VERY LOW"
        else:
            priority = None

        return int(pred), proba_yes, proba_no, priority


    def get_complete_model_info(self) -> (Dict[str, Any] | None):
        if self.model_metadata is None:
            self._load_model_metadata()
        return self.model_metadata