"""Tests for schema linking — TableIndex, find_relevant_tables, and the
text_to_sql safety cap. Provider-agnostic: assertions hold under both the
sentence-transformer and lexical-fallback embedding providers.
"""

from __future__ import annotations

from pathlib import Path

from lexsi_ds.agent.context import AgentContext, ColumnInfo, DatasetHandle, TableInfo
from lexsi_ds.agent.table_index import TableIndex
from lexsi_ds.agent.tools.find_relevant_tables import FindRelevantTablesArgs, _run
from lexsi_ds.agent.tools.text_to_sql import _render_schema


def _tbl(name, cols, comment=""):
    return TableInfo(name=name, columns=[ColumnInfo(c, "VARCHAR") for c in cols],
                     n_rows=10, comment=comment)


def _handle(tables, ds_id="syn"):
    return DatasetHandle(id=ds_id, kind="attached", duckdb_path=Path("/tmp/none.duckdb"),
                         tables=tables)


def _ctx(tables, ds_id="syn"):
    return AgentContext(dataset=_handle(tables, ds_id), run_id="t")


_THREE = [
    _tbl("loan_book", ["loan_id", "principal", "interest", "default_flag"], "loans issued to customers"),
    _tbl("weather_daily", ["date", "temperature", "rainfall"], "daily weather readings"),
    _tbl("customer_profile", ["customer_id", "age", "city"], "customer demographics"),
]


def test_table_index_ranks_relevant_first():
    idx = TableIndex.build(_THREE)
    hits = idx.search("average principal of loans that defaulted", top_k=3)
    assert hits, "expected at least one hit"
    assert hits[0][0].name == "loan_book"


def test_table_index_empty_query():
    assert TableIndex.build(_THREE).search("  ") == []


def test_find_relevant_tables_tool():
    ctx = _ctx(_THREE)
    res = _run(FindRelevantTablesArgs(question="loan principal default", top_k=2), ctx)
    assert res.ok
    assert res.payload["total_tables"] == 3
    assert len(res.payload["tables"]) <= 2
    assert "loan_book" in res.payload["tables"]


def test_find_relevant_tables_no_tables():
    res = _run(FindRelevantTablesArgs(question="x"), _ctx([]))
    assert not res.ok and res.error == "no_tables"


def test_index_cached_on_ctx():
    ctx = _ctx(_THREE)
    _run(FindRelevantTablesArgs(question="loans"), ctx)
    assert "table_index:syn" in ctx.cache


# ---- text_to_sql safety cap ----

def _many(n):
    tables = [_tbl(f"t{i:02d}", ["col_a", "col_b"]) for i in range(n)]
    tables.append(_tbl("sales_revenue", ["revenue", "region", "quarter"], "sales revenue by region"))
    return tables


def _count_rendered(schema: str) -> int:
    return schema.count("### TABLE")


def test_small_schema_renders_all_no_cap():
    schema = _render_schema(_ctx(_THREE), question="loans")
    assert _count_rendered(schema) == 3
    assert "auto-selected" not in schema


def test_large_schema_auto_capped():
    ctx = _ctx(_many(50), ds_id="big")
    schema = _render_schema(ctx, question="total revenue by region")
    n = _count_rendered(schema)
    assert n <= 21 and n < 51          # capped well below the 51 total
    assert "sales_revenue" in schema   # the relevant table survived the cap
    assert "auto-selected" in schema   # cap note present


def test_explicit_tables_subset_scopes_exactly():
    ctx = _ctx(_many(50), ds_id="big2")
    schema = _render_schema(ctx, question="anything", tables=["t00", "sales_revenue"])
    assert _count_rendered(schema) == 2
    assert "sales_revenue" in schema and "### TABLE `t00`" in schema


def test_bad_explicit_subset_falls_back():
    # nonexistent names → don't render an empty schema; fall back gracefully
    ctx = _ctx(_THREE, ds_id="fb")
    schema = _render_schema(ctx, question="loans", tables=["does_not_exist"])
    assert _count_rendered(schema) == 3