from datetime import datetime, timedelta
from typing import Literal

import pandas as pd
import pytest
from sqlalchemy import func, select

from phoenix.db import models
from phoenix.server.api.dataloaders.annotation_summaries import AnnotationSummaryDataLoader, Key
from phoenix.server.api.input_types.TimeRange import TimeRange
from phoenix.server.types import DbSessionFactory

from ...._helpers import _add_project, _add_project_session, _add_span, _add_trace


async def test_evaluation_summaries(
    db: DbSessionFactory,
    data_for_testing_dataloaders: None,
) -> None:
    start_time = datetime.fromisoformat("2021-01-01T00:00:10.000+00:00")
    end_time = datetime.fromisoformat("2021-01-01T00:10:00.000+00:00")
    pid = models.Trace.project_rowid
    async with db() as session:
        span_df = await session.run_sync(
            lambda s: pd.read_sql_query(
                select(
                    pid,
                    models.SpanAnnotation.name,
                    func.avg(models.SpanAnnotation.score).label("mean_score"),
                )
                .group_by(pid, models.SpanAnnotation.name)
                .order_by(pid, models.SpanAnnotation.name)
                .join_from(models.Trace, models.Span)
                .join_from(models.Span, models.SpanAnnotation)
                .where(models.Span.name.contains("_trace4_"))
                .where(models.SpanAnnotation.name.in_(("A", "C")))
                .where(start_time <= models.Span.start_time)
                .where(models.Span.start_time < end_time),
                s.connection(),
            )
        )
        trace_df = await session.run_sync(
            lambda s: pd.read_sql_query(
                select(
                    pid,
                    models.TraceAnnotation.name,
                    func.avg(models.TraceAnnotation.score).label("mean_score"),
                )
                .group_by(pid, models.TraceAnnotation.name)
                .order_by(pid, models.TraceAnnotation.name)
                .join_from(models.Trace, models.TraceAnnotation)
                .where(models.TraceAnnotation.name.in_(("B", "D")))
                .where(start_time <= models.Trace.start_time)
                .where(models.Trace.start_time < end_time),
                s.connection(),
            )
        )
    expected = trace_df.loc[:, "mean_score"].to_list() + span_df.loc[:, "mean_score"].to_list()
    kinds: list[Literal["span", "trace"]] = ["trace", "span"]
    session_filter_condition = None
    keys: list[Key] = [
        (
            kind,
            id_ + 1,
            TimeRange(start=start_time, end=end_time),
            "'_trace4_' in name" if kind == "span" else None,
            session_filter_condition,
            eval_name,
        )
        for kind in kinds
        for id_ in range(10)
        for eval_name in (("B", "D") if kind == "trace" else ("A", "C"))
    ]

    summaries = [summary for summary in await AnnotationSummaryDataLoader(db)._load_fn(keys)]
    actual = []
    for summary in summaries:
        assert summary is not None
        actual.append(
            summary.mean_score(),  # type: ignore[call-arg]
        )
    assert actual == pytest.approx(expected, 1e-7)


async def test_session_filter_scopes_fraction_numerator_and_denominator_alike(
    db: DbSessionFactory,
) -> None:
    """A filtered fraction is taken over the filtered sessions only.

    The denominator (distinct annotated sessions) and the numerator (sessions carrying a
    label) are built by separate statements, so a filter reaching only one of them shows up
    here as a fraction diluted by the sessions the filter excluded.
    """
    base_time = datetime.fromisoformat("2021-01-01T00:00:00+00:00")
    async with db() as session:
        project = await _add_project(session)
        # Two traces puts a session in scope of the filter below; one leaves it out.
        for trace_count, label in ((2, "good"), (2, "bad"), (1, "good"), (1, "good")):
            project_session = await _add_project_session(
                session,
                project,
                start_time=base_time,
                end_time=base_time + timedelta(seconds=10),
            )
            for _ in range(trace_count):
                await _add_trace(session, project, project_session, start_time=base_time)
            session.add(
                models.ProjectSessionAnnotation(
                    project_session_id=project_session.id,
                    name="quality",
                    label=label,
                    score=1.0,
                    explanation=None,
                    metadata_={},
                    annotator_kind="HUMAN",
                    source="APP",
                )
            )
        await session.flush()

    result = await AnnotationSummaryDataLoader(db).load(
        (
            "session",
            project.id,
            TimeRange(start=base_time, end=base_time + timedelta(hours=1)),
            None,
            "num_traces >= 2",
            "quality",
        )
    )
    assert result is not None
    label_fractions = {lf.label: lf.fraction for lf in result.label_fractions()}  # type: ignore[call-arg, attr-defined]
    # Of the two filtered sessions, one is "good" and one is "bad". Counting all four
    # sessions in the denominator would report 0.25 apiece instead.
    assert label_fractions["good"] == pytest.approx(0.5, rel=1e-4)
    assert label_fractions["bad"] == pytest.approx(0.5, rel=1e-4)


