"""Mixin extending Kili API Gateway class with label related operations."""

import json
from collections.abc import Generator
from typing import Optional

from kili_formats.tool.annotations_to_json_response import AnnotationsToJsonResponseConverter

from kili.adapters.kili_api_gateway.base import BaseOperationMixin
from kili.adapters.kili_api_gateway.helpers.queries import (
    PaginatedGraphQLQuery,
    QueryOptions,
    fragment_builder,
)
from kili.adapters.kili_api_gateway.project.common import get_project
from kili.core.constants import MUTATION_BATCH_SIZE
from kili.core.utils.pagination import batcher
from kili.domain.asset import AssetId
from kili.domain.label import LabelFilters, LabelId
from kili.domain.project import ProjectId
from kili.domain.types import ListOrTuple
from kili.utils.tqdm import tqdm

from .common import get_annotation_fragment
from .formatters import load_label_json_fields
from .mappers import append_label_data_mapper, append_to_labels_data_mapper, label_where_mapper
from .operations import (
    GQL_COPY_LABELS,
    GQL_COUNT_LABELS,
    GQL_DELETE_LABELS,
    get_append_many_labels_mutation,
    get_append_to_labels_mutation,
    get_create_honeypot_mutation,
    get_labels_query,
)
from .types import AppendManyLabelsData, AppendToLabelsData


