"""Smoke tests for ai_toolkit.nlp."""

import pandas as pd

from ai_toolkit.nlp import (
    build_bow,
    build_tfidf,
    clean_pipeline,
    create_sentiment_label,
    expand_contractions,
    extract_text_stats,
    handle_nulls,
    normalize_text,
    remove_html,
    remove_noise,
)


def test_remove_html():
    assert "hello" in remove_html("<p>hello</p>").lower()
    assert "<" not in remove_html("<p>hello</p>")


def test_expand_contractions():
    out = expand_contractions("I don't like it")
    assert "do not" in out or "don't" in out  # may skip if contractions not installed


def test_remove_noise():
    out = remove_noise("Visit https://example.com 123")
    assert "https" not in out
    assert "123" not in out or " " in out


def test_normalize_text():
    assert normalize_text("  HELLO  World  ") == "hello world"


def test_clean_pipeline():
    out = clean_pipeline("  <p>Great!</p>  ")
    assert "great" in out.lower()
    assert "<" not in out


def test_handle_nulls():
    df = pd.DataFrame({"x": ["a", None, "c"]})
    out = handle_nulls(df, "x", strategy="fill", fill_value="")
    assert out["x"].isna().sum() == 0
    out2 = handle_nulls(df, "x", strategy="drop")
    assert len(out2) == 2


def test_create_sentiment_label():
    assert create_sentiment_label(1) == "negative"
    assert create_sentiment_label(3) == "neutral"
    assert create_sentiment_label(5) == "positive"


def test_build_bow():
    vec, X_train, X_test = build_bow(
        ["hello world", "world foo"],
        ["hello foo"],
        max_features=10,
    )
    assert X_train.shape[0] == 2
    assert X_test.shape[0] == 1


def test_build_tfidf():
    vec, X_train, None_val = build_tfidf(["doc one", "doc two"], None, max_features=10)
    assert X_train.shape[0] == 2
    assert None_val is None


def test_extract_text_stats():
    df = pd.DataFrame({"text": ["Hello world!", "Hi"]})
    stats = extract_text_stats(df, "text")
    assert "word_count" in stats.columns
    assert stats.shape[0] == 2


def test_get_sbert_embeddings():
    from ai_toolkit.nlp import get_sbert_embeddings
    emb = get_sbert_embeddings(["first", "second"])
    assert emb.shape[0] == 2
    assert emb.ndim == 2