async def test_span_filter_scopes_fraction_numerator_and_denominator_alike(
    db: DbSessionFactory,
) -> None:
    base_time = datetime.fromisoformat("2021-01-01T00:00:00+00:00")
    async with db() as session:
        project = await _add_project(session)
        trace = await _add_trace(session, project, start_time=base_time)
        for span_kind, label in (("LLM", "good"), ("CHAIN", "bad")):
            span = await _add_span(session, trace, span_kind=span_kind, start_time=base_time)
            session.add(
                models.SpanAnnotation(
                    span_rowid=span.id,
                    name="quality",
                    label=label,
                    score=1.0,
                    explanation=None,
                    metadata_={},
                    annotator_kind="HUMAN",
                    source="APP",
                )
            )
        await session.flush()

    result = await AnnotationSummaryDataLoader(db).load(
        (
            "span",
            project.id,
            TimeRange(start=base_time, end=base_time + timedelta(hours=1)),
            "span_kind == 'LLM'",
            None,
            "quality",
        )
    )
    assert result is not None
    label_fractions = {lf.label: lf.fraction for lf in result.label_fractions()}  # type: ignore[call-arg, attr-defined]
    assert label_fractions["good"] == pytest.approx(1.0, rel=1e-4)
    assert sum(label_fractions.values()) == pytest.approx(1.0, rel=1e-4)


async def test_multiple_annotations_score_weighting(
    db: DbSessionFactory,
    data_with_multiple_annotations: None,
) -> None:
    # Using the "quality" annotations fixture.
    start_time = datetime.fromisoformat("2021-01-01T00:00:00.000+00:00")
    end_time = datetime.fromisoformat("2021-01-01T01:00:00.000+00:00")
    # Based on the fixture:
    # Span 1: avg score = (0.85+0.95+0.3)/3 = 0.70
    # Span 2: avg score = (0.85+0.3)/2 = 0.575
    # Span 3: avg score = 0.85
    # Overall average score = (0.70+0.575+0.85)/3 ≈ 0.70833.
    expected_avg_score = 0.70833

    async with db() as session:
        project_id = await session.scalar(
            select(models.Project.id).where(models.Project.name == "simple_multiple")
        )
        assert isinstance(project_id, int)

    loader = AnnotationSummaryDataLoader(db)
    filter_condition = None
    session_filter_condition = None
    result = await loader.load(
        (
            "span",
            project_id,
            TimeRange(start=start_time, end=end_time),
            filter_condition,
            session_filter_condition,
            "quality",
        )
    )
    assert result is not None
    assert result.mean_score() == pytest.approx(expected_avg_score, rel=1e-4)  # type: ignore[call-arg]

    # Expected fractions:
    # "good": (2/3 + 1/2 + 1) / 3 ≈ 0.722
    # "bad": (1/3 + 1/2 + 0) / 3 ≈ 0.277
    label_fracs = {lf.label: lf.fraction for lf in result.label_fractions()}  # type: ignore[call-arg, attr-defined]
    assert label_fracs["good"] == pytest.approx(0.722, rel=1e-2)
    assert label_fracs["bad"] == pytest.approx(0.277, rel=1e-2)
    assert abs(label_fracs["good"] + label_fracs["bad"] - 1.0) < 1e-2


async def test_missing_label_aggregation(
    db: DbSessionFactory,
    data_with_missing_labels: None,
) -> None:
    # Using the "distribution" annotations fixture.
    start_time = datetime.fromisoformat("2021-01-01T00:00:00.000+00:00")
    end_time = datetime.fromisoformat("2021-01-01T01:00:00.000+00:00")
    # Based on the fixture:
    # Span 1: For "distribution": "X" fraction = 2/3, "Y" fraction = 1/3.
    # Span 2: "X" fraction = 1.
    # Span 3: "X" fraction = 1.
    # Overall label fractions for "distribution" annotation:
    #   "X": (0.667 + 1 + 1) / 3 ≈ 0.889,
    #   "Y": (0.333 + 0 + 0) / 3 ≈ 0.111.
    loader = AnnotationSummaryDataLoader(db)

    async with db() as session:
        project_id = await session.scalar(
            select(models.Project.id).where(models.Project.name == "simple_missing")
        )
        assert isinstance(project_id, int)
    filter_condition = None
    session_filter_condition = None
    result = await loader.load(
        (
            "span",
            project_id,
            TimeRange(start=start_time, end=end_time),
            filter_condition,
            session_filter_condition,
            "distribution",
        )
    )
    assert result is not None

    label_fracs = {lf.label: lf.fraction for lf in result.label_fractions()}  # type: ignore[call-arg, attr-defined]
    assert label_fracs["X"] == pytest.approx(0.889, rel=1e-2)
    assert label_fracs["Y"] == pytest.approx(0.111, rel=1e-2)
    assert abs(sum(label_fracs.values()) - 1.0) < 1e-7

    # The "distribution" annotation is grouped as follows:
    # Span 1: .8, .8, .6
    # Span 2: .8
    # Span 3: .8
    # Overall average = ((0.8 + 0.8 + 0.6) / 3 + 0.8 + 0.8) / 3 ≈ 0.777
    assert result.mean_score() == pytest.approx(0.777, rel=1e-2)  # type: ignore[call-arg]


