from __future__ import annotations

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

from pydantic import BaseModel, ConfigDict, 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 BizParams(BaseModel):
    curr_base_model_name: str | None = None
    curr_model_id: str | None = None
    curr_model_ver_id: str | None = None


class MultiPart(BaseModel):
    content_type: str
    url: str | None = None
    width: int | None = None
    height: int | None = None
    text: str | None = None


class Content(BaseModel):
    multi_parts: list[MultiPart] | None = None
    parts: list[str] | None = None
    content_type: str | None = None


class Attachment(BaseModel):
    content_type: str
    url: str
    width: int | None = None
    height: int | None = None


class Metadata(BaseModel):
    attachments: list[Attachment]


class Message(BaseModel):
    role: str
    content: Content
    metadata: Metadata | None = None


class BizParam(BaseModel):
    curr_base_model_name: str | None = None
    curr_model_id: str | None = None
    curr_model_ver_id: str | None = None
    conv_type: str | None = None
    current_workspace_type: str | None = None
    video_base_model_name: str | None = None
    video_model_id: str | None = None
    video_model_ver_id: str | None = None


class SendMessageGenAgentBody(BaseModel):
    conv_id: str
    seg_id: str
    version: str | None = None
    enable_agent: bool | None = None
    start_new_topic: bool
    messages: list[Message]
    biz_param: BizParam | None = None
    safe_scene: int
    select_agent_mode: str | None = None
    is_first_msg: bool | None = None
    parent_msg_id: str | None = None


# Response models
class ContentFilterItem(APIDataModel):
    filtered: bool
    detected: bool | None = None


class ContentFilterResults(APIDataModel):
    hate: ContentFilterItem | None = None
    self_harm: ContentFilterItem | None = None
    sexual: ContentFilterItem | None = None
    violence: ContentFilterItem | None = None
    jailbreak: ContentFilterItem | None = None
    profanity: ContentFilterItem | None = None


class Delta(APIDataModel):
    content: str | None = None
    role: str | None = None


class Choice(APIDataModel):
    index: int
    delta: Delta | None = None
    finish_reason: str | None = None
    content_filter_results: ContentFilterResults | None = None


class Headers(BaseModel):  # Ignore extra fields, only extract relevant ones
    x_conv_id: str | None = None
    x_msg_id: str | None = None
    x_resp_msg_id: str | None = None


class BizInfo(APIDataModel):
    remain_free_task_generate_count: int | None = None
    show_free_task_generate_count_popup: bool | None = None


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

    event: Literal["message"]
    created: int | None = None
    model: str | None = None
    choices: list[Choice] | None = None
    headers: Headers = Headers()
    usage: dict[str, Any] | None = None
    id: str | None = None
    object: str | None = None
    system_fingerprint: str | None = None

    @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_conv = _get_header(hdrs, "x-conv-id") or _get_header(
                        hdrs, "X-Conv-Id"
                    )
                    x_msg = _get_header(hdrs, "x-msg-id") or _get_header(
                        hdrs, "X-Msg-Id"
                    )
                    x_resp = _get_header(hdrs, "x-resp-msg-id") or _get_header(
                        hdrs, "X-Resp-Msg-Id"
                    )

                    out["headers"] = {
                        "x_conv_id": x_conv,
                        "x_msg_id": x_msg,
                        "x_resp_msg_id": x_resp,
                    }

                return out
        return v

    @property
    def text(self) -> str:
        """Incremental text content of this chunk, or an empty string if absent."""
        if self.choices and self.choices[0].delta and self.choices[0].delta.content:
            return self.choices[0].delta.content
        return ""


class Status(BaseModel):
    code: int
    msg: str
    request_id: str


class AgentGenerationEnd(APIDataModel):
    event: Literal["end"] = "end"
    msg_id: str | None = None
    seg_id: str | None = None
    create_at: int | None = None
    role: str | None = None
    invoke_model: str | None = None
    content: Any | None = None
    metadata: Any | None = None
    biz_info: BizInfo | None = None
    status: Status | None = None

    @property
    def text(self) -> str:
        """Always returns an empty string — the end event carries no incremental text."""
        return ""


class AgentGenerationResult(APIDataModel):
    text: str
    value: AgentGenerationEnd


AgentGenerationStreamItem = AgentGenerationChunk | AgentGenerationEnd
