"""Unit tests for trace retention repository."""

import uuid
from datetime import datetime, timedelta, timezone
from utils.constants import DEFAULT_ORG_ID

import pytest
from sqlalchemy import delete, select
from sqlalchemy.orm import Session

from db_models import (
    DatabaseAgenticAnnotation,
    DatabaseMetric,
    DatabaseMetricResult,
    DatabaseResourceMetadata,
    DatabaseSpan,
    DatabaseTask,
    DatabaseTraceMetadata,
)
from repositories.trace_retention_repository import (
    DEFAULT_TRACE_RETENTION_BATCH_SIZE,
    delete_trace_batch,
    get_expired_trace_ids,
)
from tests.clients.base_test_client import override_get_db_session


def _cleanup_trace_data(
    db_session: Session,
    *,
    task_ids: list[str] | None = None,
    trace_ids: list[str] | None = None,
    span_ids: list[str] | None = None,
    resource_ids: list[str] | None = None,
    metric_ids: list[str] | None = None,
    annotation_ids: list[uuid.UUID] | None = None,
) -> None:
    """Delete test data in FK-safe order and close the session."""
    if annotation_ids:
        db_session.execute(
            delete(DatabaseAgenticAnnotation).where(
                DatabaseAgenticAnnotation.id.in_(annotation_ids)
            )
        )
    if metric_ids or span_ids:
        if span_ids:
            db_session.execute(
                delete(DatabaseMetricResult).where(
                    DatabaseMetricResult.span_id.in_(span_ids)
                )
            )
        if metric_ids:
            db_session.execute(
                delete(DatabaseMetric).where(DatabaseMetric.id.in_(metric_ids))
            )
    if span_ids:
        db_session.execute(delete(DatabaseSpan).where(DatabaseSpan.id.in_(span_ids)))
    if trace_ids:
        db_session.execute(
            delete(DatabaseTraceMetadata).where(
                DatabaseTraceMetadata.trace_id.in_(trace_ids)
            )
        )
    if resource_ids:
        db_session.execute(
            delete(DatabaseResourceMetadata).where(
                DatabaseResourceMetadata.id.in_(resource_ids)
            )
        )
    if task_ids:
        db_session.execute(delete(DatabaseTask).where(DatabaseTask.id.in_(task_ids)))
    db_session.commit()
    db_session.close()


def _make_task(session: Session) -> str:
    """Create a task and return its id."""
    task_id = str(uuid.uuid4())
    now = datetime.now(timezone.utc)
    task = DatabaseTask(
        id=task_id,
        name=f"trace_retention_test_{task_id[:8]}",
        created_at=now,
        updated_at=now,
        is_agentic=False,
        archived=False,
        org_id=DEFAULT_ORG_ID,
    )
    session.add(task)
    session.commit()
    return task_id


@pytest.mark.unit_tests
def test_get_expired_trace_ids_returns_only_expired_and_respects_batch_size() -> None:
    """get_expired_trace_ids returns only trace_ids with end_time < cutoff, capped by batch_size."""
    db_session: Session = override_get_db_session()
    task_id = _make_task(db_session)
    now = datetime.now(timezone.utc)
    cutoff = now - timedelta(days=1)

    # Insert trace_metadata: two expired, one not expired
    expired_1 = str(uuid.uuid4())
    expired_2 = str(uuid.uuid4())
    not_expired = str(uuid.uuid4())
    for trace_id, end_time in [
        (expired_1, now - timedelta(days=2)),
        (expired_2, now - timedelta(days=3)),
        (not_expired, now + timedelta(days=1)),
    ]:
        db_session.add(
            DatabaseTraceMetadata(
                trace_id=trace_id,
                task_id=task_id,
                org_id=DEFAULT_ORG_ID,
                start_time=end_time - timedelta(seconds=10),
                end_time=end_time,
                span_count=0,
            )
        )
    db_session.commit()

    try:
        result = get_expired_trace_ids(db_session, cutoff, batch_size=10)
        assert set(result) == {expired_1, expired_2}
        assert not_expired not in result

        # Respect batch_size; ORDER BY end_time ASC means oldest first
        result_cap = get_expired_trace_ids(db_session, cutoff, batch_size=1)
        assert len(result_cap) == 1
        assert result_cap[0] == expired_2
    finally:
        _cleanup_trace_data(
            db_session,
            task_ids=[task_id],
            trace_ids=[expired_1, expired_2, not_expired],
        )


