import sqlite3

import pandas as pd

from storai.api.story import analyze


def test_analyze_with_dataframe_inputs():
    summary_df = pd.DataFrame(
        {
            "year": [2025, 2025, 2025],
            "month": [7, 8, 9],
            "payment": [100.0, 80.0, 100.0],
            "visit_count": [50.0, 45.0, 55.0],
        }
    )
    source_df = summary_df.copy()

    result = analyze(summary_df=summary_df, source_df=source_df)

    assert result["insights"]
    assert "display_rows" in result
    assert "ml_findings" in result


def test_analyze_with_sql_inputs():
    conn = sqlite3.connect(":memory:")

    source_df = pd.DataFrame(
        {
            "year": [2025, 2025, 2025],
            "month": [7, 8, 9],
            "payment": [100.0, 80.0, 100.0],
            "visit_count": [50.0, 45.0, 55.0],
        }
    )
    source_df.to_sql("exec_summary", conn, index=False, if_exists="replace")

    summary_sql = """
        select
            year,
            month,
            sum(payment) as payment,
            sum(visit_count) as visit_count
        from exec_summary
        group by 1,2
    """
    source_sql = "select * from exec_summary"

    result = analyze(summary_sql=summary_sql, source_sql=source_sql, db_engine=conn)

    assert result["insights"]
    assert result["source_sql"] == source_sql
    assert "ml_findings" in result


def test_analyze_rejects_ambiguous_summary_inputs():
    summary_df = pd.DataFrame(
        {
            "year": [2025],
            "month": [10],
            "payment": [100.0],
        }
    )

    try:
        analyze(summary_df=summary_df, summary_sql="select 1")
        raise AssertionError("Expected ValueError for ambiguous summary inputs.")
    except ValueError:
        pass
