"""
KWS High-Level API – Lớp trừu tượng cao, đa dụng, dễ dùng cho KWS_Lib.

3 class chính:
- KWSSession: Đăng nhập & giữ client API.
- DatasetPipeline: Chia tập, lấy dữ liệu, tải và chuẩn bị cho training.
- ModelManager: Quản lý model, experiment run, metrics, artifact lên DB.

Luồng điển hình: Session → Pipeline (chia & load data) → [User train] → ModelManager (đăng ký run, metrics, artifact).
"""

from __future__ import annotations

import logging
from pathlib import Path
from typing import Any, Dict, List, Optional, Tuple, Union

import numpy as np
import pandas as pd
import requests

from . import KWSClient, build_metrics_payload
from .minio_client import DatasetSplitFilesClient
from .training import DataLoader

logger = logging.getLogger(__name__)


# ---------------------------------------------------------------------------
# 1. KWSSession – Đăng nhập
# ---------------------------------------------------------------------------


class KWSSession:
    """
    Phiên làm việc KWS: kết nối API và đăng nhập.
    Dùng làm điểm vào duy nhất cho các class cao cấp (DatasetPipeline, ModelManager).
    """

    def __init__(
        self,
        base_url: str = "http://localhost:8000",
        username: Optional[str] = None,
        password: Optional[str] = None,
        *,
        token: Optional[str] = None,
        api_key: Optional[str] = None,
    ):
        """
        Args:
            base_url: URL backend KWS (vd: http://localhost:8000).
            username: Username đăng nhập (bỏ qua nếu dùng token/api_key).
            password: Password đăng nhập.
            token: JWT sẵn có (không cần login).
            api_key: API key thay JWT.
        """
        self._client = KWSClient(
            base_url=base_url,
            token=token,
            api_key=api_key,
        )
        if username and password:
            self._client.login(username=username, password=password)
            logger.info("Đăng nhập thành công.")
        self._files_client = DatasetSplitFilesClient(self._client)

    @property
    def api(self) -> KWSClient:
        """Client API đã đăng nhập."""
        return self._client

    @property
    def client(self) -> KWSClient:
        """Alias của api."""
        return self._client

    @property
    def files_client(self) -> DatasetSplitFilesClient:
        """Client tải file split (MinIO)."""
        return self._files_client

    def login(self, username: str, password: str) -> Dict[str, Any]:
        """Đăng nhập (hoặc đăng nhập lại)."""
        return self._client.login(username=username, password=password)


# ---------------------------------------------------------------------------
# 2. DatasetPipeline – Chia tập & lấy dữ liệu
# ---------------------------------------------------------------------------


