"""Standard eval dataset schema, persistence, and task conversion.

Defines a canonical column schema for AgentByte eval datasets so that
evaluation pipelines can be built without hand-writing column mappings.

Two typed context helpers cover the two primary eval types in use:

- ``RAGEvalContext`` — document retrieval / RAG evaluation
- ``SQLEvalContext`` — text-to-SQL evaluation

Any eval type not covered by the helpers can use a plain ``dict`` for the
``context`` field of ``EvalDatasetRecord``.

Typical workflow::

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

    records = [
        EvalDatasetRecord(
            id="task_1",
            task="What is the contract status of CW386232?",
            expected_response="Published.",
            eval_criteria=["Correct status", "No hallucinations"],
            context=SQLEvalContext(
                reference_sql="SELECT DISTINCT contract_status FROM ...",
                reference_table="proc_contracting_asmd_prod",
                context_id="CW386232",
            ).model_dump(exclude_none=True),
        ),
    ]
    save_eval_dataset(records, root=Path("/data"), dataset_id="asmd/sql_eval")

    dataset = load_dataset("asmd/sql_eval", LocalDatasetConfig(local_root=Path("/data")))
    tasks = tasks_from_eval_dataset(dataset)
    # tasks[0].metadata["reference_sql"]  → the expected SQL query
    # tasks[0].metadata["_criteria"]      → ["Correct status", "No hallucinations"]
"""

from __future__ import annotations

import json
from pathlib import Path
from typing import Any

from pydantic import BaseModel, ConfigDict, Field

from agentbyte.dataset.base import BaseDataset
from agentbyte.eval.types import EvalTask


# ---------------------------------------------------------------------------
# SQLite schema — six fixed columns
# ---------------------------------------------------------------------------

STANDARD_EVAL_COLUMNS: list[dict[str, Any]] = [
    {"name": "id",                "type": "TEXT", "primary_key": True},
    {"name": "scenario",          "type": "TEXT"},
    {"name": "task",              "type": "TEXT"},
    {"name": "expected_response", "type": "TEXT"},
    {"name": "eval_criteria",     "type": "TEXT"},   # JSON list or plain string
    {"name": "context",           "type": "TEXT"},   # JSON dict — domain fields
]


# ---------------------------------------------------------------------------
# Standard eval record
# ---------------------------------------------------------------------------


class EvalDatasetRecord(BaseModel):
    """One row in a standardised AgentByte eval dataset.

    The five universal fields (``id``, ``task``, ``expected_response``,
    ``eval_criteria``, ``scenario``) cover every eval type.
    Domain-specific data — doc IDs, SQL queries, file names, etc. — goes into
    ``context`` as a plain ``dict``.  Use ``RAGEvalContext`` or
    ``SQLEvalContext`` to build a typed context, or pass a plain dict directly.

    All ``context`` fields are unpacked into ``EvalTask.metadata`` when the
    dataset is loaded with ``tasks_from_eval_dataset()``.

    Attributes:
        id: Unique task identifier.  Becomes ``EvalTask.name``.
        task: User input or prompt.  Becomes ``EvalTask.input``.
        scenario: Optional test scenario label (e.g. "Single document").
        expected_response: Reference answer.  Becomes
            ``EvalTask.expected_output``.
        eval_criteria: Per-task judging rules — a plain string or a list of
            strings.  Stored under ``EvalTask.metadata["_criteria"]`` on load.
        context: Domain-specific evaluation data.  All keys are merged into
            ``EvalTask.metadata`` on load.
    """

    model_config = ConfigDict(frozen=True)

    id: str
    task: str
    scenario: str = ""
    expected_response: str | None = None
    eval_criteria: str | list[str] | None = None
    context: dict[str, Any] = Field(default_factory=dict)


# ---------------------------------------------------------------------------
# Typed context helpers
# ---------------------------------------------------------------------------


