import logging
from datetime import datetime
from re import sub
from typing import Any, Dict, List, Optional, Union
from uuid import UUID

from pydantic import BaseModel

from highlighter.client.base_models.base_models import SubmissionType
from highlighter.client.base_models.entity import Entity
from highlighter.client.base_models.image import Image
from highlighter.core.gql_base_model import GQLBaseModel, snake_to_camel

from ..core import paginate
from .base_models import SubmissionTypeConnection as SubmissionConnection
from .gql_client import HLClient

__all__ = [
    "get_latest_assessments_gen",
    "get_assessments_gen",
    "create_assessment_with_avro_file",
    "append_data_files_to_not_finalised_assessment",
    "create_assessment_not_finalised",
    "finalise",
    "create_assessment_from_entities",
]

logger = logging.getLogger(__name__)


def get_latest_assessments_gen(
    client: HLClient,
    **kwargs,
):
    query_args = {k: v for k, v in kwargs.items() if v is not None}
    assessments_gen = paginate(
        client.latestSubmissionConnection,
        SubmissionConnection,
        **query_args,
    )
    return assessments_gen


def get_assessments_gen(
    client: HLClient,
    **kwargs,
):
    query_args = {k: v for k, v in kwargs.items() if v is not None}
    assessments_gen = paginate(
        client.assessmentConnection,
        SubmissionConnection,
        **query_args,
    )
    return assessments_gen


def create_assessment_with_avro_file(client: HLClient, workflow_id: int, file_id: int, avro_file_info: dict):
    class CreateAssessmentPayload(BaseModel):
        errors: Optional[Any] = None
        submission: Optional[SubmissionType] = None

    result = client.createSubmission(
        return_type=CreateAssessmentPayload,
        workflowId=workflow_id,
        dataFileIds=[file_id],
        backgroundInfoLayerFileData=avro_file_info,
        status="completed",
    )
    if len(result.errors) > 0:
        raise RuntimeError(f"Error creating assessment: {result.errors}")


# ToDo: Update code that calls this
# to account for the addition of case_id
def create_assessment_not_finalised(
    client: HLClient,
    task_id=UUID,
    user_id: Optional[int] = None,
    model_id: Optional[int] = None,
    training_run_id: Optional[int] = None,
    data_source_id: Optional[int] = None,
    flag_reason: Optional[str] = None,
    started_at: Optional[str] = None,
) -> SubmissionType:
    kwargs = locals()
    kwargs.pop("client")
    kwargs = {snake_to_camel(k): v for k, v in kwargs.items() if v is not None}

    class CreateSubmissionNotFinalisedPayload(GQLBaseModel):
        errors: Optional[Any] = None
        submission: Optional[SubmissionType] = None

    payload = client.createSubmissionNotFinalised(return_type=CreateSubmissionNotFinalisedPayload, **kwargs)
    if payload.errors:
        raise ValueError(payload.errors)
    return payload.submission


def append_data_files_to_not_finalised_assessment(
    client: HLClient,
    not_finalised_assessment: SubmissionType,
    data_files: List[Image],
) -> SubmissionType:
    kwargs = {
        "submissionId": not_finalised_assessment.id,
        "dataFileIds": [str(d.uuid) for d in data_files],
    }

    class Payload(BaseModel):
        errors: Optional[Any]
        submission: Optional[SubmissionType] = None

    logger.debug(f"Appending data files to submission {not_finalised_assessment.id}")
    payload = client.appendDataFilesToNotFinalisedSubmission(return_type=Payload, **kwargs)
    if payload.errors:
        raise ValueError(payload.errors)

    return payload.submission


def finalise(
    client: HLClient,
    not_finalised_assessment: SubmissionType,
) -> SubmissionType:
    kwargs = {"submissionId": not_finalised_assessment.id}

    class Payload(BaseModel):
        errors: Optional[Any]
        submission: Optional[SubmissionType] = None

    logger.debug(f"Finalising submission {not_finalised_assessment.id}")
    payload = client.finaliseSubmissionNotFinalised(return_type=Payload, **kwargs)
    if payload.errors:
        raise ValueError(payload.errors)

    return payload.submission


def create_assessment_from_entities(
    client: HLClient,
    entities: Dict[UUID, Entity],
    data_file_ids: List[Union[UUID, str]],
    task_id: Optional[UUID],
    workflow_id: Optional[int] = None,
    user_id: Optional[int] = None,
    started_at: Optional[datetime] = None,
):
    class CreateAssessmentPayload(BaseModel):
        errors: List[str]
        submission: Optional[SubmissionType] = None

    annotations_attributes = []
    eavt_attributes = []
    for entity in entities.values():
        for annotation in entity.annotations:
            annotations_attributes.append(annotation.gql_dict())
            for eavt in annotation.observations:
                eavt_attributes.append({**eavt.gql_dict(), "annotationUuid": str(annotation.id)})
        for eavt in entity.global_observations:
            eavt_attributes.append(eavt.gql_dict())

    # collect all optional kwargs to createSubmission here
    kwargs = {
        k: v
        for k, v in {
            "taskId": str(task_id) if task_id else None,
            "workflowId": workflow_id,
            "userId": user_id,
            "startedAt": started_at.isoformat() if isinstance(started_at, datetime) else started_at,
        }.items()
        if v is not None
    }

    result = client.createSubmission(
        return_type=CreateAssessmentPayload,
        status="completed",
        annotationsAttributes=annotations_attributes,
        eavtAttributes=eavt_attributes,
        dataFileIds=[str(id) for id in data_file_ids],
        **kwargs,
    )
    if len(result.errors) > 0:
        raise ValueError(f"GraphQL Error: {result.errors}")
    return result.submission