class LabelOperationMixin(BaseOperationMixin):
    """Mixin extending Kili API Gateway class with label related operations."""

    def count_labels(self, filters: LabelFilters) -> int:
        """Count labels."""
        variables = {"where": label_where_mapper(filters)}
        result = self.graphql_client.execute(GQL_COUNT_LABELS, variables)
        return result["data"]

    def list_labels(
        self,
        filters: LabelFilters,
        fields: ListOrTuple[str],
        options: QueryOptions,
    ) -> Generator[dict, None, None]:
        """List labels."""
        if "jsonResponse" in fields:
            project_info = get_project(
                self.graphql_client, filters.project_id, ("inputType", "jsonInterface")
            )
            # TODO(LAB-4269): TEMPORARY WORKAROUND - Remove when backend handles jsonResponseUrl for LLM
            # Threshold for batching based on number of annotations.
            if project_info["inputType"] in {
                "LLM_RLHF",
                "LLM_INSTR_FOLLOWING",
                "LLM_STATIC",
            }:
                yield from self.llm_list_labels_split(filters, fields, options, project_info)
                return
            if project_info["inputType"] in {"GEOSPATIAL", "VIDEO"}:
                yield from self.list_labels_split(filters, fields, options, project_info)
                return

        fragment = fragment_builder(fields)
        query = get_labels_query(fragment)
        where = label_where_mapper(filters)
        labels_gen = PaginatedGraphQLQuery(self.graphql_client).execute_query_from_paginated_call(
            query, where, options, "Retrieving labels", GQL_COUNT_LABELS
        )
        labels_gen = (
            load_label_json_fields(label, fields, self.http_client) for label in labels_gen
        )
        yield from labels_gen

    def list_labels_split(
        self,
        filters: LabelFilters,
        fields: ListOrTuple[str],
        options: QueryOptions,
        project_info,
    ) -> Generator[dict, None, None]:
        """List labels for VIDEO and GEOSPATIAL projects."""
        if project_info["inputType"] == "VIDEO":
            options = QueryOptions(
                options.disable_tqdm, options.first, options.skip, min(options.batch_size, 20)
            )

        fields = list(fields)

        if "jsonResponse" in fields and "jsonResponseUrl" not in fields:
            fields.append("jsonResponseUrl")
        fragment = fragment_builder(fields)
        query = get_labels_query(fragment)
        where = label_where_mapper(filters)
        labels_gen = PaginatedGraphQLQuery(self.graphql_client).execute_query_from_paginated_call(
            query, where, options, "Retrieving labels", GQL_COUNT_LABELS
        )
        labels_gen = (
            load_label_json_fields(label, fields, self.http_client) for label in labels_gen
        )

        yield from labels_gen

    def llm_list_labels_split(
        self,
        filters: LabelFilters,
        fields: ListOrTuple[str],
        options: QueryOptions,
        project_info,
    ) -> Generator[dict, None, None]:
        """List labels for LLM projects.

        This method handles the specific logic for LLM projects where jsonResponse
        needs to be rebuilt from annotations client-side.
        """
        needs_json_response = "jsonResponse" in fields

        fields = list(fields)

        if needs_json_response:
            # For LLM projects: fetch annotations and rebuild jsonResponse client-side
            inner_annotation_fragment = get_annotation_fragment()
            full_fragment = f"""
                {fragment_builder([f for f in fields if f not in {"jsonResponse", "jsonResponseUrl"}])}
                annotations {{
                    {inner_annotation_fragment}
                }}
            """
        else:
            full_fragment = fragment_builder(fields)

        query = get_labels_query(full_fragment)
        where = label_where_mapper(filters)
        labels_gen = PaginatedGraphQLQuery(self.graphql_client).execute_query_from_paginated_call(
            query, where, options, "Retrieving labels", GQL_COUNT_LABELS
        )
        labels_gen = (
            load_label_json_fields(label, fields, self.http_client) for label in labels_gen
        )

        if needs_json_response:
            # Rebuild jsonResponse from annotations for LLM projects
            converter = AnnotationsToJsonResponseConverter(
                json_interface=project_info["jsonInterface"],
                project_input_type=project_info["inputType"],
            )
            for label in labels_gen:
                asset = None
                converter.patch_label_json_response(asset, label, label["annotations"])
                label.pop("annotations", None)
                yield label
        else:
            yield from labels_gen

    def delete_labels(
        self, ids: ListOrTuple[LabelId], disable_tqdm: Optional[bool]
    ) -> list[LabelId]:
        """Delete labels."""
        delete_label_ids: list[LabelId] = []

        with tqdm(total=len(ids), desc="Deleting labels", disable=disable_tqdm) as pbar:
            for batch_of_label_ids in batcher(ids, batch_size=MUTATION_BATCH_SIZE):
                variables = {"ids": batch_of_label_ids}
                result = self.graphql_client.execute(GQL_DELETE_LABELS, variables)
                batch_deleted_label_ids = result["data"]
                delete_label_ids.extend(batch_deleted_label_ids)
                pbar.update(len(batch_of_label_ids))

        return delete_label_ids

    def append_many_labels(
        self,
        data: AppendManyLabelsData,
        fields: ListOrTuple[str],
        disable_tqdm: Optional[bool],
        project_id: Optional[ProjectId],
    ) -> list[dict]:
        """Append many labels."""
        nb_labels_to_add = len(data.labels_data)

        fragment = fragment_builder(fields)
        query = get_append_many_labels_mutation(fragment=fragment)

        added_labels: list[dict] = []
        with tqdm(total=nb_labels_to_add, desc="Adding labels", disable=disable_tqdm) as pbar:
            for batch_of_label_data in batcher(data.labels_data, batch_size=MUTATION_BATCH_SIZE):
                variables = {
                    "data": {
                        "labelType": data.label_type,
                        "stepName": data.step_name,
                        "overwrite": data.overwrite,
                        "labelsData": [
                            append_label_data_mapper(label) for label in batch_of_label_data
                        ],
                    },
                    "where": {
                        "idIn": [label.asset_id for label in batch_of_label_data],
                    },
                }
                if project_id is not None:
                    variables["where"]["project"] = {"id": project_id}

                # we increase the timeout because the import can take a long time
                batch_result = self.graphql_client.execute(query, variables, timeout=120)
                added_labels.extend(batch_result["data"])
                pbar.update(len(batch_of_label_data))

        return added_labels

    def append_to_labels(
        self, data: AppendToLabelsData, asset_id: AssetId, fields: ListOrTuple[str]
    ) -> dict:
        """Append to labels."""
        fragment = fragment_builder(fields)
        query = get_append_to_labels_mutation(fragment)
        variables = {
            "where": {"id": asset_id},
            "data": append_to_labels_data_mapper(data),
        }
        result = self.graphql_client.execute(query, variables)
        return result["data"]

    def create_honeypot_label(
        self, json_response: dict, asset_id: AssetId, fields: ListOrTuple[str]
    ) -> dict:
        """Create honeypot label."""
        fragment = fragment_builder(fields)
        query = get_create_honeypot_mutation(fragment)
        variables = {
            "data": {"jsonResponse": json.dumps(json_response)},
            "where": {"id": asset_id},
        }
        result = self.graphql_client.execute(query, variables)
        return result["data"]

    def copy_labels(
        self, src_asset_id: str, dst_asset_id: str, project_id: str, skip_rebuilds: bool
    ) -> bool:
        """Copy labels from one asset to another."""
        variables = {
            "data": {
                "srcAssetId": src_asset_id,
                "dstAssetId": dst_asset_id,
                "projectId": project_id,
                "skipRebuilds": skip_rebuilds,
            },
        }
        self.graphql_client.execute(GQL_COPY_LABELS, variables)
        return True
