"""
Utilities to create dataset splits using pandas, exposed as part of kwslib.training.

This is essentially the same logic as the top-level create_dataset_split.py script,
but packaged inside kwslib so that code running only with "pip install kwslib"
can import:

    from kwslib.training.create_dataset_split import get_split_data, push_splits
"""

import logging
import os
from typing import Optional, Dict, List, Any

import numpy as np
import pandas as pd

from kwslib import KWSClient

logging.basicConfig(level=logging.INFO)
logger = logging.getLogger(__name__)


def get_split_data(
    api: KWSClient,
    dataset_version_id: int,
    feature_type_id: Optional[int] = None,
    keyword_ids: Optional[List[int]] = None,
) -> pd.DataFrame:
    """
    Get data for creating dataset splits.

    Returns a DataFrame with columns:
    - feature_id: Feature ID (keyword_audio_features.id)
    - sample_id: Audio sample ID (keyword_audio_samples.id)
    - keyword_id: Keyword ID
    - keyword_name: Keyword name (label)
    - file_name: Feature file name
    - minio_path: MinIO path to feature file
    - feature_type_id: Feature type ID
    """
    logger.info(f"Fetching data for dataset version {dataset_version_id}...")

    df = api.dataset_splits.get_mfcc_files(
        dataset_version_id=dataset_version_id,
        feature_type_id=feature_type_id,
        keyword_ids=keyword_ids,
    )

    if df.empty:
        raise ValueError("No data found for the specified criteria")

    logger.info(f"Retrieved {len(df)} records")
    logger.info(f"Keywords: {df['keyword_name'].unique().tolist()}")
    logger.info(f"Keyword counts:\n{df['keyword_name'].value_counts()}")

    return df


