from typing import List, Sequence

import strawberry
from sqlalchemy import delete, insert, update
from strawberry import UNSET
from strawberry.types import Info

from phoenix.db import models
from phoenix.server.api.context import Context
from phoenix.server.api.input_types.CreateTraceAnnotationsInput import CreateTraceAnnotationsInput
from phoenix.server.api.input_types.DeleteAnnotationsInput import DeleteAnnotationsInput
from phoenix.server.api.input_types.PatchAnnotationsInput import PatchAnnotationsInput
from phoenix.server.api.mutations.auth import IsAuthenticated
from phoenix.server.api.types.node import from_global_id_with_expected_type
from phoenix.server.api.types.TraceAnnotation import TraceAnnotation, to_gql_trace_annotation


@strawberry.type
class TraceAnnotationMutationPayload:
    trace_annotations: List[TraceAnnotation]


@strawberry.type
class TraceAnnotationMutationMixin:
    @strawberry.mutation(permission_classes=[IsAuthenticated])  # type: ignore
    async def create_trace_annotations(
        self, info: Info[Context, None], input: List[CreateTraceAnnotationsInput]
    ) -> TraceAnnotationMutationPayload:
        inserted_annotations: Sequence[models.TraceAnnotation] = []
        async with info.context.db() as session:
            values_list = [
                dict(
                    trace_rowid=from_global_id_with_expected_type(annotation.trace_id, "Trace"),
                    name=annotation.name,
                    label=annotation.label,
                    score=annotation.score,
                    explanation=annotation.explanation,
                    annotator_kind=annotation.annotator_kind,
                    metadata_=annotation.metadata,
                )
                for annotation in input
            ]
            stmt = (
                insert(models.TraceAnnotation).values(values_list).returning(models.TraceAnnotation)
            )
            result = await session.scalars(stmt)
            inserted_annotations = result.all()

        return TraceAnnotationMutationPayload(
            trace_annotations=[
                to_gql_trace_annotation(annotation) for annotation in inserted_annotations
            ]
        )

    @strawberry.mutation(permission_classes=[IsAuthenticated])  # type: ignore
    async def patch_trace_annotations(
        self, info: Info[Context, None], input: List[PatchAnnotationsInput]
    ) -> TraceAnnotationMutationPayload:
        patched_annotations = []
        async with info.context.db() as session:
            for annotation in input:
                trace_annotation_id = from_global_id_with_expected_type(
                    annotation.annotation_id, "TraceAnnotation"
                )
                patch = {
                    column.key: patch_value
                    for column, patch_value, column_is_nullable in (
                        (models.TraceAnnotation.name, annotation.name, False),
                        (models.TraceAnnotation.annotator_kind, annotation.annotator_kind, False),
                        (models.TraceAnnotation.label, annotation.label, True),
                        (models.TraceAnnotation.score, annotation.score, True),
                        (models.TraceAnnotation.explanation, annotation.explanation, True),
                        (models.TraceAnnotation.metadata_, annotation.metadata, False),
                    )
                    if patch_value is not UNSET and (patch_value is not None or column_is_nullable)
                }
                trace_annotation = await session.scalar(
                    update(models.TraceAnnotation)
                    .where(models.TraceAnnotation.id == trace_annotation_id)
                    .values(**patch)
                    .returning(models.TraceAnnotation)
                )
                if trace_annotation:
                    patched_annotations.append(to_gql_trace_annotation(trace_annotation))

        return TraceAnnotationMutationPayload(trace_annotations=patched_annotations)

    @strawberry.mutation(permission_classes=[IsAuthenticated])  # type: ignore
    async def delete_trace_annotations(
        self, info: Info[Context, None], input: DeleteAnnotationsInput
    ) -> TraceAnnotationMutationPayload:
        trace_annotation_ids = [
            from_global_id_with_expected_type(global_id, "TraceAnnotation")
            for global_id in input.annotation_ids
        ]
        async with info.context.db() as session:
            stmt = (
                delete(models.TraceAnnotation)
                .where(models.TraceAnnotation.id.in_(trace_annotation_ids))
                .returning(models.TraceAnnotation)
            )
            result = await session.scalars(stmt)
            deleted_annotations = result.all()

            deleted_annotations_gql = [
                to_gql_trace_annotation(annotation) for annotation in deleted_annotations
            ]
        return TraceAnnotationMutationPayload(trace_annotations=deleted_annotations_gql)