async def test_mixed_emit_coverage_counts(
    db: DbSessionFactory,
    data_with_mixed_emit: None,
) -> None:
    start_time = datetime.fromisoformat("2021-01-01T00:00:00.000+00:00")
    end_time = datetime.fromisoformat("2021-01-01T01:00:00.000+00:00")

    async with db() as session:
        project_id = await session.scalar(
            select(models.Project.id).where(models.Project.name == "mixed_emit")
        )
        assert isinstance(project_id, int)

    result = await AnnotationSummaryDataLoader(db).load(
        (
            "span",
            project_id,
            TimeRange(start=start_time, end=end_time),
            None,
            None,
            "coverage",
        )
    )
    assert result is not None
    assert result.score_count() == 60  # type: ignore[call-arg, comparison-overlap]
    assert result.label_count() == 40  # type: ignore[call-arg, comparison-overlap]
    assert result.count() == 100  # type: ignore[call-arg, comparison-overlap]

    expected_mean = sum(i / 100.0 for i in range(60)) / 60
    assert result.mean_score() == pytest.approx(expected_mean, rel=1e-4)  # type: ignore[call-arg]

    # label_fractions are entity-weighted: each entity counts once in the denominator
    # (100 total entities), so PASS=30/100=0.30, FAIL=10/100=0.10.
    label_fracs = {lf.label: lf.fraction for lf in result.label_fractions()}  # type: ignore[call-arg, attr-defined]
    assert label_fracs["PASS"] == pytest.approx(0.30, rel=1e-4)
    assert label_fracs["FAIL"] == pytest.approx(0.10, rel=1e-4)


async def test_pure_score_only_coverage_counts(
    db: DbSessionFactory,
    data_with_null_labels: None,
) -> None:
    start_time = datetime.fromisoformat("2021-01-01T00:00:00.000+00:00")
    end_time = datetime.fromisoformat("2021-01-01T01:00:00.000+00:00")

    async with db() as session:
        project_id = await session.scalar(
            select(models.Project.id).where(models.Project.name == "null_labels")
        )
        assert isinstance(project_id, int)

    result = await AnnotationSummaryDataLoader(db).load(
        (
            "span",
            project_id,
            TimeRange(start=start_time, end=end_time),
            None,
            None,
            "unlabeled",
        )
    )
    assert result is not None
    assert result.label_count() == 0  # type: ignore[call-arg, comparison-overlap]
    assert result.label_fractions() == []  # type: ignore[call-arg, comparison-overlap]
    assert result.score_count() > 0  # type: ignore[call-arg, operator]


async def test_null_label_handling(
    db: DbSessionFactory,
    data_with_null_labels: None,
) -> None:
    """Ensure that the loader does not raise when all labels are NULL.

    The expected behavior is:
    * label_fractions() returns an empty list.
    * mean_score() computes the per-entity average score correctly.
    """
    start_time = datetime.fromisoformat("2021-01-01T00:00:00.000+00:00")
    end_time = datetime.fromisoformat("2021-01-01T01:00:00.000+00:00")

    async with db() as session:
        project_id = await session.scalar(
            select(models.Project.id).where(models.Project.name == "null_labels")
        )
        assert isinstance(project_id, int)

    loader = AnnotationSummaryDataLoader(db)
    filter_condition = None
    session_filter_condition = None
    result = await loader.load(
        (
            "span",
            project_id,
            TimeRange(start=start_time, end=end_time),
            filter_condition,
            session_filter_condition,
            "unlabeled",
        )
    )

    # Should not be None and should have no label fractions.
    assert result is not None
    assert result.label_fractions() == []  # type: ignore

    # Each span has 2 scores. Compute expected overall average.
    # Span averages: (0.5+0.9)/2 = 0.7, (0.6+0.8)/2 = 0.7, (0.4+1.0)/2 = 0.7.
    expected_avg = 0.7
    assert result.mean_score() == pytest.approx(expected_avg, rel=1e-4)  # type: ignore[call-arg]