@pytest.mark.unit_tests
def test_get_expired_trace_ids_returns_nothing_when_no_traces_expired() -> None:
    """get_expired_trace_ids returns empty when cutoff is older than all traces."""
    db_session: Session = override_get_db_session()
    task_id = _make_task(db_session)
    now = datetime.now(timezone.utc)
    trace_id = str(uuid.uuid4())
    db_session.add(
        DatabaseTraceMetadata(
            trace_id=trace_id,
            task_id=task_id,
            org_id=DEFAULT_ORG_ID,
            start_time=now,
            end_time=now,
            span_count=0,
        )
    )
    db_session.commit()
    cutoff = now - timedelta(days=1)

    try:
        result = get_expired_trace_ids(
            db_session, cutoff, batch_size=DEFAULT_TRACE_RETENTION_BATCH_SIZE
        )
        assert result == []
    finally:
        _cleanup_trace_data(
            db_session,
            task_ids=[task_id],
            trace_ids=[trace_id],
        )


@pytest.mark.unit_tests
def test_delete_trace_batch_removes_metric_results_for_deleted_spans() -> None:
    """delete_trace_batch deletes metric_results that reference spans in the batch."""
    db_session: Session = override_get_db_session()
    task_id = _make_task(db_session)
    now = datetime.now(timezone.utc)
    trace_id = str(uuid.uuid4())

    metric_id = str(uuid.uuid4())
    db_session.add(
        DatabaseMetric(
            id=metric_id,
            type="hallucination",
            name="test_metric",
            metric_metadata="{}",
        )
    )
    db_session.commit()

    db_session.add(
        DatabaseTraceMetadata(
            trace_id=trace_id,
            task_id=task_id,
            org_id=DEFAULT_ORG_ID,
            start_time=now - timedelta(days=2),
            end_time=now - timedelta(days=2),
            span_count=1,
        )
    )
    span_id = str(uuid.uuid4())
    db_session.add(
        DatabaseSpan(
            id=span_id,
            trace_id=trace_id,
            span_id=span_id,
            start_time=now - timedelta(days=2),
            end_time=now - timedelta(days=2),
            task_id=task_id,
            org_id=DEFAULT_ORG_ID,
            raw_data={},
            status_code="Unset",
        )
    )
    db_session.commit()

    mr_id_1 = str(uuid.uuid4())
    mr_id_2 = str(uuid.uuid4())
    for mr_id in (mr_id_1, mr_id_2):
        db_session.add(
            DatabaseMetricResult(
                id=mr_id,
                metric_type="hallucination",
                prompt_tokens=0,
                completion_tokens=0,
                latency_ms=0,
                span_id=span_id,
                metric_id=metric_id,
            )
        )
    db_session.commit()

    try:
        delete_trace_batch(db_session, [trace_id])
        db_session.commit()

        assert (
            db_session.execute(
                select(DatabaseSpan).where(DatabaseSpan.id == span_id)
            ).scalar_one_or_none()
            is None
        )
        assert (
            db_session.execute(
                select(DatabaseMetricResult).where(
                    DatabaseMetricResult.span_id == span_id
                )
            ).all()
            == []
        )
        assert (
            db_session.execute(
                select(DatabaseTraceMetadata).where(
                    DatabaseTraceMetadata.trace_id == trace_id
                )
            ).scalar_one_or_none()
            is None
        )
    finally:
        _cleanup_trace_data(
            db_session,
            task_ids=[task_id],
            trace_ids=[trace_id],
            span_ids=[span_id],
            metric_ids=[metric_id],
        )


@pytest.mark.unit_tests
def test_delete_trace_batch_noop_on_empty_list() -> None:
    """delete_trace_batch returns immediately when given an empty list, leaving existing data intact."""
    db_session: Session = override_get_db_session()
    task_id = _make_task(db_session)
    now = datetime.now(timezone.utc)
    trace_id = str(uuid.uuid4())
    db_session.add(
        DatabaseTraceMetadata(
            trace_id=trace_id,
            task_id=task_id,
            org_id=DEFAULT_ORG_ID,
            start_time=now,
            end_time=now,
            span_count=0,
        )
    )
    db_session.commit()

    try:
        delete_trace_batch(db_session, [])
        db_session.commit()

        assert (
            db_session.execute(
                select(DatabaseTraceMetadata).where(
                    DatabaseTraceMetadata.trace_id == trace_id
                )
            ).scalar_one_or_none()
            is not None
        )
    finally:
        _cleanup_trace_data(
            db_session,
            task_ids=[task_id],
            trace_ids=[trace_id],
        )


