from datetime import datetime

import pytest
from sqlalchemy import insert

from phoenix.config import DEFAULT_PROJECT_NAME
from phoenix.db import models
from phoenix.server.types import DbSessionFactory


@pytest.fixture
async def default_project(db: DbSessionFactory) -> None:
    async with db() as session:
        project_row_id = await session.scalar(
            insert(models.Project).values(name=DEFAULT_PROJECT_NAME).returning(models.Project.id)
        )
        trace_rowid = await session.scalar(
            insert(models.Trace)
            .values(
                trace_id="0123",
                project_rowid=project_row_id,
                start_time=datetime.fromisoformat("2021-01-01T00:00:00.000+00:00"),
                end_time=datetime.fromisoformat("2021-01-01T00:01:00.000+00:00"),
            )
            .returning(models.Trace.id)
        )
        await session.execute(
            insert(models.Span)
            .values(
                trace_rowid=trace_rowid,
                span_id="2345",
                parent_id=None,
                name="root span",
                span_kind="UNKNOWN",
                start_time=datetime.fromisoformat("2021-01-01T00:00:00.000+00:00"),
                end_time=datetime.fromisoformat("2021-01-01T00:00:30.000+00:00"),
                attributes={
                    "input": {"value": "210"},
                    "output": {"value": "321"},
                },
                events=[],
                status_code="OK",
                status_message="okay",
                cumulative_error_count=0,
                cumulative_llm_token_count_prompt=0,
                cumulative_llm_token_count_completion=0,
            )
            .returning(models.Span.id)
        )
        await session.execute(
            insert(models.Span)
            .values(
                trace_rowid=trace_rowid,
                span_id="4567",
                parent_id="2345",
                name="retriever span",
                span_kind="RETRIEVER",
                start_time=datetime.fromisoformat("2021-01-01T00:00:05.000+00:00"),
                end_time=datetime.fromisoformat("2021-01-01T00:00:20.000+00:00"),
                attributes={
                    "input": {"value": "xyz"},
                    "retrieval": {
                        "documents": [{"document": {"content": "A", "score": 1}}],
                    },
                },
                events=[],
                status_code="OK",
                status_message="okay",
                cumulative_error_count=0,
                cumulative_llm_token_count_prompt=0,
                cumulative_llm_token_count_completion=0,
            )
            .returning(models.Span.id)
        )

        await session.execute(
            insert(models.Span)
            .values(
                trace_rowid=trace_rowid,
                span_id="5678",
                parent_id="2345",
                name="retriever span",
                span_kind="RETRIEVER",
                start_time=datetime.fromisoformat("2021-01-01T00:00:05.000+00:00"),
                end_time=datetime.fromisoformat("2021-01-01T00:00:20.000+00:00"),
                attributes={
                    "input": {"value": "xyz"},
                    "retrieval": {
                        "documents": [{}, {"document": {"content": "B", "score": 2}}],
                    },
                },
                events=[],
                status_code="OK",
                status_message="okay",
                cumulative_error_count=0,
                cumulative_llm_token_count_prompt=0,
                cumulative_llm_token_count_completion=0,
            )
            .returning(models.Span.id)
        )
        await session.execute(
            insert(models.Span)
            .values(
                trace_rowid=trace_rowid,
                span_id="6789",
                parent_id="2345",
                name="retriever span",
                span_kind="RETRIEVER",
                start_time=datetime.fromisoformat("2021-01-01T00:00:05.000+00:00"),
                end_time=datetime.fromisoformat("2021-01-01T00:00:20.000+00:00"),
                attributes={
                    "input": {"value": "xyz"},
                    "retrieval": {
                        "documents": [
                            {},
                            {},
                            {"document": {"content": "C", "score": 3}},
                        ],
                    },
                },
                events=[],
                status_code="OK",
                status_message="okay",
                cumulative_error_count=0,
                cumulative_llm_token_count_prompt=0,
                cumulative_llm_token_count_completion=0,
            )
            .returning(models.Span.id)
        )
        await session.execute(
            insert(models.Span)
            .values(
                trace_rowid=trace_rowid,
                span_id="78910",
                parent_id="2345",
                name="retriever span",
                span_kind="RETRIEVER",
                start_time=datetime.fromisoformat("2021-01-01T00:00:05.000+00:00"),
                end_time=datetime.fromisoformat("2021-01-01T00:00:20.000+00:00"),
                attributes={},
                events=[],
                status_code="OK",
                status_message="okay",
                cumulative_error_count=0,
                cumulative_llm_token_count_prompt=0,
                cumulative_llm_token_count_completion=0,
            )
            .returning(models.Span.id)
        )
        await session.execute(
            insert(models.Span)
            .values(
                trace_rowid=trace_rowid,
                span_id="89101",
                parent_id="2345",
                name="llm span",
                span_kind="LLM",
                start_time=datetime.fromisoformat("2021-01-01T00:00:05.000+00:00"),
                end_time=datetime.fromisoformat("2021-01-01T00:00:20.000+00:00"),
                attributes={
                    "llm": {
                        "input_messages": [
                            {
                                "message": {
                                    "role": "user",
                                    "content": "what is 2 times 3, and what is 2 plus 3",
                                }
                            }
                        ],
                        "output_messages": [
                            {
                                "message": {
                                    "role": "assistant",
                                    "tool_calls": [
                                        {
                                            "tool_call": {
                                                "id": "a",
                                                "function": {
                                                    "name": "multiply",
                                                    "arguments": '{\n  "a": 2,\n  "b": 3\n}',
                                                },
                                            }
                                        },
                                        {
                                            "tool_call": {
                                                "id": "b",
                                                "function": {
                                                    "name": "add",
                                                    "arguments": '{\n  "a": 2,\n  "b": 3\n}',
                                                },
                                            }
                                        },
                                    ],
                                }
                            }
                        ],
                    },
                },
                events=[],
                status_code="OK",
                status_message="okay",
                cumulative_error_count=0,
                cumulative_llm_token_count_prompt=10,
                cumulative_llm_token_count_completion=5,
            )
            .returning(models.Span.id)
        )
        await session.execute(
            insert(models.Span)
            .values(
                trace_rowid=trace_rowid,
                span_id="91011",
                parent_id="2345",
                name="llm span",
                span_kind="LLM",
                start_time=datetime.fromisoformat("2021-01-01T00:00:05.000+00:00"),
                end_time=datetime.fromisoformat("2021-01-01T00:00:20.000+00:00"),
                attributes={
                    "llm": {
                        "input_messages": [{"message": {"role": "user", "content": "call foo"}}],
                        "output_messages": [
                            {
                                "message": {
                                    "role": "assistant",
                                    "tool_calls": [
                                        {
                                            "tool_call": {
                                                "id": "c",
                                                "function": {
                                                    "name": "foo",
                                                },
                                            }
                                        }
                                    ],
                                }
                            }
                        ],
                    },
                },
                events=[],
                status_code="OK",
                status_message="okay",
                cumulative_error_count=0,
                cumulative_llm_token_count_prompt=8,
                cumulative_llm_token_count_completion=4,
            )
            .returning(models.Span.id)
        )
        await session.execute(
            insert(models.Span)
            .values(
                trace_rowid=trace_rowid,
                span_id="111213",
                parent_id="2345",
                name="llm span",
                span_kind="LLM",
                start_time=datetime.fromisoformat("2021-01-01T00:00:25.000+00:00"),
                end_time=datetime.fromisoformat("2021-01-01T00:00:35.000+00:00"),
                attributes={
                    "llm": {
                        "input_messages": [{"message": {"role": "user", "content": "abc"}}],
                        "output_messages": [{"message": {"role": "assistant", "content": "xyz"}}],
                    }
                },
                events=[],
                status_code="OK",
                status_message="okay",
                cumulative_error_count=0,
                cumulative_llm_token_count_prompt=6,
                cumulative_llm_token_count_completion=15,
            )
            .returning(models.Span.id)
        )
        await session.execute(
            insert(models.Span)
            .values(
                trace_rowid=trace_rowid,
                span_id="131415",
                parent_id="2345",
                name="llm span",
                span_kind="LLM",
                start_time=datetime.fromisoformat("2021-01-01T00:00:40.000+00:00"),
                end_time=datetime.fromisoformat("2021-01-01T00:00:50.000+00:00"),
                attributes={
                    "llm": {
                        "input_messages": [
                            {
                                "message": {
                                    "role": "user",
                                    "content": "test empty output",
                                }
                            }
                        ],
                        "output_messages": None,
                    }
                },
                events=[],
                status_code="OK",
                status_message="okay",
                cumulative_error_count=0,
                cumulative_llm_token_count_prompt=5,
                cumulative_llm_token_count_completion=5,
            )
            .returning(models.Span.id)
        )
        await session.execute(
            insert(models.Span)
            .values(
                trace_rowid=trace_rowid,
                span_id="171819",
                parent_id="2345",
                name="llm span",
                span_kind="LLM",
                start_time=datetime.fromisoformat("2021-01-01T00:01:10.000+00:00"),
                end_time=datetime.fromisoformat("2021-01-01T00:01:20.000+00:00"),
                attributes={
                    "llm": {
                        "input_messages": [
                            {
                                "message": {
                                    "role": "user",
                                    "content": "test invalid tool",
                                }
                            }
                        ],
                        "output_messages": [
                            {
                                "message": {
                                    "role": "assistant",
                                    "tool_calls": [
                                        {
                                            "tool_call": {
                                                "id": "invalid",
                                            }
                                        }
                                    ],
                                }
                            }
                        ],
                    }
                },
                events=[],
                status_code="OK",
                status_message="okay",
                cumulative_error_count=0,
                cumulative_llm_token_count_prompt=5,
                cumulative_llm_token_count_completion=5,
            )
            .returning(models.Span.id)
        )


