import logging
from dataclasses import dataclass, field
from datetime import datetime, timezone
from typing import Any, Dict, Iterator, List, Optional, Union
from uuid import uuid4

from ..parsers.generation_parser import parse_model_parameters, parse_result
from ..writer import LogWriter
from .base import BaseContainer
from .types import Entity, GenerationError, object_to_dict


@dataclass
class GenerationRequestMessage:
    role: str
    content: str

    def to_dict(self) -> Dict[str, str]:
        return {"role": self.role, "content": self.content}

    @classmethod
    def from_gemini_content(cls, content: Any) -> "GenerationRequestMessage":
        if "role" not in content or "parts" not in content:
            raise ValueError("[MaximSDK] Invalid Gemini content")
        if not isinstance(content["parts"], list):
            raise ValueError("[MaximSDK] Invalid parts in Gemini content.")
        parts_content = ""
        for part in content["parts"]:
            parts_content += part["text"]
        return GenerationRequestMessage(role=content["role"], content=parts_content)


@dataclass
class GenerationConfig:
    id: str
    provider: str
    model: str
    messages: Optional[List[GenerationRequestMessage]] = field(default_factory=list)
    model_parameters: Dict[str, Any] = field(default_factory=dict)
    span_id: Optional[str] = None
    name: Optional[str] = None
    maxim_prompt_id: Optional[str] = None
    tags: Optional[Dict[str, str]] = None


valid_providers = [
    "openai",
    "azure",
    "anthropic",
    "huggingface",
    "together",
    "google",
    "groq",
    "bedrock",
    "cohere",
    "unknown",
]


@dataclass
class GenerationToolCallFunction:
    name: str
    arguments: Optional[str]


@dataclass
class GenerationToolCall:
    id: str
    type: str
    function: GenerationToolCallFunction


@dataclass
class GenerationResultMessage:
    role: str
    content: Optional[str]
    tool_calls: Optional[List[GenerationToolCall]]


@dataclass
class GenerationResultChoice:
    index: int
    message: GenerationResultMessage
    logprobs: Optional[Any] = None
    finish_reason: Optional[str] = None


@dataclass
class GenerationUsage:
    prompt_tokens: int
    completion_tokens: int
    total_tokens: int


@dataclass
class GenerationResult:
    id: str
    object: str
    created: int
    model: str
    choices: List[GenerationResultChoice]
    usage: GenerationUsage


