"""Tests for the standard eval dataset schema (Spec 0010)."""

from __future__ import annotations

import json
from pathlib import Path

import pytest

from agentbyte.dataset import load_dataset
from agentbyte.dataset.config import LocalDatasetConfig
from agentbyte.eval import (
    EvalDatasetRecord,
    RAGEvalContext,
    SQLEvalContext,
    save_eval_dataset,
    tasks_from_eval_dataset,
)


# ---------------------------------------------------------------------------
# Fixtures
# ---------------------------------------------------------------------------


def _make_sql_record(idx: int = 1) -> EvalDatasetRecord:
    return EvalDatasetRecord(
        id=f"asmd_{idx}",
        task=f"What is the status of CW{idx}?",
        scenario="No Files Selected",
        expected_response=f"Status {idx}.",
        eval_criteria=["Correct status", "No hallucinations"],
        context=SQLEvalContext(
            reference_sql=f"SELECT status FROM table WHERE id={idx}",
            reference_table="contracts",
            context_id=f"CW{idx}",
        ).model_dump(exclude_none=True),
    )


def _make_rag_record(idx: int = 1) -> EvalDatasetRecord:
    return EvalDatasetRecord(
        id=f"rag_{idx}",
        task=f"What is mentioned in document {idx}?",
        scenario="Single document selected",
        expected_response=f"Answer {idx}.",
        eval_criteria=["Identifies key fact", "No hallucinations"],
        context=RAGEvalContext(
            doc_ids=[f"Doc{idx}"],
            file_names=[f"file_{idx}.pdf"],
            reference_section=f"Page {idx}",
            context_id=f"CW{idx}",
        ).model_dump(exclude_none=True),
    )


# ---------------------------------------------------------------------------
# EvalDatasetRecord defaults
# ---------------------------------------------------------------------------


def test_eval_dataset_record_defaults() -> None:
    record = EvalDatasetRecord(id="t1", task="hello")
    assert record.scenario == ""
    assert record.expected_response is None
    assert record.eval_criteria is None
    assert record.context == {}


def test_eval_dataset_record_is_frozen() -> None:
    record = EvalDatasetRecord(id="t1", task="hello")
    with pytest.raises(Exception):
        record.task = "changed"  # type: ignore[misc]


# ---------------------------------------------------------------------------
# Context helpers
# ---------------------------------------------------------------------------


def test_rag_eval_context_roundtrip() -> None:
    ctx = RAGEvalContext(
        doc_ids=["Doc1", "Doc2"],
        file_names=["a.pdf"],
        reference_section="Page 3",
        context_id="CW999",
    )
    dumped = ctx.model_dump(exclude_none=True)
    restored = RAGEvalContext(**dumped)
    assert restored.doc_ids == ["Doc1", "Doc2"]
    assert restored.context_id == "CW999"


def test_sql_eval_context_roundtrip() -> None:
    ctx = SQLEvalContext(
        reference_sql="SELECT 1",
        reference_table="my_table",
        context_id="CW42",
    )
    dumped = ctx.model_dump(exclude_none=True)
    restored = SQLEvalContext(**dumped)
    assert restored.reference_sql == "SELECT 1"
    assert restored.db_schema is None


# ---------------------------------------------------------------------------
# save_eval_dataset
# ---------------------------------------------------------------------------


def test_save_eval_dataset_writes_correct_files(tmp_path: Path) -> None:
    records = [_make_sql_record()]
    save_eval_dataset(records, root=tmp_path, dataset_id="test/sql")

    dataset_dir = tmp_path / "test" / "sql"
    assert (dataset_dir / "config.json").exists()
    assert (dataset_dir / "data.json").exists()
    assert (dataset_dir / "schema.json").exists()


def test_save_eval_dataset_config_has_correct_engine(tmp_path: Path) -> None:
    save_eval_dataset([_make_sql_record()], root=tmp_path, dataset_id="test/sql")
    config = json.loads((tmp_path / "test" / "sql" / "config.json").read_text())
    assert config["engine"] == "sqlite"
    assert config["table_name"] == "eval"
    assert config["in_memory"] is True


def test_save_eval_dataset_custom_table_name(tmp_path: Path) -> None:
    save_eval_dataset(
        [_make_sql_record()], root=tmp_path, dataset_id="test/sql", table_name="tasks"
    )
    config = json.loads((tmp_path / "test" / "sql" / "config.json").read_text())
    assert config["table_name"] == "tasks"


def test_save_eval_dataset_serialises_context_as_json_string(tmp_path: Path) -> None:
    save_eval_dataset([_make_sql_record()], root=tmp_path, dataset_id="test/sql")
    data = json.loads((tmp_path / "test" / "sql" / "data.json").read_text())
    row = data["records"][0]
    # context must be a JSON string (TEXT column), not a nested dict
    assert isinstance(row["context"], str)
    parsed = json.loads(row["context"])
    assert parsed["reference_sql"] == "SELECT status FROM table WHERE id=1"


def test_save_eval_dataset_serialises_list_criteria_as_json_string(tmp_path: Path) -> None:
    save_eval_dataset([_make_sql_record()], root=tmp_path, dataset_id="test/sql")
    data = json.loads((tmp_path / "test" / "sql" / "data.json").read_text())
    row = data["records"][0]
    assert isinstance(row["eval_criteria"], str)
    parsed = json.loads(row["eval_criteria"])
    assert parsed == ["Correct status", "No hallucinations"]


