"""
Includes coverage for:
- API balance and cut logic
- Error and validation handling
- Builder and sentence generation variants
"""

import io

import pandas as pd
import pytest

import flownarrator as fn
from flownarrator.errors import (
    ColumnValidationError,
    CSVReadError,
    FileNotFoundErrorFN,
    FlowNarratorError,
)
from flownarrator.flow_sentences.builder import build_sentence
from flownarrator.flow_sentences.io import coerce_to_df
from flownarrator.io.csv_loader import load_csv
from flownarrator.utils.validation import ensure_required_columns_exist, validate_dataframe_columns

# ----------------------------------------------------------------------------
# API balance / cap path
# ----------------------------------------------------------------------------

def test_api_balance_cap_and_cut():
    """Tests per-class cap and cut behavior in generate_sentences API."""
    df = pd.DataFrame({"protocol": ["TCP"] * 6, "label": ["a"] * 3 + ["b"] * 2 + ["c"] * 1})
    out = fn.generate_sentences(df, balance="cap", per_class_cap=1, max_items=2, random_state=0)
    assert len(out) == 2
    # Expect at least one unique label, ideally multiple
    assert out["label"].nunique() >= 1


# ----------------------------------------------------------------------------
# Error hierarchy and validation
# ----------------------------------------------------------------------------

def test_error_hierarchy_and_instantiation():
    for cls in (FileNotFoundErrorFN, CSVReadError, ColumnValidationError):
        e = cls("msg")
        assert isinstance(e, FlowNarratorError)


def test_coerce_to_df_all_branches(tmp_path):
    """Covers all input forms: DataFrame, path, and file-like object."""
    path = tmp_path / "x.csv"
    path.write_text("protocol,label\nTCP,benign\n", encoding="utf-8")
    buf = io.StringIO("protocol,label\nUDP,malicious\n")

    df1 = coerce_to_df(pd.DataFrame({"protocol": ["TCP"], "label": ["benign"]}))
    df2 = coerce_to_df(str(path))
    df3 = coerce_to_df(buf)

    assert all(len(df) >= 1 for df in (df1, df2, df3))

    with pytest.raises(TypeError):
        coerce_to_df(object())


def test_load_csv_branches(tmp_path, monkeypatch):
    """Tests CSV loading paths: happy path, missing file, parse errors, missing columns."""
    csv_path = tmp_path / "good.csv"
    csv_path.write_text("protocol,label,dst_port\nTCP,benign,443\n", encoding="utf-8")
    df = load_csv(csv_path, columns=["protocol", "label"])
    assert list(df.columns) == ["protocol", "label"]

    # Missing file: raise and warn behavior
    missing = tmp_path / "nope.csv"
    with pytest.raises(FileNotFoundErrorFN):
        load_csv(missing)
    load_csv(missing, on_error="warn")  # Should not raise

    # Simulate parse error
    original_read_csv = pd.read_csv

    def boom(*_, **__):
        raise ValueError("mock parse error")

    monkeypatch.setattr(pd, "read_csv", boom, raising=True)
    with pytest.raises(CSVReadError):
        load_csv(csv_path)

    # Restore original pandas.read_csv
    monkeypatch.setattr(pd, "read_csv", original_read_csv, raising=True)

    # Missing columns: raise and warn path
    csv2 = tmp_path / "cols.csv"
    csv2.write_text("protocol,label\nTCP,benign\n", encoding="utf-8")
    with pytest.raises(ColumnValidationError):
        load_csv(csv2, columns=["protocol", "label", "dst_port"])

    df2 = load_csv(csv2, columns=["protocol", "label", "dst_port"], on_error="warn")
    assert list(df2.columns) == ["protocol", "label"]


def test_validation_helpers():
    """Ensure validation utilities detect missing columns and raise as expected."""
    cols = ["protocol", "label"]
    assert ensure_required_columns_exist(cols, cols) == []

    missing = ensure_required_columns_exist(cols, [*cols, "extra"])
    assert "extra" in missing

    df_ok = pd.DataFrame({"protocol": [6], "label": ["benign"]})
    validate_dataframe_columns(df_ok, cols)

    with pytest.raises(ColumnValidationError):
        validate_dataframe_columns(pd.DataFrame({"protocol": [6]}), cols)


# ----------------------------------------------------------------------------
# Builder and sentence generation
# ----------------------------------------------------------------------------

def test_builder_variants_optional_fields():
    """Check multiple builder scenarios with varying completeness."""
    # Full TCP example
    df1 = pd.DataFrame({
        "protocol": ["TCP"],
        "label": ["benign"],
        "src_ip": ["10.0.0.5"],
        "dst_ip": ["8.8.8.8"],
        "src_port": [12345],
        "dst_port": [443],
        "duration": [0.5],
        "total_payload_bytes": [1024],
        "packets_count": [10],
        "bytes_rate": [5_000.0],
        "packets_rate": [5.0],
        "tcp_flags": ["SYN,ACK"],
        "handshake_state": ["handshake completed"],
    })
    s1 = build_sentence(next(df1.itertuples(index=False)))
    assert isinstance(s1, str) and len(s1) > 10

    # UDP minimal fields
    df2 = pd.DataFrame({"protocol": ["UDP"], "label": ["malicious"]})
    s2 = build_sentence(next(df2.itertuples(index=False)))
    assert isinstance(s2, str)

    # ICMPv6 high payload variant
    df3 = pd.DataFrame({
        "protocol": ["ICMPv6"],
        "label": ["benign"],
        "total_payload_bytes": [20_000],
        "packets_count": [400],
        "duration": [120.0],
    })
    s3 = build_sentence(next(df3.itertuples(index=False)))
    assert isinstance(s3, str)