@pytest.fixture
async def abc_project(db: DbSessionFactory) -> None:
    async with db() as session:
        project_row_id = await session.scalar(
            insert(models.Project).values(name="abc").returning(models.Project.id)
        )
        trace_rowid = await session.scalar(
            insert(models.Trace)
            .values(
                trace_id="012",
                project_rowid=project_row_id,
                start_time=datetime.fromisoformat("2021-01-01T00:00:00.000+00:00"),
                end_time=datetime.fromisoformat("2021-01-01T00:01:00.000+00:00"),
            )
            .returning(models.Trace.id)
        )
        await session.execute(
            insert(models.Span)
            .values(
                trace_rowid=trace_rowid,
                span_id="234",
                parent_id="123",
                name="root span",
                span_kind="UNKNOWN",
                start_time=datetime.fromisoformat("2021-01-01T00:00:00.000+00:00"),
                end_time=datetime.fromisoformat("2021-01-01T00:00:30.000+00:00"),
                attributes={
                    "input": {"value": "xy%z*"},
                    "output": {"value": "321"},
                },
                events=[],
                status_code="OK",
                status_message="okay",
                cumulative_error_count=1,
                cumulative_llm_token_count_prompt=100,
                cumulative_llm_token_count_completion=200,
            )
            .returning(models.Span.id)
        )
        span_rowid = await session.scalar(
            insert(models.Span)
            .values(
                trace_rowid=trace_rowid,
                span_id="345",
                parent_id="234",
                name="embedding span",
                span_kind="EMBEDDING",
                start_time=datetime.fromisoformat("2021-01-01T00:00:00.000+00:00"),
                end_time=datetime.fromisoformat("2021-01-01T00:00:05.000+00:00"),
                attributes={
                    "input": {
                        "value": "XY%*Z",
                    },
                    "metadata": {
                        "a.b.c": 123,
                        "1.2.3": "abc",
                        "x.y": {"z.a": {"b.c": 321}},
                    },
                    "embedding": {
                        "model_name": "xyz",
                        "embeddings": [
                            {"embedding": {"vector": [1, 2, 3], "text": "123"}},
                            {"embedding": {"vector": [2, 3, 4], "text": "234"}},
                        ],
                    },
                },
                events=[],
                status_code="OK",
                status_message="no problemo",
                cumulative_error_count=0,
                cumulative_llm_token_count_prompt=0,
                cumulative_llm_token_count_completion=0,
            )
            .returning(models.Span.id)
        )
        await session.execute(
            insert(models.SpanAnnotation).values(
                span_rowid=span_rowid,
                annotator_kind="LLM",
                name="0",
                score=0,
                metadata_={},
                identifier="",
                source="APP",
                user_id=None,
            )
        )
        span_rowid = await session.scalar(
            insert(models.Span)
            .values(
                trace_rowid=trace_rowid,
                span_id="456",
                parent_id="234",
                name="retriever span",
                span_kind="RETRIEVER",
                start_time=datetime.fromisoformat("2021-01-01T00:00:05.000+00:00"),
                end_time=datetime.fromisoformat("2021-01-01T00:00:20.000+00:00"),
                attributes={
                    "attributes": "attributes",
                    "input": {
                        "value": "xy%*z",
                    },
                    "retrieval": {
                        "documents": [
                            {"document": {"content": "A", "score": 1}},
                            {"document": {"content": "B", "score": 2}},
                            {"document": {"content": "C", "score": 3}},
                        ],
                    },
                },
                events=[],
                status_code="OK",
                status_message="okay",
                cumulative_error_count=0,
                cumulative_llm_token_count_prompt=0,
                cumulative_llm_token_count_completion=0,
            )
            .returning(models.Span.id)
        )
        await session.execute(
            insert(models.SpanAnnotation).values(
                span_rowid=span_rowid,
                annotator_kind="LLM",
                name="0",
                score=1,
                metadata_={},
                identifier="",
                source="APP",
                user_id=None,
            )
        )
        await session.execute(
            insert(models.SpanAnnotation).values(
                span_rowid=span_rowid,
                annotator_kind="LLM",
                name="1",
                label="1",
                metadata_={},
                identifier="",
                source="APP",
                user_id=None,
            )
        )
        span_rowid = await session.scalar(
            insert(models.Span)
            .values(
                trace_rowid=trace_rowid,
                span_id="567",
                parent_id="234",
                name="llm span",
                span_kind="LLM",
                start_time=datetime.fromisoformat("2021-01-01T00:00:20.000+00:00"),
                end_time=datetime.fromisoformat("2021-01-01T00:00:30.000+00:00"),
                attributes={
                    "attributes": {"attributes": "attributes"},
                    "llm": {
                        "token_count": {
                            "prompt": 100,
                            "completion": 200,
                        },
                    },
                },
                events=[],
                status_code="ERROR",
                status_message="uh-oh",
                cumulative_error_count=1,
                cumulative_llm_token_count_prompt=100,
                cumulative_llm_token_count_completion=200,
            )
            .returning(models.Span.id)
        )
        await session.execute(
            insert(models.SpanAnnotation).values(
                span_rowid=span_rowid,
                annotator_kind="LLM",
                name="1",
                label="0",
                metadata_={},
                identifier="",
                source="APP",
                user_id=None,
            )
        )
