from __future__ import annotations

from collections.abc import Generator
from datetime import UTC, datetime, timedelta

from auth_helpers import install_legacy_header_auth_override
from fastapi.testclient import TestClient
from sqlalchemy import create_engine
from sqlalchemy.orm import Session, sessionmaker
from sqlalchemy.pool import StaticPool

from app.api.deps import get_db
from app.core.agent_enums import AgentName, AgentRunSource, AgentRunStatus, AgentToolType
from app.db.base import Base
from app.main import create_app
from app.models.agent_conversation import AgentConversation, AgentConversationMessage
from app.services.agent_runs import AgentRunService
from app.services.agent_traces import AgentTraceService


def build_session() -> Session:
    engine = create_engine(
        "sqlite+pysqlite:///:memory:",
        connect_args={"check_same_thread": False},
        poolclass=StaticPool,
    )
    Base.metadata.create_all(bind=engine)
    session_factory = sessionmaker(bind=engine, autoflush=False, autocommit=False)
    return session_factory()


def build_client() -> tuple[TestClient, sessionmaker[Session]]:
    engine = create_engine(
        "sqlite+pysqlite:///:memory:",
        connect_args={"check_same_thread": False},
        poolclass=StaticPool,
    )
    Base.metadata.create_all(bind=engine)
    session_factory = sessionmaker(bind=engine, autoflush=False, autocommit=False)
    app = create_app()
    install_legacy_header_auth_override(app)

    def override_db() -> Generator[Session, None, None]:
        db = session_factory()
        try:
            yield db
        finally:
            db.close()

    app.dependency_overrides[get_db] = override_db
    return TestClient(app), session_factory


def test_agent_trace_service_records_events_and_reads_detail() -> None:
    with build_session() as db:
        run_service = AgentRunService(db)
        trace_service = AgentTraceService(db)
        started_at = datetime.now(UTC) - timedelta(seconds=2)
        run = run_service.create_run(
            agent=AgentName.ORCHESTRATOR.value,
            source=AgentRunSource.USER_MESSAGE.value,
            status=AgentRunStatus.SUCCEEDED.value,
            route_json={"conversation_id": "conv-trace-1"},
            result_summary="expense answer ready",
            started_at=started_at,
            finished_at=started_at + timedelta(seconds=1),
        )
        db.add(
            AgentConversation(
                conversation_id="conv-trace-1",
                user_id="u-1",
                source=AgentRunSource.USER_MESSAGE.value,
            )
        )
        db.add(
            AgentConversationMessage(
                conversation_id="conv-trace-1",
                run_id=run.run_id,
                role="user",
                content="帮我看报销风险",
                message_json={"source": "test"},
            )
        )
        db.commit()

        first = trace_service.record_event(
            run_id=run.run_id,
            conversation_id="conv-trace-1",
            stage="orchestrator",
            event_name="request_received",
            title="接收请求",
            summary="用户消息进入编排",
            input_json={"message": "帮我看报销风险"},
            output_json={"run_id": run.run_id},
            started_at=started_at,
            finished_at=started_at + timedelta(milliseconds=20),
        )
        second = trace_service.record_event(
            run_id=run.run_id,
            conversation_id="conv-trace-1",
            stage="response",
            event_name="response_built",
            title="生成回复",
            status=AgentRunStatus.SUCCEEDED.value,
            output_json={"message": "已完成"},
            started_at=started_at + timedelta(milliseconds=500),
            finished_at=started_at + timedelta(milliseconds=650),
        )

        items = trace_service.list_traces(keyword=run.run_id, limit=10)
        detail = trace_service.get_trace(run.run_id)

        assert first.sequence == 1
        assert second.sequence == 2
        assert len(items) == 1
        assert items[0].event_count == 2
        assert detail is not None
        assert detail.fallback_generated is False
        assert [event.event_name for event in detail.events] == [
            "request_received",
            "response_built",
        ]
        assert detail.conversation_id == "conv-trace-1"
        assert detail.conversation_messages[0].content == "帮我看报销风险"


def test_agent_trace_service_builds_fallback_timeline_for_legacy_runs() -> None:
    with build_session() as db:
        run_service = AgentRunService(db)
        trace_service = AgentTraceService(db)
        run = run_service.create_run(
            agent=AgentName.HERMES.value,
            source=AgentRunSource.SCHEDULE.value,
            status=AgentRunStatus.FAILED.value,
            route_json={"conversation_id": "conv-trace-legacy", "stage": "tooling"},
            result_summary="sync failed",
            error_message="boom",
        )
        run_service.record_semantic_parse(
            run_id=run.run_id,
            user_id="u-1",
            raw_query="同步知识库",
            scenario="knowledge",
            intent="sync",
            confidence=0.92,
        )
        run_service.record_tool_call(
            run_id=run.run_id,
            tool_type=AgentToolType.LLM.value,
            tool_name="lightrag.index_documents",
            request_json={"document_ids": ["doc-1"]},
            response_json={"fallback": True},
            status=AgentRunStatus.FAILED.value,
            duration_ms=31,
            error_message="boom",
        )

        detail = trace_service.get_trace(run.run_id)
        conversation_detail = trace_service.get_conversation_trace("conv-trace-legacy")

        assert detail is not None
        assert detail.fallback_generated is True
        assert detail.conversation_id == "conv-trace-legacy"
        assert "semantic_parsed" in [event.event_name for event in detail.events]
        assert "tool_invoked" in [event.event_name for event in detail.events]
        assert detail.events[-1].event_name == "failed"
        assert detail.tool_calls[0].tool_name == "lightrag.index_documents"
        assert [item.run.run_id for item in conversation_detail.runs] == [run.run_id]


def test_agent_trace_endpoints_return_admin_trace_detail() -> None:
    client, session_factory = build_client()

    with session_factory() as db:
        run_service = AgentRunService(db)
        trace_service = AgentTraceService(db)
        run = run_service.create_run(
            agent=AgentName.ORCHESTRATOR.value,
            source=AgentRunSource.USER_MESSAGE.value,
            status=AgentRunStatus.SUCCEEDED.value,
            route_json={"conversation_id": "conv-api-trace"},
            result_summary="api trace ready",
        )
        trace_service.record_event(
            run_id=run.run_id,
            conversation_id="conv-api-trace",
            stage="response",
            event_name="response_built",
            title="生成回复",
            status=AgentRunStatus.SUCCEEDED.value,
            output_json={"message": "ok"},
        )

    headers = {
        "x-auth-username": "admin",
        "x-auth-name": "admin",
        "x-auth-is-admin": "true",
    }
    list_response = client.get("/api/v1/agent-traces", headers=headers)
    detail_response = client.get(f"/api/v1/agent-traces/{run.run_id}", headers=headers)

    assert list_response.status_code == 200
    assert any(item["run_id"] == run.run_id for item in list_response.json())
    assert detail_response.status_code == 200
    payload = detail_response.json()
    assert payload["run"]["run_id"] == run.run_id
    assert payload["events"][0]["event_name"] == "response_built"
