from __future__ import annotations

from collections.abc import Sequence
from pathlib import Path
from typing import Any, Literal

import pandas as pd
from tqdm import tqdm

from ..utils.balancing import balance_dataframe, stratified_cut
from .builder import build_sentence
from .io import coerce_to_df

REQUIRED = {"protocol", "label"}

BalanceStrategy = Literal["none", "downsample", "cap"]


def _describe_label(label: str) -> str:
    """
    Convert a label into a human-readable description.
    Back-compat:
      - 'benign'/'normal' -> benign description
      - anything else      -> malicious description
    """
    if isinstance(label, str) and label.strip().lower() in {"benign", "normal"}:
        return "Benign flow"
    return "Malicious flow"


def generate_sentences(
    data: pd.DataFrame | str | bytes | Path,
    *,
    sep: str = ",",
    encoding: str | None = "utf-8",
    dtype: dict[str, Any] | None = None,
    parse_dates: list[str] | None = None,
    na_values: list[str] | None = None,
    nrows: int | None = None,
    # balancing
    balance: BalanceStrategy | bool = "none",
    balance_labels: Sequence[str] | None = None,
    per_class_cap: int | None = None,
    random_state: int = 42,
    # others
    max_items: int | None = None,
    show_progress: bool = False,
) -> pd.DataFrame:
    """
    Build a DataFrame with exactly three columns:
      - sentence
      - description
      - label

    Processing order:
      1) Load (CSV/DataFrame)
      2) Validate required columns
      3) Balance (if requested)
      4) Deterministic shuffle
      5) Stratified cut (if max_items is provided)
      6) Build sentences
    """
    # 1) CSV/DataFrame → DataFrame (uses pandas kwargs) 
    df = coerce_to_df(
        data,
        sep=sep,
        encoding=encoding,
        dtype=dtype,
        parse_dates=parse_dates,
        na_values=na_values,
        nrows=nrows,
    )

    # 2) Required columns present? 
    if not set(df.columns) >= REQUIRED:
        miss = sorted(REQUIRED - set(df.columns))
        raise ValueError(f"Missing required columns: {miss}")

    # 3) Normalize balance option
    if isinstance(balance, bool):
        strategy: BalanceStrategy = "downsample" if balance else "none"
    else:
        strategy = balance

    # 4) Balance FIRST (no target_total here), then deterministic shuffle
    if strategy != "none":
        df = balance_dataframe(
            df,
            strategy=strategy,
            labels=balance_labels,
            per_class_cap=per_class_cap,
            random_state=random_state,
        )
        df = df.sample(frac=1.0, random_state=random_state).reset_index(drop=True)

    # 5) STRATIFIED cut AFTER balancing (so even totals split evenly when possible)
    if max_items is not None and len(df) > max_items:
        df = stratified_cut(
            df,
            labels=balance_labels,      # if provided, respect that subset
            total=max_items,
            random_state=random_state,
        )

    # 6) Sentence + description generation (progress bar preserved) 
    iterator = df.itertuples(index=False)
    if show_progress:
        iterator = tqdm(iterator, total=len(df), desc="Generating sentences", unit="row")

    sentences = [build_sentence(r) for r in iterator]
    descriptions = [_describe_label(lbl) for lbl in df["label"]]

    return pd.DataFrame({
        "sentence": sentences,
        "response": descriptions,
        "label": df["label"].values,
    })