class DatasetPipeline:
    """
    Pipeline chia tập và lấy dữ liệu: lấy raw data → lọc class → chia train/test → đẩy DB → tải mảng & chuẩn bị CNN.
    """

    def __init__(self, session: KWSSession):
        self._session = session
        self._api = session.api
        self._data_loader = DataLoader(self._api, session.files_client)

    def get_data(
        self,
        dataset_version_id: int,
        feature_type_id: Optional[int] = None,
        keyword_ids: Optional[List[int]] = None,
    ) -> pd.DataFrame:
        """
        Lấy toàn bộ dữ liệu feature (MFCC) từ API.
        Returns:
            DataFrame với cột: feature_id, sample_id, keyword_id, keyword_name, file_name, minio_path, feature_type_id.
        """
        df = self._api.dataset_splits.get_mfcc_files(
            dataset_version_id=dataset_version_id,
            feature_type_id=feature_type_id,
            keyword_ids=keyword_ids,
        )
        logger.info(f"Lấy dữ liệu: {len(df)} mẫu.")
        return df

    def filter_classes(
        self,
        df: pd.DataFrame,
        class_names: List[str],
        label_column: str = "derivative_label",
    ) -> pd.DataFrame:
        """
        Chỉ giữ các dòng có keyword_name nằm trong class_names.
        Tạo cột derivative_label = keyword_name (để dùng khi push split).
        """
        out = df[df["keyword_name"].isin(class_names)].copy()
        out[label_column] = out["keyword_name"]
        logger.info(f"Sau lọc {len(class_names)} class: {len(out)} mẫu.")
        return out

    def split_and_push(
        self,
        df: pd.DataFrame,
        config_name: str,
        dataset_version_id: int,
        train_size: float = 0.8,
        test_size: float = 0.2,
        val_size: float = 0.0,
        random_state: int = 42,
        stratify: bool = True,
        create_merged_keywords: Optional[Dict[str, List[int]]] = None,
        auto_detect_merged_keywords: bool = True,
    ) -> Dict[str, int]:
        """
        Chia df thành train/test (và val nếu val_size > 0), đẩy lên DB.
        Args:
            df: DataFrame có cột feature_id, keyword_id, derivative_label (và keyword_name để auto-detect merge).
            config_name: Tên cấu hình split (vd: config_10class_80_20).
            dataset_version_id: ID phiên bản dataset.
            train_size, test_size, val_size: Tỷ lệ (train_size + test_size + val_size = 1.0).
            random_state, stratify: Cho train_test_split.
        Returns:
            Dict split_name -> split_id, vd: {"train": 123, "test": 124} hoặc {"train": 123, "val": 124, "test": 125}.
        """
        if abs(train_size + test_size + val_size - 1.0) > 1e-6:
            raise ValueError("train_size + test_size + val_size phải bằng 1.0")

        stratify_col = df["derivative_label"] if stratify else None

        if val_size <= 0:
            from sklearn.model_selection import train_test_split
            train_df, test_df = train_test_split(
                df, train_size=train_size, test_size=test_size,
                random_state=random_state, stratify=stratify_col,
            )
            splits = {"train": train_df, "test": test_df}
        else:
            from sklearn.model_selection import train_test_split
            # train vs (val+test)
            train_df, rest_df = train_test_split(
                df, train_size=train_size, test_size=(val_size + test_size),
                random_state=random_state, stratify=stratify_col,
            )
            rest_stratify = rest_df["derivative_label"] if stratify else None
            val_ratio = val_size / (val_size + test_size) if (val_size + test_size) > 0 else 0
            val_df, test_df = train_test_split(
                rest_df, train_size=val_ratio, test_size=(1 - val_ratio),
                random_state=random_state, stratify=rest_stratify,
            )
            splits = {"train": train_df, "val": val_df, "test": test_df}

        created = self._push_splits(
            dataset_version_id=dataset_version_id,
            config_name=config_name,
            splits=splits,
            create_merged_keywords=create_merged_keywords,
            auto_detect_merged_keywords=auto_detect_merged_keywords,
        )
        return created

    def _push_splits(
        self,
        dataset_version_id: int,
        config_name: str,
        splits: Dict[str, pd.DataFrame],
        create_merged_keywords: Optional[Dict[str, List[int]]] = None,
        auto_detect_merged_keywords: bool = True,
    ) -> Dict[str, int]:
        """Đẩy các split lên DB (gọi create_split_from_list từng split)."""
        required = ["feature_id", "keyword_id", "derivative_label"]
        for name, frame in splits.items():
            if frame.empty:
                continue
            missing = [c for c in required if c not in frame.columns]
            if missing:
                raise ValueError(f"Split '{name}' thiếu cột: {missing}")

        merged = create_merged_keywords
        if auto_detect_merged_keywords and merged is None:
            first_name = next(iter(splits))
            first_df = splits[first_name]
            if not first_df.empty and "keyword_name" in first_df.columns:
                groups = first_df.groupby("keyword_name")["keyword_id"].unique()
                merged = {
                    k: sorted(v.tolist())
                    for k, v in groups.items()
                    if len(v) > 1
                }
                if merged:
                    logger.info(f"Auto-detected merged keywords: {list(merged)}")

        created_splits: Dict[str, int] = {}
        for split_name, split_df in splits.items():
            if split_df.empty:
                continue
            files = []
            for _, row in split_df.iterrows():
                d = {
                    "feature_id": int(row["feature_id"]),
                    "keyword_id": int(row["keyword_id"]),
                }
                dl = row.get("derivative_label")
                if pd.notna(dl) and str(dl).strip():
                    d["derivative_label"] = str(dl).strip()
                else:
                    d["derivative_label"] = None
                files.append(d)

            use_merged = merged if split_name == next(iter(splits)) else None
            resp = self._api.dataset_splits.create_split_from_list(
                dataset_version_id=dataset_version_id,
                split_name=split_name,
                config_name=config_name,
                files=files,
                create_merged_keywords=use_merged,
            )
            sid = resp.get("dataset_splits_id") or resp.get("id")
            created_splits[split_name] = sid
            logger.info(f"Đã tạo split '{split_name}' (ID: {sid}).")
        return created_splits

    def load_arrays(
        self,
        train_split_id: int,
        test_split_id: int,
        base_dir: str = "data",
    ) -> Tuple[np.ndarray, np.ndarray, np.ndarray, np.ndarray, Dict[str, int]]:
        """
        Tải dữ liệu train/test từ split đã tạo, trả về mảng và label_to_id.
        Returns:
            (X_train, y_train, X_test, y_test, label_to_id).
        """
        return self._data_loader.download_and_load_data(
            train_split_id=train_split_id,
            test_split_id=test_split_id,
            base_dir=base_dir,
        )

    def prepare_for_cnn(
        self,
        X_train: np.ndarray,
        X_test: np.ndarray,
    ) -> Tuple[np.ndarray, np.ndarray]:
        """Chuẩn hóa và reshape cho CNN."""
        return self._data_loader.prepare_for_cnn(X_train, X_test)

    def run(
        self,
        dataset_version_id: int,
        config_name: str,
        class_names: List[str],
        feature_type_id: Optional[int] = None,
        train_size: float = 0.8,
        test_size: float = 0.2,
        val_size: float = 0.0,
        random_state: int = 42,
        base_dir: str = "data",
        max_samples_per_class: Optional[int] = None,
        max_total_samples: Optional[int] = None,
        split_ids: Optional[Dict[str, int]] = None,
    ) -> Tuple[
        np.ndarray, np.ndarray, np.ndarray, np.ndarray,
        Dict[str, int],
        Dict[str, int],
    ]:
        """
        One-shot: (nếu chưa có split_ids) get_data → filter_classes → split_and_push;
        sau đó lấy danh sách file từ DB, so với thư mục local, thiếu thì tải thêm → load_arrays → prepare_for_cnn.

        - Lần đầu: split_ids=None → tạo split mới, tải NPZ về base_dir.
        - Lần sau: truyền split_ids={"train": <id>, "test": <id>} → chỉ lấy dữ liệu từ DB,
          so với file đã có trong base_dir, thiếu thì tải thêm, rồi load.

        max_samples_per_class / max_total_samples: chỉ áp dụng khi split_ids=None.
        Returns:
            (X_train, y_train, X_test, y_test, label_to_id, split_ids).
        """
        if split_ids is not None:
            if "train" not in split_ids or "test" not in split_ids:
                raise ValueError("split_ids phải có key 'train' và 'test'")
            train_id = split_ids["train"]
            test_id = split_ids["test"]
            logger.info(f"Dùng split có sẵn: train={train_id}, test={test_id}. Chỉ tải file thiếu.")
        else:
            df = self.get_data(dataset_version_id, feature_type_id=feature_type_id)
            df = self.filter_classes(df, class_names)
            if max_samples_per_class is not None and max_samples_per_class > 0 and "derivative_label" in df.columns:
                df = df.groupby("derivative_label", group_keys=False).apply(
                    lambda g: g.sample(n=min(len(g), max_samples_per_class), random_state=random_state)
                ).reset_index(drop=True)
                logger.info(f"Giới hạn tối đa {max_samples_per_class} mẫu/class → {len(df)} mẫu.")
            if max_total_samples is not None and max_total_samples > 0 and len(df) > max_total_samples:
                df = df.sample(n=max_total_samples, random_state=random_state).reset_index(drop=True)
                logger.info(f"Giới hạn tổng {max_total_samples} mẫu → {len(df)} mẫu.")
            split_ids = self.split_and_push(
                df, config_name, dataset_version_id,
                train_size=train_size, test_size=test_size, val_size=val_size,
                random_state=random_state,
            )
            train_id = split_ids["train"]
            test_id = split_ids["test"]
        X_train, y_train, X_test, y_test, label_to_id = self.load_arrays(
            train_id, test_id, base_dir=base_dir,
        )
        X_train, X_test = self.prepare_for_cnn(X_train, X_test)
        return X_train, y_train, X_test, y_test, label_to_id, split_ids


