import pandas as pd
import pytest

import flownarrator as fn
from flownarrator.errors import FileNotFoundErrorFN
from flownarrator.flow_sentences import (
    builder as bld,
    protocol as prt,
    quantitative as qnt,
    roles as rl,
    services as srv,
    tcp as tcpmod,
)
from flownarrator.io import csv_loader as csvio
from flownarrator.utils import balancing as bal

# ---------------------------------------------------------------------------
# roles.py — cover True / False / None branches for src/dst phrases
# ---------------------------------------------------------------------------

@pytest.mark.parametrize(
    "ip,exp_snip_src,exp_snip_dst",
    [
        ("10.0.0.1", "internal host", "internal server"),   # private -> True
        ("8.8.8.8", "external system", "server on the Internet"),  # public -> False
        ("not_an_ip", "unidentified host", "unidentified server"), # invalid -> None path
    ],
)
def test_roles_phrase_all_branches(ip, exp_snip_src, exp_snip_dst):
    s = rl.src_role_phrase(ip)
    d = rl.dst_role_phrase(ip)
    assert isinstance(s, str) and isinstance(d, str)
    assert exp_snip_src in s
    assert exp_snip_dst in d


# ---------------------------------------------------------------------------
# builder.join_clauses — empty, whitespace, punctuation trimming
# NOTE: your implementation uses 'and' before the last item.
# ---------------------------------------------------------------------------

@pytest.mark.parametrize(
    "parts, expected",
    [
        ([], ""),
        (["  ", ""], ""),
        (["Hello.", " world  "], "Hello and world"),
        (["One.", "Two.", "  Three  "], "One, Two and Three"),
    ],
)
def test_join_clauses_variants(parts, expected):
    assert bld.join_clauses(parts) == expected


# ---------------------------------------------------------------------------
# protocol.py — normalization + context flags
# NOTE: normalize_protocol expects strings; unknowns get upper-cased as-is.
# protocol_context provides article/flags/ports/handshake booleans.
# ---------------------------------------------------------------------------

@pytest.mark.parametrize("raw, norm", [("tcp", "TCP"), ("UDP", "UDP"), ("ICMP", "ICMP"), ("weird", "WEIRD")])
def test_protocol_normalize_and_context(raw, norm):
    p = prt.normalize_protocol(raw)
    assert p == norm
    ctx = prt.protocol_context(raw)
    for k in ("article", "has_ports", "has_tcp_flags", "has_tcp_handshake"):
        assert k in ctx


# ---------------------------------------------------------------------------
# quantitative.py — boundary buckets (hit the tricky edges)
# ---------------------------------------------------------------------------

@pytest.mark.parametrize("dur, expect_nonempty", [(None, False), (0.0, True), (0.049, True), (0.05, True), (10.0, True)])
def test_duration_phrase_edges(dur, expect_nonempty):
    s = qnt.duration_phrase(dur)
    assert isinstance(s, str)
    if expect_nonempty:
        assert s != ""

@pytest.mark.parametrize("val", [0, 1, 512, 4096, 10_000_000])
def test_bytes_and_packets_phrases(val):
    assert isinstance(qnt.bytes_phrase(val), str)
    assert isinstance(qnt.packets_phrase(val), str)

@pytest.mark.parametrize("br, pr", [(0, 0), (100, 1), (10_000, 100), (2_000_000, 5_000)])
def test_rate_phrases(br, pr):
    s1 = qnt.rate_phrase(bytes_rate=br, packets_rate=pr)
    assert isinstance(s1, str)


# ---------------------------------------------------------------------------
# services.py — port categories + dst service phrase coverage
# ---------------------------------------------------------------------------

@pytest.mark.parametrize("p", [None, 0, 22, 53, 80, 123, 443, 8080, 5353, 3389, 9999])
def test_services_src_port_category_and_dst_phrase(p):
    cat = srv.src_port_category(p)
    assert isinstance(cat, str)
    # Ensure dst_service_phrase doesn't crash on common and odd ports
    out = srv.dst_service_phrase(p)
    assert isinstance(out, str)


# ---------------------------------------------------------------------------
# tcp.py — flags and handshake stories including empties
# ---------------------------------------------------------------------------

@pytest.mark.parametrize("flags", [None, "", "SYN", "SYN,ACK", "FIN,ACK", "PSH,ACK", "X"])
def test_tcp_flags_story_variants(flags):
    s = tcpmod.tcp_flags_story(flags)
    assert isinstance(s, str)

