from __future__ import annotations

import json
import logging
from pathlib import Path
from unittest.mock import AsyncMock

import pytest
from pydantic import BaseModel

from aiops.infra.agent.middlewares.data_warehouse.data_warehouse import (
    DataWarehouse,
    DataWarehouseDataFile,
    DataWarehouseDatasetDB,
    DataWarehouseDatasetMetadata,
)
from aiops.infra.agent.middlewares.data_warehouse.middleware import (
    DataWarehouseMiddleware,
)
from tests.infra.agent.middlewares.helpers import DummyModelRequest


class SampleRow(BaseModel):
    id: int
    name: str


class ContextWithWarehouseRoot:
    def __init__(self, root: str | None) -> None:
        self.data_warehouse_root = root


@pytest.fixture
def test_logger() -> logging.Logger:
    return logging.getLogger("tests.datawarehouse")


def test_datawarehouse_datafile_text_and_binary(
    tmp_path: Path, test_logger: logging.Logger
) -> None:
    file_path = tmp_path / "files" / "record.txt"
    datafile = DataWarehouseDataFile(file_path, test_logger)

    datafile.write("hello world")
    assert datafile.path == file_path
    assert datafile.read() == "hello world"

    payload = b"raw-bytes"
    datafile.write_bytes(file_path, payload)
    assert datafile.read_bytes() == payload

    datafile.delete()
    assert datafile.path.exists() is False


@pytest.mark.asyncio
async def test_datawarehouse_datafile_async_text_and_binary(
    tmp_path: Path, test_logger: logging.Logger
) -> None:
    file_path = tmp_path / "files" / "async_record.txt"
    datafile = DataWarehouseDataFile(file_path, test_logger)

    await datafile.awrite("hello async")
    assert await datafile.aread() == "hello async"

    payload = b"async-bytes"
    await datafile.awrite_bytes(file_path, payload)
    assert await datafile.aread_bytes() == payload


def test_datawarehouse_metadata_and_listing(tmp_path: Path) -> None:
    warehouse = DataWarehouse(str(tmp_path))

    metadata = warehouse.get_metadata()
    assert metadata.root == str(warehouse.root)

    rel_path = "data/session.log"
    datafile = warehouse.get_datafile(rel_path)
    datafile.write("entry")

    files = warehouse.list_datafiles(pattern="**/*.log")
    assert len(files) == 1
    assert files[0].path == warehouse.datafiles_dir / rel_path


def test_datawarehouse_get_db(tmp_path: Path) -> None:
    warehouse = DataWarehouse(str(tmp_path))
    db = warehouse.get_db()

    assert isinstance(db, DataWarehouseDatasetDB)
    assert (warehouse.database_dir / DataWarehouseDatasetDB.DATABASE_FILE).exists()

    assert db.execute("SELECT 42").fetchone()[0] == 42
    db.close()


def test_datawarehouse_db_write_dataset_and_metadata(
    tmp_path: Path, test_logger: logging.Logger
) -> None:
    db_dir = tmp_path / "db"
    db = DataWarehouseDatasetDB(db_dir, test_logger)

    assert db.get_dataset_metadata("events") is None

    first_batch = [
        SampleRow(id=1, name="alpha"),
        SampleRow(id=2, name="beta"),
    ]
    second_batch = [SampleRow(id=3, name="gamma")]

    db.write_dataset("events", first_batch, table_description="event rows")
    db.write_dataset("events", second_batch)

    total = db.execute("SELECT COUNT(*) FROM events").fetchone()[0]
    assert total == len(first_batch) + len(second_batch)

    metadata = db.get_dataset_metadata("events")
    assert isinstance(metadata, DataWarehouseDatasetMetadata)
    assert metadata.description == "event rows"
    schema = json.loads(metadata.field_description or "{}")
    assert "properties" in schema
    assert metadata.table_schema is not None
    assert any(column.column_name == "id" for column in metadata.table_schema)

    metadatas = db.list_dataset_metadatas()
    assert any(item.table_name == "events" for item in metadatas)
    assert all(item.table_schema is not None for item in metadatas)

    db.close()


@pytest.mark.asyncio
async def test_datawarehouse_db_async_interfaces(
    tmp_path: Path, test_logger: logging.Logger
) -> None:
    db_dir = tmp_path / "db_async"
    db = DataWarehouseDatasetDB(db_dir, test_logger)

    try:
        await db.awrite_dataset(
            "events",
            [SampleRow(id=10, name="async-alpha")],
            table_description="async rows",
        )

        metadata = await db.aget_dataset_metadata("events")
        assert metadata is not None
        assert metadata.description == "async rows"
        assert metadata.table_schema is not None

        metadatas = await db.list_dataset_metadatas_async()
        assert any(item.table_name == "events" for item in metadatas)
        assert all(item.table_schema is not None for item in metadatas)

        count = db.execute("SELECT COUNT(*) FROM events").fetchone()[0]
        assert count == 1
    finally:
        db.close()


def test_datawarehouse_middleware_build_from_context(tmp_path: Path) -> None:
    context = ContextWithWarehouseRoot(str(tmp_path))
    warehouse = DataWarehouseMiddleware.build_from_context(context)

    assert isinstance(warehouse, DataWarehouse)
    assert str(warehouse.root) == str(tmp_path.resolve())

    context_without_root = ContextWithWarehouseRoot(None)
    assert DataWarehouseMiddleware.build_from_context(context_without_root) is None
    assert DataWarehouseMiddleware.build_from_context(object()) is None


def test_datawarehouse_middleware_available_dataset_prompt_empty(
    tmp_path: Path,
) -> None:
    middleware = DataWarehouseMiddleware()
    db = DataWarehouse(str(tmp_path)).get_db()

    try:
        prompt = middleware._build_available_dataset_prompt(db)
    finally:
        db.close()

    assert prompt == "No datasets are available in the Data Warehouse."


@pytest.mark.asyncio
async def test_datawarehouse_middleware_appends_prompt(
    tmp_path: Path,
    dummy_model_request,
) -> None:
    middleware = DataWarehouseMiddleware()
    warehouse = DataWarehouse(str(tmp_path))
    db = warehouse.get_db()

    try:
        db.write_dataset(
            "events",
            [SampleRow(id=1, name="alpha")],
            table_description="Event rows",
        )
    finally:
        db.close()

    handler = AsyncMock(return_value="handled")
    request = dummy_model_request(
        system_prompt="base",
        runtime_context=ContextWithWarehouseRoot(str(tmp_path)),
    )

    result = await middleware.awrap_model_call(request, handler)

    assert result == "handled"
    handler.assert_awaited_once()
    called_request = handler.await_args[0][0]
    assert isinstance(called_request, DummyModelRequest)
    assert "Data Warehouse Tools Usage" in called_request.system_prompt
    assert '"dataset_name":"events"' in called_request.system_prompt


@pytest.mark.asyncio
async def test_datawarehouse_middleware_skips_without_root(dummy_model_request) -> None:
    middleware = DataWarehouseMiddleware()
    handler = AsyncMock(return_value="handled")
    request = dummy_model_request(
        system_prompt="base",
        runtime_context=ContextWithWarehouseRoot(None),
    )

    result = await middleware.awrap_model_call(request, handler)

    assert result == "handled"
    handler.assert_awaited_once()
    called_request = handler.await_args[0][0]
    assert called_request is request