"""Stream truncation event persistence + model-layer classification tests."""

from __future__ import annotations

import asyncio
from datetime import UTC, datetime, timedelta

import pytest
import pytest_asyncio
from sqlalchemy.ext.asyncio import async_sessionmaker, create_async_engine

import deerflow.persistence.stream_truncation as stream_truncation_module
from deerflow.models.patched_openai import _ToolCallStreamDiagnostics
from deerflow.persistence.base import Base
from deerflow.persistence.stream_truncation import (
    SqlStreamTruncationStore,
    StreamTruncationEvent,
    StreamTruncationQuery,
)


@pytest_asyncio.fixture
async def session_factory(tmp_path):
    engine = create_async_engine(f"sqlite+aiosqlite:///{tmp_path / 'stream-truncation.db'}")
    async with engine.begin() as conn:
        await conn.run_sync(Base.metadata.create_all)
    factory = async_sessionmaker(engine, expire_on_commit=False)
    yield factory
    await engine.dispose()


def _event(**overrides) -> StreamTruncationEvent:
    payload = dict(
        id="evt-1",
        created_at=datetime.now(UTC),
        model_name="relay-model",
        run_id="run-1",
        event_type="silent_no_terminal",
        finish_reason="none",
        status="success",
        chunk_count=12,
        duration_ms=3_000,
        first_chunk_ms=120,
    )
    payload.update(overrides)
    return StreamTruncationEvent(**payload)


@pytest.mark.asyncio
async def test_record_list_count_and_breakdown(session_factory) -> None:
    store = SqlStreamTruncationStore(session_factory)
    await store.record(_event())
    await store.record(
        _event(
            id="evt-2",
            model_name="relay-model",
            event_type="length_truncated",
            finish_reason="length",
        )
    )
    await store.record(_event(id="evt-3", model_name="other-model"))

    items = await store.list(StreamTruncationQuery(limit=10))
    assert len(items) == 3
    total = await store.count(StreamTruncationQuery(model_name="relay-model"))
    assert total == 2

    filtered = await store.list(
        StreamTruncationQuery(model_name="relay-model", event_type="length_truncated")
    )
    assert [item["id"] for item in filtered] == ["evt-2"]

    breakdown = await store.model_breakdown(StreamTruncationQuery())
    by_model = {item["model_name"]: item for item in breakdown}
    assert by_model["relay-model"]["total"] == 2
    assert by_model["relay-model"]["length_truncated"] == 1
    assert by_model["relay-model"]["silent_no_terminal"] == 1
    assert by_model["other-model"]["total"] == 1
    assert by_model["relay-model"]["last_at"]


@pytest.mark.asyncio
async def test_retention_cap_deletes_rows_older_than_one_month(session_factory) -> None:
    store = SqlStreamTruncationStore(session_factory)
    old = _event(
        id="evt-old",
        created_at=datetime.now(UTC) - timedelta(days=40),
    )
    await store.record(old)
    # The sweep after every insert must already have removed the 40-day-old row.
    assert await store.count(StreamTruncationQuery()) == 0

    fresh = _event(id="evt-fresh", created_at=datetime.now(UTC) - timedelta(days=3))
    await store.record(fresh)
    assert await store.count(StreamTruncationQuery()) == 1

    purged = await store.purge_before(datetime.now(UTC) - timedelta(days=30))
    assert purged == 0
    assert await store.count(StreamTruncationQuery()) == 1


def _diagnostics(**overrides) -> _ToolCallStreamDiagnostics:
    diagnostics = _ToolCallStreamDiagnostics(
        model_name="relay-model",
        run_id="run-abc",
        strict_stream_integrity=False,
    )
    for key, value in overrides.items():
        setattr(diagnostics, key, value)
    return diagnostics


def _capture_recorded(monkeypatch: pytest.MonkeyPatch):
    captured: list[dict] = []

    def fake_record(**kwargs):
        captured.append(kwargs)

    monkeypatch.setattr(
        stream_truncation_module,
        "record_stream_truncation_event",
        fake_record,
    )
    return captured


def _flush(diagnostics: _ToolCallStreamDiagnostics, *, status: str = "success", error=None, error_type: str = "none") -> None:
    diagnostics._record_truncation_event(
        status=status,
        error=error,
        error_type=error_type,
        total_ms=1_500,
        first_chunk_ms=100,
    )


def test_classification_silent_no_terminal(monkeypatch: pytest.MonkeyPatch) -> None:
    captured = _capture_recorded(monkeypatch)
    diagnostics = _diagnostics(chunk_count=8)
    _flush(diagnostics)
    assert len(captured) == 1
    assert captured[0]["event_type"] == "silent_no_terminal"
    assert captured[0]["finish_reason"] == "none"
    assert captured[0]["model_name"] == "relay-model"
    assert captured[0]["strict_mode"] is False
    assert captured[0]["run_id"] == "run-abc"


def test_classification_skips_clean_and_empty_streams(monkeypatch: pytest.MonkeyPatch) -> None:
    captured = _capture_recorded(monkeypatch)
    # Normal stop with terminal frame -> nothing to record.
    _flush(_diagnostics(chunk_count=8, finish_reason="stop"))
    # No terminal frame but zero chunks -> empty response, not a half answer.
    _flush(_diagnostics(chunk_count=0))
    # Cancelled by the user -> not a truncation.
    _flush(_diagnostics(chunk_count=8), status="cancelled", error=asyncio.CancelledError())
    assert captured == []


def test_classification_intentional_bridge_stop_is_skipped(monkeypatch: pytest.MonkeyPatch) -> None:
    captured = _capture_recorded(monkeypatch)
    _flush(_diagnostics(chunk_count=8, intentional_early_finish=True))
    assert captured == []


def test_classification_strict_incomplete_and_midstream_errors(monkeypatch: pytest.MonkeyPatch) -> None:
    captured = _capture_recorded(monkeypatch)
    _flush(
        _diagnostics(chunk_count=8, strict_stream_integrity=True),
        status="error",
        error=RuntimeError("incomplete"),
        error_type="IncompleteChatStreamError",
    )
    _flush(
        _diagnostics(chunk_count=30),
        status="error",
        error=TimeoutError(),
        error_type="ReadTimeout",
    )
    # Errors before the first chunk are connection failures, not truncations.
    _flush(
        _diagnostics(chunk_count=0),
        status="error",
        error=TimeoutError(),
        error_type="ReadTimeout",
    )
    assert [item["event_type"] for item in captured] == [
        "incomplete_stream_error",
        "midstream_error",
    ]
    assert captured[0]["strict_mode"] is True


def test_classification_length_and_content_filter(monkeypatch: pytest.MonkeyPatch) -> None:
    captured = _capture_recorded(monkeypatch)
    _flush(_diagnostics(chunk_count=99, finish_reason="length"))
    _flush(_diagnostics(chunk_count=99, finish_reason="content_filter"))
    assert [item["event_type"] for item in captured] == [
        "length_truncated",
        "content_filter",
    ]