def test_save_eval_dataset_raises_on_invalid_dataset_id(tmp_path: Path) -> None:
    with pytest.raises(ValueError, match="category/name"):
        save_eval_dataset([], root=tmp_path, dataset_id="no_slash")


def test_save_eval_dataset_accepts_json_engine(tmp_path: Path) -> None:
    save_eval_dataset(
        [_make_sql_record()],
        root=tmp_path,
        dataset_id="test/sql",
        engine="json",
    )
    config = json.loads((tmp_path / "test" / "sql" / "config.json").read_text())
    assert config["engine"] == "json"
    # schema.json is not written for non-SQLite engines
    assert not (tmp_path / "test" / "sql" / "schema.json").exists()


# ---------------------------------------------------------------------------
# tasks_from_eval_dataset — core column mapping
# ---------------------------------------------------------------------------


def test_tasks_from_eval_dataset_maps_core_columns(tmp_path: Path) -> None:
    save_eval_dataset([_make_sql_record(1)], root=tmp_path, dataset_id="t/s")
    dataset = load_dataset("t/s", LocalDatasetConfig(local_root=tmp_path))
    tasks = tasks_from_eval_dataset(dataset)

    assert len(tasks) == 1
    t = tasks[0]
    assert t.name == "asmd_1"
    assert t.input == "What is the status of CW1?"
    assert t.expected_output == "Status 1."


def test_tasks_from_eval_dataset_criteria_stored_under_underscore_key(tmp_path: Path) -> None:
    save_eval_dataset([_make_sql_record()], root=tmp_path, dataset_id="t/s")
    dataset = load_dataset("t/s", LocalDatasetConfig(local_root=tmp_path))
    tasks = tasks_from_eval_dataset(dataset)

    assert tasks[0].metadata["_criteria"] == ["Correct status", "No hallucinations"]


def test_tasks_from_eval_dataset_scenario_in_metadata(tmp_path: Path) -> None:
    save_eval_dataset([_make_sql_record()], root=tmp_path, dataset_id="t/s")
    dataset = load_dataset("t/s", LocalDatasetConfig(local_root=tmp_path))
    tasks = tasks_from_eval_dataset(dataset)

    assert tasks[0].metadata["scenario"] == "No Files Selected"


def test_tasks_from_eval_dataset_context_unpacked_into_metadata(tmp_path: Path) -> None:
    save_eval_dataset([_make_sql_record(1)], root=tmp_path, dataset_id="t/s")
    dataset = load_dataset("t/s", LocalDatasetConfig(local_root=tmp_path))
    tasks = tasks_from_eval_dataset(dataset)

    assert tasks[0].metadata["reference_sql"] == "SELECT status FROM table WHERE id=1"
    assert tasks[0].metadata["reference_table"] == "contracts"
    assert tasks[0].metadata["context_id"] == "CW1"


def test_tasks_from_eval_dataset_rag_doc_ids_accessible(tmp_path: Path) -> None:
    save_eval_dataset([_make_rag_record(7)], root=tmp_path, dataset_id="t/r")
    dataset = load_dataset("t/r", LocalDatasetConfig(local_root=tmp_path))
    tasks = tasks_from_eval_dataset(dataset)

    assert tasks[0].metadata["doc_ids"] == ["Doc7"]
    assert tasks[0].metadata["file_names"] == ["file_7.pdf"]
    assert tasks[0].metadata["reference_section"] == "Page 7"


def test_tasks_from_eval_dataset_missing_optional_fields_produce_none(tmp_path: Path) -> None:
    record = EvalDatasetRecord(id="bare", task="hello")
    save_eval_dataset([record], root=tmp_path, dataset_id="t/b")
    dataset = load_dataset("t/b", LocalDatasetConfig(local_root=tmp_path))
    tasks = tasks_from_eval_dataset(dataset)

    t = tasks[0]
    assert t.expected_output is None
    assert "_criteria" not in t.metadata
    assert "scenario" not in t.metadata


def test_tasks_from_eval_dataset_plain_string_criteria_preserved(tmp_path: Path) -> None:
    record = EvalDatasetRecord(
        id="t1", task="q", eval_criteria="Must be correct"
    )
    save_eval_dataset([record], root=tmp_path, dataset_id="t/s")
    dataset = load_dataset("t/s", LocalDatasetConfig(local_root=tmp_path))
    tasks = tasks_from_eval_dataset(dataset)

    assert tasks[0].metadata["_criteria"] == "Must be correct"


def test_tasks_from_eval_dataset_works_with_json_engine(tmp_path: Path) -> None:
    save_eval_dataset([_make_sql_record()], root=tmp_path, dataset_id="t/j", engine="json")
    dataset = load_dataset("t/j", LocalDatasetConfig(local_root=tmp_path))
    tasks = tasks_from_eval_dataset(dataset)
    assert len(tasks) == 1
    assert tasks[0].name == "asmd_1"
    assert tasks[0].input == "What is the status of CW1?"


def test_tasks_from_eval_dataset_raises_on_missing_required_columns(tmp_path: Path) -> None:
    from agentbyte.dataset.base import DatasetConfig
    from agentbyte.dataset.json import JSONDataset

    config = DatasetConfig(engine="json", table_name="eval", in_memory=True)
    dataset = JSONDataset(
        dataset_id="t/bad",
        config=config,
        data_text=json.dumps({"records": [{"col_a": "v1", "col_b": "v2"}]}),
        source_uri="memory",
    )
    with pytest.raises(ValueError, match="missing required eval columns"):
        tasks_from_eval_dataset(dataset)


# ---------------------------------------------------------------------------