class RAGEvalContext(BaseModel):
    """Context for RAG / document retrieval evaluation.

    Pass as ``context=RAGEvalContext(...).model_dump(exclude_none=True)`` when
    constructing an ``EvalDatasetRecord``.

    On load, all fields become available in ``EvalTask.metadata``:
    ``task.metadata["doc_ids"]``, ``task.metadata["file_names"]``, etc.

    Attributes:
        doc_ids: Relevant document IDs (ground truth for retrieval P/R).
        file_names: Source document file names.
        reference_section: Location of the answer in the source document.
        context_id: Business context identifier, e.g. a contract work ID.
    """

    model_config = ConfigDict(frozen=True)

    doc_ids: list[str] = Field(default_factory=list)
    file_names: list[str] = Field(default_factory=list)
    reference_section: str | None = None
    context_id: str | None = None


class SQLEvalContext(BaseModel):
    """Context for text-to-SQL evaluation.

    Pass as ``context=SQLEvalContext(...).model_dump(exclude_none=True)`` when
    constructing an ``EvalDatasetRecord``.

    On load, all fields become available in ``EvalTask.metadata``:
    ``task.metadata["reference_sql"]``, ``task.metadata["reference_table"]``,
    etc.

    Attributes:
        reference_sql: Expected / reference SQL query.
        reference_table: Database table or view name.
        db_schema: Optional schema hint for the target database.
        context_id: Business context identifier, e.g. a contract work ID.
    """

    model_config = ConfigDict(frozen=True)

    reference_sql: str | None = None
    reference_table: str | None = None
    db_schema: str | None = None
    context_id: str | None = None


# ---------------------------------------------------------------------------
# Persistence
# ---------------------------------------------------------------------------


def save_eval_dataset(
    records: list[EvalDatasetRecord],
    root: Path,
    dataset_id: str,
    *,
    engine: str = "sqlite",
    table_name: str = "eval",
    description: str = "",
) -> None:
    """Persist eval records to a local dataset directory.

    Writes files in the AgentByte dataset layout so the result can be loaded
    with ``load_dataset()`` and converted to tasks via
    ``tasks_from_eval_dataset()``:

    - ``config.json`` — engine, table name, and description
    - ``schema.json`` — the six standard column definitions (SQLite only)
    - ``data.json``   — records with list/dict fields JSON-encoded

    Args:
        records: Eval records to persist.
        root: Local root directory (same value as
            ``LocalDatasetConfig.local_root``).
        dataset_id: Dataset identifier in ``"category/name"`` format, e.g.
            ``"asmd/sql_eval"``.
        engine: Dataset engine to write.  Defaults to ``"sqlite"``.  Use
            ``"json"`` for a schema-free JSON dataset.
        table_name: Table/collection name stored in ``config.json``.
            Defaults to ``"eval"``.
        description: Human-readable description stored in ``config.json``.

    Raises:
        ValueError: If ``dataset_id`` is not in ``"category/name"`` format.
    """
    parts = [p for p in dataset_id.split("/") if p]
    if len(parts) != 2:  # noqa: PLR2004
        raise ValueError(
            f"dataset_id must be 'category/name', got: {dataset_id!r}"
        )

    dataset_dir = root / dataset_id
    dataset_dir.mkdir(parents=True, exist_ok=True)

    # config.json
    config_payload = {
        "engine": engine,
        "table_name": table_name,
        "in_memory": True,
        "description": description,
    }
    (dataset_dir / "config.json").write_text(
        json.dumps(config_payload, indent=2), encoding="utf-8"
    )

    # schema.json — only written for SQLite (schema-driven engines)
    if engine == "sqlite":
        (dataset_dir / "schema.json").write_text(
            json.dumps(STANDARD_EVAL_COLUMNS, indent=2), encoding="utf-8"
        )

    # data.json — JSON-encode list and dict fields for TEXT storage
    record_rows: list[dict[str, Any]] = []
    for rec in records:
        criteria_value: str | None
        if isinstance(rec.eval_criteria, list):
            criteria_value = json.dumps(rec.eval_criteria)
        else:
            criteria_value = rec.eval_criteria  # str or None

        record_rows.append(
            {
                "id": rec.id,
                "scenario": rec.scenario,
                "task": rec.task,
                "expected_response": rec.expected_response,
                "eval_criteria": criteria_value,
                "context": json.dumps(rec.context),
            }
        )

    (dataset_dir / "data.json").write_text(
        json.dumps({"records": record_rows}, indent=2), encoding="utf-8"
    )