class Generation(BaseContainer):
    def __init__(self, config: GenerationConfig, writer: LogWriter):
        super().__init__(Entity.GENERATION, config.__dict__, writer)
        self.model = config.model
        self.maxim_prompt_id = config.maxim_prompt_id
        self.messages = []
        self.provider = config.provider
        self.provider = config.provider.lower()
        if self.provider not in valid_providers:
            self.provider = "unknown"
        if config.messages is not None:
            self.messages.extend([m.to_dict() for m in config.messages])
        self.model_parameters = parse_model_parameters(config.model_parameters)

    def _set_provider(self, writer: LogWriter, id: str, provider: str):
        if provider not in valid_providers:
            raise ValueError(
                f"Invalid provider: {self.provider}. Must be one of {', '.join(valid_providers)}."
            )
        BaseContainer._commit_(
            writer, Entity.GENERATION, id, "update", {"provider": provider}
        )

    def set_provider(self, provider: str):
        if provider not in valid_providers:
            raise ValueError(
                f"Invalid provider: {self.provider}. Must be one of {', '.join(valid_providers)}."
            )
        self.provider = provider

    @staticmethod
    def set_model_(writer: LogWriter, id: str, model: str):
        BaseContainer._commit_(
            writer, Entity.GENERATION, id, "update", {"model": model}
        )

    def set_model(self, model: str):
        self.model = model
        self._commit("update", {"model": model})

    @staticmethod
    def add_message_(writer: LogWriter, id: str, message: GenerationRequestMessage):
        BaseContainer._commit_(
            writer, Entity.GENERATION, id, "update", {"messages": [message.to_dict()]}
        )

    def add_message(self, message: GenerationRequestMessage):
        self.messages.append(message)
        self._commit("update", {"messages": [message.to_dict()]})

    @staticmethod
    def set_model_parameters_(
        writer: LogWriter, id: str, model_parameters: Dict[str, Any]
    ):
        BaseContainer._commit_(
            writer,
            Entity.GENERATION,
            id,
            "update",
            {"model_parameters": model_parameters},
        )

    def set_model_parameters(self, model_parameters: Dict[str, Any]):
        self.model_parameters = model_parameters
        self._commit("update", {"model_parameters": model_parameters})

    @staticmethod
    def result_(
        writer: LogWriter, id: str, result: Union[GenerationResult, Dict[str, Any]]
    ):
        try:
            # Checking the type
            result = Generation.convert_result(result)
            # Validating the result
            parse_result(result)
            BaseContainer._commit_(
                writer, Entity.GENERATION, id, "result", {"result": result}
            )
            BaseContainer._end_(
                writer,
                Entity.GENERATION,
                id,
                {
                    "endTimestamp": datetime.now(timezone.utc),
                },
            )
        except Exception as e:
            logging.error(
                f"[MaximSDK] Invalid result. You can pass OpenAI/Azure ChatCompletion or Langchain LLMResult,AIMessage,ToolMessage or LiteLLM ModelResponse: {str(e)}"
            )

    @staticmethod
    def end_(writer: LogWriter, id: str, data: Optional[Dict[str, Any]] = None):
        if data is None:
            data = {}
        BaseContainer._end_(
            writer,
            Entity.GENERATION,
            id,
            {
                "endTimestamp": datetime.now(timezone.utc),
                **data,
            },
        )

    @staticmethod
    def add_tag_(writer: LogWriter, id: str, key: str, value: str):
        BaseContainer._add_tag_(writer, Entity.GENERATION, id, key, value)

    @staticmethod
    def convert_chat_completion(chat_completion: Dict[str, Any]):
        return {
            "id": chat_completion.get("id", str(uuid4())),
            "created": chat_completion.get("created", datetime.now(timezone.utc)),
            "choices": [
                {
                    "index": choice.get("index", 0),
                    "message": {
                        "role": choice.get("message").get("role", "assistant"),
                        "content": choice.get("message").get("content", ""),
                        "tool_calls": choice.get("message").get("tool_calls", None),
                        "function_calls": choice.get("message").get(
                            "function_calls", None
                        ),
                    },
                    "finish_reason": choice.get("finish_reason", None),
                    "logprobs": choice.get("logprobs", None),
                }
                for choice in chat_completion.get("choices", [])
            ],
            "usage": chat_completion.get("usage", {}),
        }

    @staticmethod
    def convert_result(
        result: Union[Any, GenerationResult, Dict[str, Any]],
    ) -> Union[Any, GenerationResult, Dict[str, Any]]:
        try:
            parse_result(result)
            return result
        except Exception:
            if isinstance(result, object):
                # trying for langchain first
                try:
                    from langchain_core.messages import AIMessage  # type: ignore
                    from langchain_core.outputs import LLMResult  # type: ignore

                    from ..langchain.utils import (
                        parse_base_message_to_maxim_generation,
                        parse_langchain_llm_result,
                    )

                    if isinstance(result, AIMessage):
                        return parse_base_message_to_maxim_generation(result)
                    elif isinstance(result, LLMResult):
                        return parse_langchain_llm_result(result)
                except ImportError:
                    pass
                # trying for litellm
                try:
                    from litellm.types.utils import ModelResponse  # type: ignore

                    from ..litellm.parser import parse_litellm_model_response

                    if isinstance(result, ModelResponse):
                        return parse_litellm_model_response(result)
                except ImportError:
                    pass
                # trying gemini response
                try:
                    from google.genai.types import GenerateContentResponse

                    from ..gemini.utils import GeminiUtils

                    if isinstance(result, GenerateContentResponse):
                        return GeminiUtils.parse_gemini_generation_content(result)
                    elif isinstance(result, Iterator):
                        return GeminiUtils.parse_gemini_generation_content_iterator(
                            result
                        )
                except ImportError:
                    pass
            result_dict = object_to_dict(result)
            if isinstance(result_dict, Dict):
                # Checking if its Azure or OpenAI result
                if (
                    "object" in result_dict
                    and result_dict["object"] == "chat.completion"
                ):
                    return Generation.convert_chat_completion(result_dict)
                elif (
                    "object" in result_dict
                    and result_dict["object"] == "text.completion"
                ):
                    raise ValueError("Text completion is not yet supported.")

            return result

    def result(self, result: Any):
        try:
            # Checking the type
            result = Generation.convert_result(result)
            # Validating the result
            parse_result(result)
            # Logging the result
            self._commit("result", {"result": result})
            self.end()
        except ValueError as e:
            logging.error(
                f"[MaximSDK] Invalid result. You can pass OpenAI/Azure ChatCompletion or Langchain LLMResult,AIMessage,ToolMessage or LiteLLM ModelResponse: {str(e)}"
            )

    def error(self, error: GenerationError):
        if not error.code:
            error.code = ""
        if not error.type:
            error.type = ""
        self._commit(
            "result",
            {
                "result": {
                    "error": {
                        "message": error.message,
                        "code": error.code,
                        "type": error.type,
                    },
                    "id": str(uuid4()),
                }
            },
        )
        self.end()

    @staticmethod
    def error_(writer: LogWriter, id: str, error: GenerationError):
        if not error.code:
            error.code = ""
        if not error.type:
            error.type = ""
        BaseContainer._commit_(
            writer,
            Entity.GENERATION,
            id,
            "result",
            {
                "result": {
                    "error": {
                        "message": error.message,
                        "code": error.code,
                        "type": error.type,
                    },
                    "id": str(uuid4()),
                }
            },
        )
        BaseContainer._end_(
            writer,
            Entity.GENERATION,
            id,
            {
                "endTimestamp": datetime.now(timezone.utc),
            },
        )

    def data(self) -> Dict[str, Any]:
        base_data = super().data()
        return {
            **base_data,
            "model": self.model,
            "provider": self.provider,
            "maximPromptId": self.maxim_prompt_id,
            "messages": self.messages,
            "modelParameters": self.model_parameters,
        }
