from __future__ import annotations

import io
import json
from pathlib import Path

import pytest

from agentbyte.dataset import JSONDataset, SQLiteDataset, load_dataset
from agentbyte.dataset.config import LocalDatasetConfig, S3DatasetConfig
from agentbyte.dataset.sources import S3DatasetSource

# Root of the bundled sample datasets shipped with the repo.
_SAMPLES_ROOT = Path(__file__).parents[2] / "data" / "samples"


def _write_dataset(
    dataset_dir: Path,
    *,
    engine: str = "sqlite",
    table_name: str = "contracts_asmd",
    schema: list[dict[str, object]] | None = None,
    data: object | None = None,
) -> None:
    dataset_dir.mkdir(parents=True, exist_ok=True)
    (dataset_dir / "config.json").write_text(
        json.dumps(
            {
                "engine": engine,
                "table_name": table_name,
                "in_memory": True,
                "description": "Test dataset",
            }
        ),
        encoding="utf-8",
    )

    if schema is not None:
        (dataset_dir / "schema.json").write_text(json.dumps(schema), encoding="utf-8")

    if data is None:
        data = {"records": [{"id": "1", "name": "Acme"}]}
    (dataset_dir / "data.json").write_text(json.dumps(data), encoding="utf-8")


class FakeS3Client:
    def __init__(self, payloads: dict[str, str]) -> None:
        self.payloads = payloads

    def get_object(self, *, Bucket: str, Key: str) -> dict[str, io.BytesIO]:
        del Bucket
        if Key not in self.payloads:
            raise _FakeClientError("NoSuchKey")
        return {"Body": io.BytesIO(self.payloads[Key].encode("utf-8"))}


class _FakeClientError(Exception):
    def __init__(self, code: str) -> None:
        super().__init__(code)
        self.response = {"Error": {"Code": code}}


def test_load_dataset_local_source() -> None:
    config = LocalDatasetConfig(local_root=_SAMPLES_ROOT)
    ds = load_dataset("contracts10/asmd", config)

    assert isinstance(ds, SQLiteDataset)
    assert ds.table_name == "contracts_asmd"

    conn = ds.connect()
    rows = conn.execute("SELECT * FROM contracts_asmd").fetchall()
    assert len(rows) == 10


def test_load_dataset_json_from_local_source(tmp_path: Path) -> None:
    dataset_root = tmp_path / "samples"
    _write_dataset(
        dataset_root / "documents" / "wiki",
        engine="json",
        table_name="documents_wiki",
        schema=None,
        data={"records": [{"title": "Example", "url": "https://example.com"}]},
    )

    config = LocalDatasetConfig(local_root=dataset_root)
    ds = load_dataset("documents/wiki", config)

    assert isinstance(ds, JSONDataset)
    data = ds.connect()
    assert data["records"][0]["title"] == "Example"


def test_load_dataset_rejects_invalid_dataset_id() -> None:
    config = LocalDatasetConfig(local_root=_SAMPLES_ROOT)
    with pytest.raises(ValueError, match="Dataset ID must be 'category/name'"):
        load_dataset("contracts10", config)


def test_load_dataset_from_s3_without_prefix(monkeypatch: pytest.MonkeyPatch) -> None:
    payloads = {
        "contracts10/asmd/config.json": json.dumps(
            {
                "engine": "sqlite",
                "table_name": "contracts_asmd",
                "in_memory": True,
                "description": "ASMD contracts",
            }
        ),
        "contracts10/asmd/schema.json": json.dumps(
            [
                {"name": "id", "type": "TEXT", "primary_key": True},
                {"name": "name", "type": "TEXT"},
            ]
        ),
        "contracts10/asmd/data.json": json.dumps(
            {"records": [{"id": "1", "name": "Acme"}]}
        ),
    }
    fake_client = FakeS3Client(payloads)
    monkeypatch.setattr(S3DatasetSource, "_get_client", lambda self: fake_client)

    config = S3DatasetConfig(s3_bucket="datapsycho-datasets")
    ds = load_dataset("contracts10/asmd", config)

    assert isinstance(ds, SQLiteDataset)
    rows = ds.connect().execute("SELECT * FROM contracts_asmd").fetchall()
    assert len(rows) == 1
    assert rows[0]["name"] == "Acme"


def test_load_dataset_from_s3_with_prefix(monkeypatch: pytest.MonkeyPatch) -> None:
    payloads = {
        "datasets/documents/wiki/config.json": json.dumps(
            {
                "engine": "json",
                "table_name": "documents_wiki",
                "in_memory": True,
                "description": "Wiki docs",
            }
        ),
        "datasets/documents/wiki/data.json": json.dumps(
            {"records": [{"title": "Doc", "url": "https://example.com"}]}
        ),
    }
    fake_client = FakeS3Client(payloads)
    monkeypatch.setattr(S3DatasetSource, "_get_client", lambda self: fake_client)

    config = S3DatasetConfig(s3_bucket="datapsycho-datasets", s3_prefix="datasets")
    ds = load_dataset("documents/wiki", config)

    assert isinstance(ds, JSONDataset)
    data = ds.connect()
    assert data["records"][0]["title"] == "Doc"


def test_load_dataset_from_s3_raises_clear_error_for_missing_object(
    monkeypatch: pytest.MonkeyPatch,
) -> None:
    fake_client = FakeS3Client({})
    monkeypatch.setattr(S3DatasetSource, "_get_client", lambda self: fake_client)

    config = S3DatasetConfig(s3_bucket="datapsycho-datasets")
    with pytest.raises(
        FileNotFoundError,
        match=r"s3://datapsycho-datasets/contracts10/asmd/config.json",
    ):
        load_dataset("contracts10/asmd", config)


def test_load_dataset_from_s3_requires_bucket_setting(
    monkeypatch: pytest.MonkeyPatch,
) -> None:
    monkeypatch.delenv("AGENTBYTE_DATASET_S3_BUCKET", raising=False)

    with pytest.raises(Exception, match="s3_bucket|S3_BUCKET|validation"):
        S3DatasetConfig()


def test_sqlite_dataset_from_s3_requires_schema(
    monkeypatch: pytest.MonkeyPatch,
) -> None:
    payloads = {
        "contracts10/asmd/config.json": json.dumps(
            {
                "engine": "sqlite",
                "table_name": "contracts_asmd",
                "in_memory": True,
                "description": "ASMD contracts",
            }
        ),
        "contracts10/asmd/data.json": json.dumps({"records": [{"id": "1"}]}),
    }
    fake_client = FakeS3Client(payloads)
    monkeypatch.setattr(S3DatasetSource, "_get_client", lambda self: fake_client)

    config = S3DatasetConfig(s3_bucket="datapsycho-datasets")
    ds = load_dataset("contracts10/asmd", config)

    with pytest.raises(FileNotFoundError, match="schema.json"):
        ds.connect()
