import base64
from typing import Any, Optional

try:
    import magic
    from google.genai import Client, types
except ImportError:
    pass

from pydantic import BaseModel, Field

from llmhub.core.llms.template import BaseClientAsync
from llmhub.core.llms.template.response import (
    BaseGenerationModel,
    BaseGenerationResponseModel,
)
from llmhub.core.secret_managers.base import BaseSecretManager


def get_mimetype(file_bytes: bytes) -> str:
    return magic.from_buffer(file_bytes[:2048], mime=True)


class Part(BaseModel):
    """A datatype containing media content.

    Exactly one field within a Part should be set, representing the specific type
    of content being conveyed. Using multiple fields within the same `Part`
    instance is considered invalid.
    """

    video_metadata: Optional[types.VideoMetadata] = Field(
        default=None, description="""Metadata for a given video."""
    )
    thought: Optional[bool] = Field(
        default=None,
        description="""Indicates if the part is thought from the model.""",
    )
    inline_data: Optional[types.Blob] = Field(
        default=None, description="""Optional. Inlined bytes data."""
    )
    file_data: Optional[types.FileData] = Field(
        default=None, description="""Optional. URI based data."""
    )
    thought_signature: Optional[bytes] = Field(
        default=None,
        description="""An opaque signature for the thought so it can be reused in subsequent requests.""",
    )
    code_execution_result: Optional[types.CodeExecutionResult] = Field(
        default=None,
        description="""Optional. Result of executing the [ExecutableCode].""",
    )
    executable_code: Optional[types.ExecutableCode] = Field(
        default=None,
        description="""Optional. Code generated by the model that is meant to be executed.""",
    )
    function_call: Optional[types.FunctionCall] = Field(
        default=None,
        description="""Optional. A predicted [FunctionCall] returned from the model that contains a string representing the [FunctionDeclaration.name] with the parameters and their values.""",
    )
    function_response: Optional[types.FunctionResponse] = Field(
        default=None,
        description="""Optional. The result output of a [FunctionCall] that contains a string representing the [FunctionDeclaration.name] and a structured JSON object containing any output from the function call. It is used as context to the model.""",
    )
    text: Optional[str] = Field(
        default=None, description="""Optional. Text part (can be code)."""
    )

    @classmethod
    def from_uri(cls, *, file_uri: str, mime_type: Optional[str] = None) -> "Part":
        """Creates a Part from a file uri.

        Args:
          file_uri (str): The uri of the file
          mime_type (str): mime_type: The MIME type of the image. If not provided,
            the MIME type will be automatically determined.
        """
        if mime_type is None:
            import mimetypes

            mime_type, _ = mimetypes.guess_type(file_uri)
            if not mime_type:
                raise ValueError(f"Failed to determine mime type for file: {file_uri}")
        file_data = types.FileData(file_uri=file_uri, mime_type=mime_type)
        return cls(file_data=file_data)

    @classmethod
    def from_text(cls, *, text: str, **args) -> "Part":
        return cls(text=text)

    @classmethod
    def from_base64(cls, *, file_data: str, **args) -> "Part":
        file_data = base64.b64decode(file_data)
        inline_data = types.Blob(
            data=file_data,
            mime_type=get_mimetype(file_data),
        )
        return cls(inline_data=inline_data)

    @classmethod
    def from_function_call(cls, *, name: str, args: dict[str, Any]) -> "Part":
        function_call = types.FunctionCall(name=name, args=args)
        return cls(function_call=function_call)

    @classmethod
    def from_function_response(cls, *, name: str, response: dict[str, Any]) -> "Part":
        function_response = types.FunctionResponse(name=name, response=response)
        return cls(function_response=function_response)

    @classmethod
    def from_executable_code(cls, *, code: str, language: types.Language) -> "Part":
        executable_code = types.ExecutableCode(code=code, language=language)
        return cls(executable_code=executable_code)

    @classmethod
    def from_code_execution_result(
        cls, *, outcome: types.Outcome, output: str
    ) -> "Part":
        code_execution_result = types.CodeExecutionResult(
            outcome=outcome, output=output
        )
        return cls(code_execution_result=code_execution_result)


class GeminiClientAsync(BaseClientAsync):
    def __init__(self, secret_manager: BaseSecretManager):
        self.client = Client(api_key=secret_manager.get_secret("GEMINI_API_KEY"))
        self.aio_client = self.client.aio
        self.PART_TYPE_CACTORY = {
            "text": Part.from_text,
            "input_file": Part.from_base64,
        }

    def format_content(self, contents):
        return [
            types.Content(
                role=content["role"],
                parts=[
                    self.PART_TYPE_CACTORY[part["type"]](**part)
                    if not isinstance(part, str)
                    else Part(text=part)
                    for part in content["content"]
                ],
            )
            for content in contents
        ]

    async def create_generation(self, input_model: BaseGenerationModel):
        if input_model.response_schema:
            schema = input_model.response_schema.schema()
            schema["propertyOrdering"] = list(schema["properties"].keys())
        response = await self.aio_client.models.generate_content(
            model=input_model.model,
            contents=self.format_content(input_model.messages),
            config=types.GenerateContentConfig(
                system_instruction=input_model.system,
                response_mime_type="application/json"
                if input_model.response_schema
                else None,
                response_schema=schema if input_model.response_schema else None,
            ),
        )

        parsed_response = BaseGenerationResponseModel(
            created_at=int(response.create_time or 0),
            instructions=input_model.system,
            model=response.model_version,
            object="response",
            output=[
                {
                    "content": [
                        {
                            "annotations": [],
                            "text": j.text,
                            "type": "output_text",
                            "logprobs": [],
                        }
                        for j in i.content.parts
                    ],
                    "id": i.index,
                    "role": i.content.role,
                    "status": "completed",
                    "type": "message",
                }
                for i in response.candidates
            ],
            status="completed",
        )
        return parsed_response
