from __future__ import annotations

from collections.abc import Mapping
from typing import Any, Literal, cast

from pydantic import BaseModel, ConfigDict, Field, model_validator

from ..base import APIDataModel


# Helper functions
def _get_header(h: Mapping[str, Any], key: str) -> str | None:
    v = h.get(key)
    if v is None:
        return None
    return v if isinstance(v, str) else str(v)


# Request models
class HistoryCharacterMessagesBody(BaseModel):
    session_id: str
    offset_to: int
    page_size: int
    send_ss: int | None = None
    ss: int | None = None


# Response models
class Message(BaseModel):
    role: str
    content: str
    split_msg_ids: list[str] | None = None


class Item(APIDataModel):
    id: str
    num_of_token: int | None = None
    content: str
    original_content: str
    json_content: Any
    create_at: int
    role: int
    support_voice: int
    voice_duration: int
    think_duration: int
    status: int
    active_idx: int
    gpt_model: str
    sub_msg: list[str] | None = None
    file_info: Any
    is_active_upload: int
    search_results: Any
    on_line_search: int
    scene: str
    is_like: int
    indistinct: int
    handle_button_text: str
    handle_msg_type: int
    is_first_msg: int | None = None


class HistoryCharacterMessages(APIDataModel):
    items: list[Item]
    offset: int
    has_more: bool


class Choice(APIDataModel):
    index: int
    finish_reason: str
    message: Message


class Usage(APIDataModel):
    prompt_tokens: int
    completion_tokens: int
    search_input_tokens: int
    search_output_tokens: int
    total_tokens: int
    total_used: int
    cost: int


class Headers(BaseModel):  # Ignore extra fields, only care about msg ids
    x_msg_id: str | None = None
    x_response_msg_id: str | None = None


class _CharacterChatCompletionBase(APIDataModel):
    model_config = ConfigDict(from_attributes=True)

    id: str
    object: Literal["chat.completion.chunk"]
    created: int
    model: str
    session_id: str
    msg_id: str
    choices: list[Choice]
    usage: Usage
    headers: Headers = Headers()

    @model_validator(mode="before")
    @classmethod
    def _unwrap_env(cls, v: Any) -> Any:
        if hasattr(v, "event") and hasattr(v, "data"):
            data_any = v.data
            if isinstance(data_any, Mapping):
                out: dict[str, Any] = dict(cast("Mapping[str, Any]", data_any))
                out["event"] = cast("str", v.event)

                hdrs_any = getattr(v, "headers", None)
                if isinstance(hdrs_any, Mapping):
                    hdrs = cast("Mapping[str, Any]", hdrs_any)

                    x_msg = _get_header(hdrs, "x-msg-id") or _get_header(
                        hdrs, "X-Msg-Id"
                    )
                    x_resp = _get_header(hdrs, "x-response-msg-id") or _get_header(
                        hdrs, "X-Response-Msg-Id"
                    )

                    out["headers"] = {
                        "x_msg_id": x_msg,
                        "x_response_msg_id": x_resp,
                    }

                return out
        return v

    @property
    def content(self) -> str:
        """Text content of the first choice in this chunk or completion."""
        if self.choices and self.choices[0].message and self.choices[0].message.content:
            return self.choices[0].message.content
        return ""

    @property
    def my_msg_id(self) -> str | None:
        """ID of the user's message, from the ``x-msg-id`` SSE response header."""
        return self.headers.x_msg_id

    @property
    def response_msg_id(self) -> str | None:
        """ID of the AI's response message, from the ``x-response-msg-id`` SSE header.

        Pass this as ``parent_msg_id`` in the next :meth:`stream` /
        :meth:`send` call to maintain conversation threading.
        """
        return self.headers.x_response_msg_id


class CharacterChatCompletionChunk(_CharacterChatCompletionBase):
    event: Literal["chunk"]


class CharacterChatCompletionEnd(_CharacterChatCompletionBase):
    event: Literal["end"]


class CharacterChatVoiceChunk(APIDataModel):
    model_config = ConfigDict(from_attributes=True, populate_by_name=True)

    event: Literal["chunk"]
    audio_data: str = Field(alias="audioData")
    audio_id: str

    @model_validator(mode="before")
    @classmethod
    def _unwrap_env(cls, v: Any) -> Any:
        if hasattr(v, "event") and hasattr(v, "data"):
            data_any = v.data

            if isinstance(data_any, Mapping):
                data = cast("Mapping[str, Any]", data_any)
                out: dict[str, Any] = dict(data)
                out["event"] = cast("str", v.event)
                return out

        return v


class Voice(APIDataModel):
    id: str
    create_at: int | None = None
    update_at: int | None = None
    account_id: str | None = None
    voice_text: str | None = None
    channel: str | None = None
    lang: str | None = None
    is_chargeable: bool | None = None
    temp_hashrate: int | None = None
    after_hashrate: int | None = None
    after_temp_hashrate: int | None = None
    bytes: int | None = None
    duration: int | None = None
    status: int | None = None
    msg_id: str | None = None
    deadline: int | None = None


class CharacterChatVoiceEnd(APIDataModel):
    model_config = ConfigDict(from_attributes=True)
    event: Literal["end"]
    voice: Voice

    @model_validator(mode="before")
    @classmethod
    def _unwrap_env(cls, v: Any) -> Any:
        if hasattr(v, "event") and hasattr(v, "data"):
            data_any = v.data

            if isinstance(data_any, Mapping):
                data = cast("Mapping[str, Any]", data_any)
                out: dict[str, Any] = dict(data)
                out["event"] = cast("str", v.event)
                return out

        return v


class CharacterChatCompletionResult(APIDataModel):
    text: str
    value: CharacterChatCompletionEnd


StreamItem = CharacterChatCompletionChunk | CharacterChatCompletionEnd
VoiceStreamItem = CharacterChatVoiceChunk | CharacterChatVoiceEnd