def push_splits(
    api: KWSClient,
    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,
    batch_size: int = 1000,
) -> Dict[str, int]:
    """
    Push split assignments to the database.

    Args:
        api: KWSClient instance (must be authenticated)
        dataset_version_id: Dataset version ID
        config_name: Configuration name (e.g., "config_70_15_15")
        splits: Dictionary mapping split names to DataFrames.
                Each DataFrame MUST have 'feature_id', 'keyword_id', and 'derivative_label' columns.
        create_merged_keywords: Optional dict mapping new keyword names to
            list of source keyword IDs to merge.
        auto_detect_merged_keywords: If True, automatically detect merged keywords.
        batch_size: Number of files to send per batch.
    """
    created_splits: Dict[str, int] = {}

    required_columns = ["feature_id", "keyword_id", "derivative_label"]
    for split_name, split_df in splits.items():
        if split_df.empty:
            continue
        missing_columns = [col for col in required_columns if col not in split_df.columns]
        if missing_columns:
            raise ValueError(
                f"DataFrame for split '{split_name}' is missing required columns: {missing_columns}. "
                f"Required columns: {required_columns}"
            )

    detected_merged_keywords = create_merged_keywords
    if auto_detect_merged_keywords and create_merged_keywords is None:
        first_split_name = list(splits.keys())[0]
        first_split_df = splits[first_split_name]

        if not first_split_df.empty and "keyword_name" in first_split_df.columns:
            keyword_groups = first_split_df.groupby("keyword_name")["keyword_id"].unique()

            merged_keywords: Dict[str, List[int]] = {}
            for keyword_name, keyword_ids in keyword_groups.items():
                unique_ids = sorted(keyword_ids.tolist())
                if len(unique_ids) > 1:
                    merged_keywords[keyword_name] = unique_ids
                    logger.info(
                        "Auto-detected merged keyword '%s' from source IDs: %s",
                        keyword_name,
                        unique_ids,
                    )

            if merged_keywords:
                detected_merged_keywords = merged_keywords
                logger.info("Auto-detected %d merged keyword(s)", len(merged_keywords))

    for split_name, split_df in splits.items():
        if split_df.empty:
            logger.warning("Skipping empty %s split", split_name)
            continue

        logger.info("\nCreating %s split (%d files)...", split_name, len(split_df))

        files: List[Dict[str, Any]] = []
        for idx, (_, row) in enumerate(split_df.iterrows()):
            try:
                if "feature_id" not in row.index:
                    raise KeyError("feature_id")
                if "keyword_id" not in row.index:
                    raise KeyError("keyword_id")
                if "derivative_label" not in row.index:
                    raise KeyError("derivative_label")

                file_assignment: Dict[str, Any] = {
                    "feature_id": int(row["feature_id"]),
                    "keyword_id": int(row["keyword_id"]),
                }

                derivative_label_value = row["derivative_label"]
                if pd.notna(derivative_label_value) and str(derivative_label_value).strip():
                    file_assignment["derivative_label"] = str(derivative_label_value).strip()
                else:
                    logger.warning(
                        "Row %d has NaN or empty derivative_label, setting to None. Value was: %r",
                        idx,
                        derivative_label_value,
                    )
                    file_assignment["derivative_label"] = None

                if idx < 3:
                    logger.info(
                        "Row %d: feature_id=%s, keyword_id=%s, derivative_label=%s",
                        idx,
                        file_assignment["feature_id"],
                        file_assignment["keyword_id"],
                        file_assignment["derivative_label"],
                    )

                files.append(file_assignment)
            except KeyError as e:
                logger.error("Error: Missing required column '%s' in row %d", e, idx)
                logger.error("Available columns: %s", list(row.index))
                raise ValueError(
                    f"DataFrame must contain 'derivative_label' column. "
                    f"Available columns: {list(split_df.columns)}"
                ) from e
            except Exception as e:  # pragma: no cover - logging path
                logger.error("Error converting row %d: %s", idx, e)
                logger.error("Row data: %s", row.to_dict())
                raise

        total_files = len(files)
        original_timeout = api.timeout
        try:
            if total_files > 1000:
                estimated_timeout = max(60, (total_files // 100) + 30)
                api.timeout = estimated_timeout
                logger.info("  Increased timeout to %ds for %d files", estimated_timeout, total_files)

            if files:
                logger.info("DEBUG: First file in request: %s", files[0])

            merged_keywords = detected_merged_keywords if split_name == list(splits.keys())[0] else None

            response = 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=merged_keywords,
            )
            created_splits[split_name] = response["dataset_splits_id"]
            logger.info("  ✅ Created %s split (ID: %s)", split_name, created_splits[split_name])
        finally:
            api.timeout = original_timeout

    return created_splits


def create_class_splits_with_cap(
    api: KWSClient,
    dataset_version_id: int,
    feature_type_id: Optional[int],
    config_name: str,
    max_per_class: int = 1000,
    train_ratio: float = 0.7,
    val_ratio: float = 0.15,
    test_ratio: float = 0.15,
    class_names: Optional[List[str]] = None,
) -> Dict[str, int]:
    """
    Wrapper high-level:

    1. Gọi get_split_data(...) để lấy DataFrame feature/keyword.
    2. Giới hạn tối đa max_per_class mẫu / keyword (derivative_label).
    3. Chia class thành train/val/test theo tỷ lệ train_ratio/val_ratio/test_ratio.
    4. Gọi push_splits(...) để tạo split trên server.

    Trả về:
        {"train": train_split_id, "val": val_split_id, "test": test_split_id}
    """
    if abs(train_ratio + val_ratio + test_ratio - 1.0) > 1e-6:
        raise ValueError("train_ratio + val_ratio + test_ratio must equal 1.0")

    df = get_split_data(
        api=api,
        dataset_version_id=dataset_version_id,
        feature_type_id=feature_type_id,
        keyword_ids=None,
    )

    # Áp dụng class_names nếu có
    if class_names:
        df = df[df["keyword_name"].isin(class_names)].copy()
        logger.info("Filtered to %d classes from CLASS_NAMES", len(class_names))

    # derivative_label mặc định = keyword_name
    if "derivative_label" not in df.columns:
        df["derivative_label"] = df["keyword_name"]

    # Giới hạn số mẫu mỗi class
    df = (
        df.sort_values("sample_id")
        .groupby("derivative_label", group_keys=False)
        .head(max_per_class)
    )
    logger.info("After cap %d samples/class:", max_per_class)
    logger.info("%s", df["derivative_label"].value_counts())

    # Chia class thành train/val/test
    unique_classes = df["derivative_label"].unique().tolist()
    n_total = len(unique_classes)
    n_train = max(1, int(n_total * train_ratio))
    n_val = max(1, int(n_total * val_ratio))
    n_test = max(1, n_total - n_train - n_val)

    train_classes = unique_classes[:n_train]
    val_classes = unique_classes[n_train : n_train + n_val]
    test_classes = unique_classes[n_train + n_val : n_train + n_val + n_test]

    train_df = df[df["derivative_label"].isin(train_classes)].copy()
    val_df = df[df["derivative_label"].isin(val_classes)].copy()
    test_df = df[df["derivative_label"].isin(test_classes)].copy()

    logger.info("Train classes (%d): %s", len(train_classes), train_classes)
    logger.info("Val classes   (%d): %s", len(val_classes), val_classes)
    logger.info("Test classes  (%d): %s", len(test_classes), test_classes)

    created = push_splits(
        api=api,
        dataset_version_id=dataset_version_id,
        config_name=config_name,
        splits={"train": train_df, "val": val_df, "test": test_df},
    )
    return created

