"""transform_column: derive a structured column from free text via the LLM,
materialize it, then query with SQL. Scripted fake LLM for determinism."""

from __future__ import annotations

import duckdb
import pytest

from lexsi_ds.agent.context import AgentContext, ColumnInfo, DatasetHandle, TableInfo
from lexsi_ds.agent.tools.transform_column import TransformColumnArgs, _run
from lexsi_ds.llm.client import LLMResult


class FakeLLM:
    name = "fake"

    def __init__(self, *texts):
        self.texts = list(texts)
        self.i = 0

    def complete(self, system, user):
        t = self.texts[min(self.i, len(self.texts) - 1)]
        self.i += 1
        return LLMResult(text=t)


def _ctx(tmp_path, rows, llm):
    p = tmp_path / "d.duckdb"
    con = duckdb.connect(str(p))
    con.execute("CREATE TABLE proj (repo VARCHAR, prose VARCHAR)")
    con.executemany("INSERT INTO proj VALUES (?, ?)", rows)
    con.close()
    ds = DatasetHandle(id="t", kind="attached", duckdb_path=p, tables=[
        TableInfo(name="proj", columns=[ColumnInfo("repo", "VARCHAR"),
                                        ColumnInfo("prose", "VARCHAR")])])
    return AgentContext(dataset=ds, run_id="t", llm=llm)


def test_extracts_number_from_prose_as_integer(tmp_path):
    rows = [("a/b", "has 38715 stars and 10 forks"), ("c/d", "currently 0 stars")]
    ctx = _ctx(tmp_path, rows, FakeLLM('["38715", "0"]'))
    res = _run(TransformColumnArgs(table="proj", source_column="prose",
                                   new_column="stars", instruction="extract the integer star count",
                                   output_type="integer"), ctx)
    assert res.ok, res.summary
    con = duckdb.connect(str(ctx.dataset.duckdb_path), read_only=True)
    rows_out = con.execute('SELECT repo, stars FROM "proj__stars" ORDER BY stars DESC').fetchall()
    typ = con.execute("SELECT data_type FROM information_schema.columns "
                      "WHERE table_name='proj__stars' AND column_name='stars'").fetchone()[0]
    con.close()
    assert rows_out[0] == ("a/b", 38715)        # numeric, sortable
    assert "INT" in typ.upper() or "BIGINT" in typ.upper()


def test_where_scopes_rows(tmp_path):
    rows = [("a/b", "100 stars"), ("c/d", "200 stars")]
    ctx = _ctx(tmp_path, rows, FakeLLM('["100"]'))
    res = _run(TransformColumnArgs(table="proj", source_column="prose", new_column="stars",
                                   instruction="extract stars", where="repo = 'a/b'",
                                   output_type="integer"), ctx)
    assert res.ok and res.payload["n_rows"] == 1


def test_classification_labels(tmp_path):
    rows = [("x", "stocks tumble on wall street"), ("y", "team wins the cup")]
    ctx = _ctx(tmp_path, rows, FakeLLM('["Business", "Sports"]'))
    res = _run(TransformColumnArgs(table="proj", source_column="prose", new_column="category",
                                   instruction="classify into World/Sports/Business/SciTech"), ctx)
    assert res.ok
    con = duckdb.connect(str(ctx.dataset.duckdb_path), read_only=True)
    cats = dict(con.execute('SELECT repo, category FROM "proj__category"').fetchall())
    con.close()
    assert cats == {"x": "Business", "y": "Sports"}


def test_no_llm_errors(tmp_path):
    ctx = _ctx(tmp_path, [("a", "x")], FakeLLM("[]"))
    ctx.llm = None
    res = _run(TransformColumnArgs(table="proj", source_column="prose",
                                   new_column="c", instruction="x"), ctx)
    assert not res.ok and res.error == "no_llm"


def test_bad_new_column_rejected(tmp_path):
    ctx = _ctx(tmp_path, [("a", "x")], FakeLLM("[]"))
    res = _run(TransformColumnArgs(table="proj", source_column="prose",
                                   new_column="bad name", instruction="x"), ctx)
    assert not res.ok and res.error == "bad_identifier"


def test_short_array_padded_with_none(tmp_path):
    # LLM returns fewer outputs than inputs → missing rows become null, no crash
    rows = [("a", "10 stars"), ("b", "20 stars")]
    ctx = _ctx(tmp_path, rows, FakeLLM('["10"]'))
    res = _run(TransformColumnArgs(table="proj", source_column="prose", new_column="stars",
                                   instruction="extract stars", output_type="integer"), ctx)
    assert res.ok and res.payload["n_nonnull"] == 1