# ---------------------------------------------------------------------------
# Task conversion
# ---------------------------------------------------------------------------


def _parse_criteria(raw: str | None) -> str | list[str] | None:
    """Parse eval_criteria from SQLite TEXT storage.

    JSON-encoded lists are parsed back to ``list[str]``.
    Plain strings are returned as-is.
    ``None`` is returned when the column is NULL.
    """
    if raw is None:
        return None
    try:
        parsed = json.loads(raw)
        if isinstance(parsed, list):
            return parsed
    except (json.JSONDecodeError, ValueError):
        pass
    return raw


def _parse_context(raw: str | None) -> dict[str, Any]:
    """Parse the context JSON column back to a dict.

    Returns an empty dict when the column is NULL or not valid JSON.
    """
    if not raw:
        return {}
    try:
        result = json.loads(raw)
        return result if isinstance(result, dict) else {}
    except (json.JSONDecodeError, ValueError):
        return {}


def tasks_from_eval_dataset(dataset: BaseDataset) -> list[EvalTask]:
    """Convert a standard-schema eval dataset to a list of ``EvalTask`` objects.

    Works with any registered engine. Rows are fetched via the engine-agnostic
    ``fetch_all()`` method — no engine-specific query code is required.

    Column → ``EvalTask`` mapping:

    ==================  =============================================
    ``id``              ``task.name``
    ``task``            ``task.input``
    ``expected_response`` ``task.expected_output``
    ``eval_criteria``   ``task.metadata["_criteria"]`` (list or str)
    ``scenario``        ``task.metadata["scenario"]``
    ``context``         unpacked into ``task.metadata`` (all keys)
    ==================  =============================================

    Args:
        dataset: A ``BaseDataset`` loaded from a standard eval dataset
            directory.

    Returns:
        List of ``EvalTask`` objects ready for ``EvalRunner.evaluate()``.

    Raises:
        ValueError: If the dataset rows are missing required eval columns
            (``id`` and ``task``).
    """
    rows = dataset.fetch_all()

    # Validate that the dataset follows the standard eval schema
    required = {"id", "task"}
    if rows and not required.issubset(rows[0].keys()):
        missing = required - rows[0].keys()
        raise ValueError(
            f"Dataset rows are missing required eval columns: {sorted(missing)}. "
            f"Got columns: {sorted(rows[0].keys())}. "
            f"Ensure the dataset was saved with save_eval_dataset()."
        )

    tasks: list[EvalTask] = []
    for row_dict in rows:
        context = _parse_context(row_dict.get("context"))
        criteria = _parse_criteria(row_dict.get("eval_criteria"))

        metadata: dict[str, Any] = {}
        scenario = row_dict.get("scenario")
        if scenario:
            metadata["scenario"] = scenario
        if criteria is not None:
            metadata["_criteria"] = criteria
        metadata.update(context)

        tasks.append(
            EvalTask(
                name=row_dict["id"],
                input=row_dict["task"],
                expected_output=row_dict.get("expected_response"),
                metadata=metadata,
            )
        )

    return tasks


__all__ = [
    "STANDARD_EVAL_COLUMNS",
    "EvalDatasetRecord",
    "RAGEvalContext",
    "SQLEvalContext",
    "save_eval_dataset",
    "tasks_from_eval_dataset",
]