@pytest.mark.unit_tests
def test_delete_trace_batch_removes_agentic_annotations() -> None:
    """delete_trace_batch deletes agentic annotations for targeted traces and preserves others."""
    db_session: Session = override_get_db_session()
    task_id = _make_task(db_session)
    now = datetime.now(timezone.utc)

    trace_to_delete = str(uuid.uuid4())
    trace_survivor = str(uuid.uuid4())
    for tid in (trace_to_delete, trace_survivor):
        db_session.add(
            DatabaseTraceMetadata(
                trace_id=tid,
                task_id=task_id,
                org_id=DEFAULT_ORG_ID,
                start_time=now - timedelta(days=2),
                end_time=now - timedelta(days=2),
                span_count=0,
            )
        )
    db_session.commit()

    annotation_deleted = uuid.uuid4()
    annotation_survivor = uuid.uuid4()
    db_session.add(
        DatabaseAgenticAnnotation(
            id=annotation_deleted,
            annotation_type="human",
            annotation_score=1,
            trace_id=trace_to_delete,
            created_at=now,
            updated_at=now,
            org_id=DEFAULT_ORG_ID,
        )
    )
    db_session.add(
        DatabaseAgenticAnnotation(
            id=annotation_survivor,
            annotation_type="human",
            annotation_score=0,
            trace_id=trace_survivor,
            created_at=now,
            updated_at=now,
            org_id=DEFAULT_ORG_ID,
        )
    )
    db_session.commit()

    try:
        delete_trace_batch(db_session, [trace_to_delete])
        db_session.commit()

        assert (
            db_session.execute(
                select(DatabaseAgenticAnnotation).where(
                    DatabaseAgenticAnnotation.id == annotation_deleted
                )
            ).scalar_one_or_none()
            is None
        )
        assert (
            db_session.execute(
                select(DatabaseAgenticAnnotation).where(
                    DatabaseAgenticAnnotation.id == annotation_survivor
                )
            ).scalar_one_or_none()
            is not None
        )
    finally:
        _cleanup_trace_data(
            db_session,
            task_ids=[task_id],
            trace_ids=[trace_to_delete, trace_survivor],
            annotation_ids=[annotation_deleted, annotation_survivor],
        )


@pytest.mark.unit_tests
def test_delete_trace_batch_preserves_resource_metadata() -> None:
    """delete_trace_batch never deletes resource_metadata, even when no traces reference it."""
    db_session: Session = override_get_db_session()
    task_id = _make_task(db_session)
    now = datetime.now(timezone.utc)

    resource_id = str(uuid.uuid4())
    db_session.add(
        DatabaseResourceMetadata(
            id=resource_id,
            service_name="test-service",
            resource_attributes={},
        )
    )
    db_session.commit()

    trace_id = str(uuid.uuid4())
    db_session.add(
        DatabaseTraceMetadata(
            trace_id=trace_id,
            task_id=task_id,
            org_id=DEFAULT_ORG_ID,
            root_span_resource_id=resource_id,
            start_time=now - timedelta(days=2),
            end_time=now - timedelta(days=2),
            span_count=1,
        )
    )
    span_id = str(uuid.uuid4())
    db_session.add(
        DatabaseSpan(
            id=span_id,
            trace_id=trace_id,
            span_id=span_id,
            start_time=now - timedelta(days=2),
            end_time=now - timedelta(days=2),
            task_id=task_id,
            org_id=DEFAULT_ORG_ID,
            resource_id=resource_id,
            raw_data={},
            status_code="Unset",
        )
    )
    db_session.commit()

    try:
        delete_trace_batch(db_session, [trace_id])
        db_session.commit()

        assert (
            db_session.execute(
                select(DatabaseTraceMetadata).where(
                    DatabaseTraceMetadata.trace_id == trace_id
                )
            ).scalar_one_or_none()
            is None
        )
        assert (
            db_session.execute(
                select(DatabaseSpan).where(DatabaseSpan.id == span_id)
            ).scalar_one_or_none()
            is None
        )
        assert (
            db_session.execute(
                select(DatabaseResourceMetadata).where(
                    DatabaseResourceMetadata.id == resource_id
                )
            ).scalar_one_or_none()
            is not None
        )
    finally:
        _cleanup_trace_data(
            db_session,
            task_ids=[task_id],
            trace_ids=[trace_id],
            span_ids=[span_id],
            resource_ids=[resource_id],
        )