@pytest.mark.parametrize("hs", [None, "", "handshake completed", "HANDSHAKE PARTIAL", "no handshake", "error", "X"])
def test_tcp_handshake_story_variants(hs):
    s = tcpmod.handshake_story(hs)
    assert isinstance(s, str)


# ---------------------------------------------------------------------------
# io.csv_loader — errors & minimal happy path (using actual filesystem paths)
# ---------------------------------------------------------------------------

def test_csv_loader_error_and_happy_path(tmp_path):
    # Non-existent path -> should raise specific project error
    with pytest.raises(FileNotFoundErrorFN):
        csvio.load_csv(tmp_path / "missing.csv")


# ---------------------------------------------------------------------------
# utils.balancing — extra edge cases
# NOTE: stratified_cut in your code may *increase* size to reach 'total'.
# We only assert basic invariants instead of exact length.
# ---------------------------------------------------------------------------

def test_balance_cap_and_downsample_edge_cases():
    df = pd.DataFrame({"protocol": ["TCP"]*6, "label": ["a","a","a","b","b","c"]})
    # cap=1 -> at most one per class
    out = bal.balance_cap(df, per_class_cap=1, random_state=0)
    assert set(out["label"]) <= {"a","b","c"} and out["label"].value_counts().max() <= 1

    # invalid labels subset should raise
    with pytest.raises(ValueError):
        bal.balance_downsample(df, labels=["a", "d"], random_state=0)

    # stratified_cut: result should not exceed the requested total and must preserve labels set
    out2 = bal.stratified_cut(df, labels=df["label"], total=14, random_state=0)
    assert set(out2["label"]) == set(df["label"])
    assert len(out2) <= 14


# ---------------------------------------------------------------------------
# flow_sentences.api.generate_sentences — hit progress + response mapping
# ---------------------------------------------------------------------------

def test_generate_sentences_variants_and_responses():
    df = pd.DataFrame({
        "protocol": ["TCP","UDP","TCP","UDP"],
        "label": ["benign","malicious","malicious","benign"],
        "src_ip": ["10.0.0.5","8.8.8.8","10.0.0.7","1.2.3.4"],
        "dst_ip": ["8.8.8.8","10.0.0.1","1.1.1.1","10.0.0.9"],
        "src_port": [12345, 34567, 22, 80],
        "dst_port": [443, 53, 3389, 9999],
        "duration": [0.05, 3.0, 0.5, 15.0],
        "total_payload_bytes": [200, 5000, 2048, 0],
        "packets_count": [5, 50, 200, 0],
        "bytes_rate": [1000, 500_000, 2_000_000, 0],
        "packets_rate": [10, 5, 20, 0],
        "tcp_flags": ["SYN,ACK", "", "FIN,ACK", None],
        "handshake_state": ["handshake completed", "no handshake", "HANDSHAKE PARTIAL", "error"],
    })
    out = fn.generate_sentences(df, balance="cap", per_class_cap=10, show_progress=False, random_state=7)
    assert {"sentence","response","label"} <= set(out.columns)
    assert all(isinstance(x, str) for x in out["response"])
    # Ensure both classes are present after cap
    assert set(out["label"]) == {"benign","malicious"}


def test_generate_sentences_e2e_and_balance():
    df = pd.DataFrame({
        "protocol": ["TCP"]*4,
        "label": ["benign", "benign", "malicious", "malicious"],
        "src_ip": ["10.0.0.5"]*4,
        "dst_ip": ["8.8.8.8"]*4,
        "src_port": [12345]*4,
        "dst_port": [443, 8080, 53, 9999],
        "duration": [0.05, 3.0, 15.0, 0.5],
        "total_payload_bytes": [200, 5000, 0, 2048],
        "packets_count": [5, 50, 0, 200],
        "bytes_rate": [1000, 500_000, 0, 2_000_000],
        "packets_rate": [10, 5, 0, 20],
        "tcp_flags": ["SYN,ACK", "FIN,ACK", "", None],
        "handshake_state": ["handshake completed", "HANDSHAKE PARTIAL", "no handshake", "error"],
    })
    out = fn.generate_sentences(df, balance=True, max_items=2, random_state=0)
    assert len(out) <= 2
    assert {"sentence", "response", "label"} <= set(out.columns)
    assert all(isinstance(x, str) for x in out["sentence"])
