import logging
from abc import ABC, abstractmethod
from typing import Optional, Any, Callable, Literal, TypeVar, Self, Protocol, runtime_checkable
from enum import Enum

from pydantic import BaseModel, model_validator, ConfigDict


logger = logging.getLogger(__name__)

type MessagesDict = list[dict[str, str]]
type Role = Literal["user", "model"]
type Json = list | dict
type JsonType = Literal["string", "number", "integer", "boolean", "array", "object"]
PydanticStructure = TypeVar("PydanticStructure", bound=BaseModel)

type ResultType = Literal["json"] | type[PydanticStructure] | None


@runtime_checkable
class PartLike(Protocol):
    """Protocol for Part-like objects that have the same interface as Part."""
    text: Optional[str]
    function_call: Optional[Any]  # Using Any to allow different FunctionCall types
    function_response: Optional[Any]  # Using Any to allow different FunctionResponse types
    thought: Optional[bool]
    inline_data: Optional[Any]  # Using Any to allow different Blob types


class CustomApiKey(ABC):
    @abstractmethod
    def __hash__(self):
        pass

type ApiKey = str | CustomApiKey

type MessageDict = dict[str, str | Any | None]

class Message(BaseModel):
    role: str
    content: str | dict[str, Any] | None = None

class Choice(BaseModel):
    message: Message

class Usage(BaseModel):
    prompt_tokens: int
    completion_tokens: int
    total_tokens: int

class Completion(BaseModel):
    choices: list[Choice]
    usage: Optional[Usage] = None

class FunctionCall(BaseModel):
    id: Optional[str] = None
    args: Optional[dict[str, Any]] = None
    name: Optional[str] = None

class FunctionResponse(BaseModel):
    id: Optional[str] = None
    name: Optional[str] = None
    response: Optional[dict[str, Any]] = None

class Blob(BaseModel):
    display_name: Optional[str] = None
    data: Optional[bytes] = None
    mime_type: Optional[str] = None


class Part(BaseModel):
    text: Optional[str] = None
    function_call: Optional[FunctionCall] = None
    function_response: Optional[FunctionResponse] = None
    thought: Optional[bool] = None
    inline_data: Optional[Blob] = None
    
    def as_str(self) -> str:
        if self.text is not None:
            return self.text
        return ""

    @classmethod
    def from_bytes(cls, *, data: bytes, mime_type: str, display_name: str | None = None) -> 'Part':
        inline_data = Blob(
            data=data,
            mime_type=mime_type,
            display_name=display_name
        )
        return cls(inline_data=inline_data)

class Content(BaseModel):
    model_config = ConfigDict(arbitrary_types_allowed=True)
    
    parts: list[Part | PartLike] | None = None
    role: Role | None = None
    
    def as_str(self) -> str:
        if self.parts is None:
            return ""
        else:
            return "\n".join([(part.text or "") for part in self.parts])


class FinishReason(Enum):
  """Output only. The reason why the model stopped generating tokens.

  If empty, the model has not stopped generating the tokens.
  """

  FINISH_REASON_UNSPECIFIED = 'FINISH_REASON_UNSPECIFIED'
  """The finish reason is unspecified."""
  STOP = 'STOP'
  """Token generation reached a natural stopping point or a configured stop sequence."""
  MAX_TOKENS = 'MAX_TOKENS'
  """Token generation reached the configured maximum output tokens."""
  SAFETY = 'SAFETY'
  """Token generation stopped because the content potentially contains safety violations. NOTE: When streaming, [content][] is empty if content filters blocks the output."""
  RECITATION = 'RECITATION'
  """The token generation stopped because of potential recitation."""
  LANGUAGE = 'LANGUAGE'
  """The token generation stopped because of using an unsupported language."""
  OTHER = 'OTHER'
  """All other reasons that stopped the token generation."""
  BLOCKLIST = 'BLOCKLIST'
  """Token generation stopped because the content contains forbidden terms."""
  PROHIBITED_CONTENT = 'PROHIBITED_CONTENT'
  """Token generation stopped for potentially containing prohibited content."""
  SPII = 'SPII'
  """Token generation stopped because the content potentially contains Sensitive Personally Identifiable Information (SPII)."""
  MALFORMED_FUNCTION_CALL = 'MALFORMED_FUNCTION_CALL'
  """The function call generated by the model is invalid."""
  IMAGE_SAFETY = 'IMAGE_SAFETY'
  """Token generation stopped because generated images have safety violations."""
  UNEXPECTED_TOOL_CALL = 'UNEXPECTED_TOOL_CALL'
  """The tool call generated by the model is invalid."""

class Candidate(BaseModel):
    content: Optional[Content] = None
    finish_reason: Optional[FinishReason] = None

class UsageMetadata(BaseModel):
    cached_content_token_count: Optional[int] = None
    candidates_token_count: Optional[int] = None
    prompt_token_count: Optional[int] = None
    thoughts_token_count: Optional[int] = None
    total_token_count: Optional[int] = None

class ThinkingConfig(BaseModel):
    include_thoughts: Optional[bool] = None
    thinking_budget: Optional[int] = None

class Response(BaseModel):
    candidates: Optional[list[Candidate]] = None
    usage_metadata: Optional[UsageMetadata] = None
    parsed: Optional[Json | PydanticStructure] = None

    @property
    def text(self) -> Optional[str]:
        """Returns the concatenation of all text parts in the response."""
        if (
            not self.candidates
            or not self.candidates[0].content
            or not self.candidates[0].content.parts
        ):
            return None
        if len(self.candidates) > 1:
            logger.warning(
                f"there are {len(self.candidates)} candidates, returning text from"
                " the first candidate.Access response.candidates directly to get"
                " text from other candidates."
            )
        text = ""
        any_text_part_text = False
        for part in self.candidates[0].content.parts:
            if isinstance(part.text, str):
                if isinstance(part.thought, bool) and part.thought:
                    continue
                any_text_part_text = True
                text += part.text
        # part.text == "" is different from part.text is None
        return text if any_text_part_text else None

class Schema(BaseModel):
    example: Optional[Any] = None
    pattern: Optional[str] = None
    minimum: Optional[float] = None
    default: Optional[Any] = None
    any_of: Optional[list["Schema"]] = None
    max_length: Optional[int] = None
    title: Optional[str] = None
    min_length: Optional[int] = None
    min_properties: Optional[int] = None
    maximum: Optional[float] = None
    max_properties: Optional[int] = None
    description: Optional[str] = None
    enum: Optional[list[str]] = None
    format: Optional[str] = None
    items: Optional["Schema"] = None
    max_items: Optional[int] = None
    min_items: Optional[int] = None
    nullable: Optional[bool] = None
    properties: Optional[dict[str, "Schema"]] = None
    property_ordering: Optional[list[str]] = None
    required: Optional[list[str]] = None
    type: Optional[JsonType] = None

class FunctionDeclaration(BaseModel):
    response: Optional[Schema] = None
    description: Optional[str] = None
    name: Optional[str] = None
    parameters: Optional[Schema] = None

class Tool(BaseModel):
    function_declarations: Optional[list[FunctionDeclaration]] = None
    callable: Optional[Callable] = None

class FunctionCallingConfig(BaseModel):
    mode: Optional[Literal["AUTO", "ANY", "NONE"]] = None
    allowed_function_names: Optional[list[str]] = None

class ToolConfig(BaseModel):
    function_calling_config: Optional[FunctionCallingConfig] = None

class Model(BaseModel):
    full_model_name: str
    provider: str
    model: str
    display_name: str | None = None