# ---------------------------------------------------------------------------
# 3. ModelManager – Quản lý model, metrics, artifact
# ---------------------------------------------------------------------------


class ModelManager:
    """
    Quản lý model trên DB: get/create model init, model, experiment, run; đẩy metrics và artifact.
    """

    def __init__(self, session: KWSSession):
        self._session = session
        self._api = session.api

    def register_run(
        self,
        model_name: str,
        num_classes: int,
        experiment_name: str,
        train_split_id: int,
        config: Dict[str, Any],
        *,
        model_init_id: Optional[int] = None,
        model_init_type: Optional[str] = None,
        model_description: Optional[str] = None,
        experiment_description: Optional[str] = None,
        git_commit: str = "manual-training",
        wait_job_timeout: int = 120,
        label_to_id: Optional[Dict[str, int]] = None,
    ) -> Tuple[int, int]:
        """
        Tạo model (và model_init nếu chưa có), experiment; tạo run và chờ job (nếu có).
        - label_to_id: mapping nhãn thực -> id (vd: {'one': 2, 'two': 4}); lưu vào run.config để predict trả nhãn thực.
        - Nếu truyền model_init_id: dùng luôn model_init đó, chỉ tạo model (không cần model_init_type).
        - Nếu không truyền model_init_id: cần model_init_type để tìm hoặc tạo model_init.
        Returns:
            (model_id, experiment_run_id).
        """
        api = self._api
        if model_init_id is not None:
            try:
                model_init = api.models.get_model_init(model_init_id)
            except requests.HTTPError as e:
                if e.response is not None and e.response.status_code == 404:
                    # Backend chưa có GET /model-init/{id} hoặc 404 → lấy từ list và tìm theo id
                    inits = api.models.list_model_inits(page=1, page_size=500)
                    items = inits.get("items", [])
                    model_init = next(
                        (m for m in items if (m.get("model_init_id") or m.get("id")) == model_init_id),
                        None,
                    )
                    if not model_init:
                        raise ValueError(f"model_init_id={model_init_id} không tồn tại trong danh sách.")
                else:
                    raise
            model_init_id = model_init.get("model_init_id") or model_init.get("id")
        else:
            if not model_init_type:
                raise ValueError("Cần model_init_id hoặc model_init_type")
            inits = api.models.list_model_inits(page=1, page_size=100)
            items = inits.get("items", [])
            model_init = next((m for m in items if m.get("type") == model_init_type), None)
            if not model_init:
                model_init = api.models.create_model_init(
                    name=f"{model_init_type} Architecture",
                    type=model_init_type,
                )
            model_init_id = model_init.get("model_init_id") or model_init.get("id")

        # Model
        models = api.models.list(page=1, page_size=100, name=model_name)
        model_items = models.get("items", [])
        model_obj = next((m for m in model_items if m.get("name") == model_name), None)
        if not model_obj:
            model_obj = api.models.create(
                name=model_name,
                model_init_id=model_init_id,
                num_classes=num_classes,
                description=model_description or f"{model_name}",
            )
        model_id = model_obj.get("models_id") or model_obj.get("id")

        # Experiment
        exps = api.experiments.list(page=1, page_size=100)
        exp_items = exps.get("items", [])
        exp_obj = next((e for e in exp_items if e.get("name") == experiment_name), None)
        if not exp_obj:
            exp_obj = api.experiments.create(
                name=experiment_name,
                description=experiment_description or experiment_name,
            )
        experiment_id = exp_obj.get("experiments_id") or exp_obj.get("id")

        # Run config: thêm label_to_id để predict trả nhãn thực (one, two,...) thay vì class_0, class_1
        run_config = {**config, "label_to_id": label_to_id} if label_to_id is not None else config

        # Run (async job)
        run_resp = api.experiments.create_run(
            experiment_id=experiment_id,
            model_id=model_id,
            dataset_split_id=train_split_id,
            config=run_config,
            git_commit=git_commit,
        )
        job_id = run_resp.get("job_id")
        experiment_run_id = None
        if job_id:
            status = api.jobs.wait_for_completion(job_id, timeout=wait_job_timeout)
            if status.get("status") != "failed":
                runs = api.experiments.list_runs(experiment_id)
                runs_list = runs if isinstance(runs, list) else runs.get("items", [])
                if runs_list:
                    experiment_run_id = runs_list[0].get("experiment_runs_id") or runs_list[0].get("id")

        if experiment_run_id is None:
            sync_run = api.experiments.create_run_sync(
                experiment_id=experiment_id,
                model_id=model_id,
                dataset_split_id=train_split_id,
                config=run_config,
                git_commit=git_commit,
            )
            experiment_run_id = sync_run.get("experiment_runs_id") or sync_run.get("id")
            logger.info(f"Đã tạo run đồng bộ: experiment_run_id={experiment_run_id}.")

        return model_id, experiment_run_id

    def push_metrics(
        self,
        model_id: int,
        dataset_split_id: int,
        experiment_run_id: int,
        metrics: Dict[str, Any],
        require_standard: bool = True,
    ) -> int:
        """
        Đẩy metrics lên DB (chuẩn: accuracy, precision, recall, f1_score, confusion_matrix).
        Returns:
            metrics_id.
        """
        payload = build_metrics_payload(
            model_id=model_id,
            dataset_split_id=dataset_split_id,
            experiment_run_id=experiment_run_id,
            metrics=metrics,
            require_standard=require_standard,
        )
        try:
            resp = self._api.metrics.create(payload=payload)
        except TypeError:
            resp = self._api.metrics.create(**payload)
        return resp.get("metrics_id") or resp.get("id")

    def push_artifact(
        self,
        model_id: int,
        run_id: int,
        file_path: Union[str, Path],
        artifact_type: str = "tflite",
        framework: str = "tensorflow",
        target_device: str = "CPU",
        is_default: bool = True,
        storage_path: Optional[str] = None,
    ) -> int:
        """
        Upload file artifact lên MinIO và tạo bản ghi artifact trên DB.
        storage_path: Đường dẫn object trên MinIO (vd: models/13/model.tflite). Nếu None thì dùng models/{model_id}/model.{artifact_type}.
        Returns:
            model_artifacts_id (hoặc id).
        """
        path = Path(file_path)
        if not path.exists():
            raise FileNotFoundError(f"File không tồn tại: {path}")
        art = self._api.artifacts.upload(
            file_path=path,
            model_id=model_id,
            run_id=run_id,
            artifact_type=artifact_type,
            framework=framework,
            target_device=target_device,
            is_default=is_default,
            storage_path=storage_path,
        )
        return art.get("model_artifacts_id") or art.get("id")

    def compare_metrics(
        self,
        model_ids: List[int],
        split_id: int,
    ) -> Any:
        """So sánh metrics giữa các model trên một split."""
        return self._api.metrics.compare_metrics(model_ids=model_ids, split_id=split_id)


__all__ = [
    "KWSSession",
    "DatasetPipeline",
    "ModelManager",
]
