"""Shared LLM data models and provider contracts."""

from abc import ABC, abstractmethod
from typing import AsyncGenerator, BinaryIO, Literal, TypedDict

from msgspec import UNSET, convert, field
from msgspec.structs import asdict
from typing_extensions import Required, Self

from aceai.llm.errors import AceAIImplementationError
from aceai.llm.interface import MessageRole, Record, StrDict, Struct, Unset
from .tool_spec import IToolSpec


class LLMMessagePart(TypedDict, total=False):
    """Structured message chunk describing a single modality."""

    type: Required[Literal["text", "image", "audio", "file"]]
    """Semantic modality of this part."""

    data: str
    """Inline payload for text parts."""

    binary: bytes
    """Inline payload for binary modalities."""

    mime_type: str
    """Optional MIME type describing the payload."""

    url: str
    """External reference (HTTP URL, provider handle, etc.)."""

    metadata: StrDict
    """Provider- or caller-specific metadata."""


def value_to_msgpart(val: str | LLMMessagePart) -> LLMMessagePart:
    if isinstance(val, str):
        return LLMMessagePart(type="text", data=val)
    else:
        return val


type SupportedValueType = str | LLMMessagePart | list[LLMMessagePart]


def parts_factory(val: SupportedValueType):
    if isinstance(val, str):
        return [value_to_msgpart(val)]
    elif isinstance(val, dict):
        return [value_to_msgpart(val)]
    else:
        return [value_to_msgpart(v) for v in val]


class LLMUploadedAsset(Record, kw_only=True):
    """Descriptor for an uploaded asset that can be re-used by providers."""

    id: str
    """Provider-specific identifier for the uploaded asset."""

    mime_type: str
    """MIME type of the uploaded asset."""

    url: str | None = None
    """Optional URL referencing the uploaded asset."""


class LLMToolCall(Record, kw_only=True):
    """Normalized representation of an LLM-triggered tool/function call."""

    type: Literal["function_call", "mcp", "custom"] = "function_call"
    """Provider-declared call category; defaults to OpenAI-style function calls."""

    name: str
    """Registered tool name the model wants to invoke."""

    arguments: str
    """Raw JSON payload emitted by the model for the tool invocation."""

    call_id: str
    """Stable identifier used to correlate streaming deltas and tool outputs."""

    @classmethod
    def from_payload(cls, payload: StrDict) -> Self:
        return cls(
            type=payload["type"],
            name=payload["name"],
            arguments=payload["arguments"],
            call_id=payload["call_id"],
        )

    @classmethod
    def list_from_payload(cls, payload: StrDict) -> list[Self]:
        if "tool_calls" not in payload:
            return []
        return [cls.from_payload(call) for call in payload["tool_calls"]]


class LLMToolCallDelta(Record):
    """Incremental tool call update emitted during streaming."""

    id: str
    """Identifier of the tool call being incrementally populated."""

    arguments_delta: str
    """Slice of JSON arguments appended to the buffered payload."""


class LLMHostedToolSpec(Record, kw_only=True):
    """Provider-native tool executed by the model provider, not AceAI."""

    provider_name: str
    """Provider adapter that knows how to serialize this native tool."""

    native_name: str
    """Provider-native tool selector, such as OpenAI's `web_search`."""

    native_config: StrDict = field(default_factory=dict[str, object])
    """Provider-native configuration passed through by that adapter."""


type LLMToolSpec = IToolSpec | LLMHostedToolSpec


type LLMHostedToolSourceType = Literal["url", "api"]


class LLMHostedToolSource(Record, kw_only=True):
    """Provider-neutral source returned by a provider-hosted tool action."""

    url: Unset[str] = UNSET
    """Source URL reported by the provider, when this source has one."""

    title: Unset[str] = UNSET
    """Optional source title reported by the provider."""

    type: LLMHostedToolSourceType = "url"
    """Source kind reported by the provider."""


class LLMHostedToolAction(Record, kw_only=True):
    """Provider-neutral description of provider-hosted tool activity."""

    type: Literal["search", "open_page", "find_in_page"]
    """Semantic action performed by the hosted tool."""

    query: Unset[str] = UNSET
    """Single query string for search-like actions."""

    queries: Unset[list[str]] = UNSET
    """Provider-reported query list for search-like actions."""

    domains: Unset[list[str]] = UNSET
    """Optional domain restrictions used for the action."""

    sources: Unset[list[LLMHostedToolSource]] = UNSET
    """Sources returned by the action."""

    url: Unset[str] = UNSET
    """URL used by page-opening actions."""

    pattern: Unset[str] = UNSET
    """Find pattern used by page-search actions."""

    @classmethod
    def from_payload(cls, payload: StrDict) -> Self:
        """Decode a provider action payload into the normalized hosted-tool shape."""

        return convert(payload, type=cls)

    @property
    def first_query(self) -> str | None:
        """Return the primary query for display or indexing."""

        if self.query is not UNSET and self.query != "":
            return self.query
        if self.queries is not UNSET:
            for query in self.queries:
                if query != "":
                    return query
        return None

    @property
    def first_source_url(self) -> str | None:
        """Return the first source URL, when one is available."""

        if self.sources is UNSET:
            return None
        for source in self.sources:
            if source.url is not UNSET and source.url != "":
                return source.url
        return None


class LLMMessage(Struct, kw_only=True):
    """Typed chat message that backs `LLMInput.messages`."""

    role: MessageRole
    """Canonical chat role assigned to this turn."""

    content: list[LLMMessagePart] = field(default_factory=list[LLMMessagePart])
    """Structured message content; plain strings are not accepted."""

    def __ior__(self, other: Self) -> Self:
        if self.role != other.role:
            raise ValueError(f"Can't merge {other}, {self.role=}, {other.role=}")
        self.content.extend(other.content)
        return self

    def asdict(self) -> StrDict:
        # TODO: exclude UNSET values
        return asdict(self)

    @classmethod
    def build(cls, role: MessageRole, content: SupportedValueType):
        if isinstance(content, str):
            return cls(role=role, content=[value_to_msgpart(content)])
        elif isinstance(content, dict):
            return cls(role=role, content=[content])
        else:
            return cls(role=role, content=content)


class LLMToolCallMessage(LLMMessage, kw_only=True):
    """Assistant message variant that batches tool call declarations."""

    type: Literal["function_call"] = "function_call"
    """Assistant subtype emitted when tool calls are present."""

    role: Literal["system", "user", "assistant", "tool"] = "assistant"
    """Fixed assistant role for compatibility with chat providers."""

    tool_calls: list[LLMToolCall] = field(default_factory=list[LLMToolCall])
    """Always a list; empty list indicates a bug in the caller."""

    reasoning_content: str | None = None
    """Provider-specific thinking content needed to continue tool-call sub-turns."""

    def asdict(self) -> StrDict:
        res = super().asdict()
        res["tool_calls"] = [tc.asdict() for tc in self.tool_calls]
        return res

    @classmethod
    def from_content(
        cls,
        content: SupportedValueType,
        tool_calls: list[LLMToolCall],
        reasoning_content: str | None = None,
    ) -> "LLMToolCallMessage":
        content = parts_factory(content)
        return cls(
            content=content,
            tool_calls=tool_calls,
            reasoning_content=reasoning_content,
        )


class LLMToolUseMessage(LLMMessage, kw_only=True):
    """Synthetic tool response message sent back to the model."""

    role: Literal["system", "user", "assistant", "tool"] = "tool"
    """Marks the message as tool output when relayed to the model."""

    name: str
    """tool/function label associated with the message."""

    call_id: str
    """Identifier of the `LLMToolCall` whose output is being returned."""

    @classmethod
    def from_content(
        cls, content: SupportedValueType, name: str, call_id: str
    ) -> "LLMToolUseMessage":
        return cls(content=parts_factory(content), name=name, call_id=call_id)


class LLMResponseFormat(Record):
    """Response format hint shared between adapters and providers."""

    type: Literal["json_object", "text", "json_schema"] = "text"
    """Desired response style: plain text, generic JSON, or schema-constrained JSON."""

    schema: Unset[StrDict] = UNSET
    """JSON Schema payload for providers that expect a `schema` field."""

    json_schema: Unset[StrDict] = UNSET
    """Alternate schema slot for providers that expect `json_schema`."""


class ReasoningConfig(TypedDict, total=False):
    """Reasoning configuration supported by providers like OpenAI Responses."""

    effort: Literal["low", "medium", "high", "max"]
    """Requested reasoning depth; providers map onto model-specific tiers."""

    summary: Literal["auto", "none"]
    """Whether the provider should emit a sanitized reasoning summary."""

    generate_summary: Literal["auto", "none"]
    """Alias supported by some providers; kept for forward compatibility."""

    encrypted_content: bool
    """Whether encrypted reasoning blobs should be included in the payload."""


class LLMReasoningConfigSnapshot(Record, kw_only=True):
    """Provider-reported reasoning configuration applied to a response."""

    effort: (
        Literal["none", "minimal", "low", "medium", "high", "max", "xhigh"] | None
    ) = None
    """Requested reasoning depth tier (provider-specific semantics)."""

    summary: Literal["auto", "concise", "detailed"] | None = None
    """Whether the provider emitted a reasoning summary."""

    generate_summary: Literal["auto", "concise", "detailed"] | None = None
    """Alias used by some providers; preserved for compatibility."""

    encrypted_content: bool | None = None
    """Whether encrypted reasoning blobs were included by the provider."""


class LLMReasoningMeta(Record, kw_only=True):
    """Reasoning-related metadata surfaced by the provider."""

    config: LLMReasoningConfigSnapshot | None = None
    """Provider-reported reasoning config for this response."""

    tokens: int | None = None
    """Provider-reported reasoning token count (if available)."""


class LLMProviderModality(Record, kw_only=True):
    """Simple capability flags describing provider modality support."""

    text_in: bool = True
    text_out: bool = True
    image_in: bool = False
    image_out: bool = False
    audio_in: bool = False
    audio_out: bool = False
    file_in: bool = False
    file_out: bool = False


class LLMRequestMeta(TypedDict, total=False):
    """Adapter-specific metadata attached to a request."""

    model: str
    """Provider-facing model identifier; adapter defaults must supply this."""

    stream_model: str
    """Optional streaming override used when the provider exposes separate models."""

    reasoning: ReasoningConfig
    """Provider-specific reasoning options (e.g., OpenAI Responses `reasoning`)."""


class LLMInput(TypedDict, total=False):
    """Unified request bag accepted by top-level LLM APIs."""

    messages: Required[list[LLMMessage]]
    """Ordered chat history supplied to the provider."""

    temperature: float
    """Optional override for randomness; provider defaults apply when omitted."""

    top_p: float
    """Optional nucleus sampling override for providers that expose it."""

    top_k: int
    """Optional top-k setting; only a subset of providers honour it."""

    max_tokens: int
    """Optional completion cap; provider fallback limits apply when absent."""

    stop: list[str]
    """Optional stop-sequence list when the caller needs early termination."""

    tools: list[LLMToolSpec]
    """Optional local or provider-hosted tools for the model."""

    tool_choice: Literal["auto", "none"] | str
    """Optional override when the caller must pin or disable tool usage."""

    response_format: LLMResponseFormat
    """Optional hint for providers that support JSON/structured output."""

    stream: bool
    """Optional streaming flag; defaults to provider-specific behaviour."""

    metadata: LLMRequestMeta
    """Optional adapter overrides (e.g., forcing a model name)."""


class LLMUsage(Record):
    """Token accounting for a single completion."""

    input_tokens: int | None = None
    """Optional because some providers never report prompt token usage."""

    cached_input_tokens: int | None = None
    """Input tokens reported by the provider as prompt-cache hits."""

    cache_miss_input_tokens: int | None = None
    """Input tokens reported by the provider as prompt-cache misses."""

    input_cache_hit_rate: float | None = None
    """Input cache hits divided by cache-accounted input tokens."""

    output_tokens: int | None = None
    """Optional when providers omit completion-side accounting."""

    total_tokens: int | None = None
    """Optional; only present if the provider supplies the aggregate."""

    @classmethod
    def from_payload(cls, payload: StrDict) -> Self:
        input_tokens = payload["input_tokens"]
        cached_input_tokens = payload["cached_input_tokens"]
        cache_miss_input_tokens = (
            payload["cache_miss_input_tokens"]
            if "cache_miss_input_tokens" in payload
            else _cache_miss_input_tokens(input_tokens, cached_input_tokens)
        )
        input_cache_hit_rate = (
            payload["input_cache_hit_rate"]
            if "input_cache_hit_rate" in payload
            else _input_cache_hit_rate(cached_input_tokens, cache_miss_input_tokens)
        )
        return cls(
            input_tokens=input_tokens,
            cached_input_tokens=cached_input_tokens,
            cache_miss_input_tokens=cache_miss_input_tokens,
            input_cache_hit_rate=input_cache_hit_rate,
            output_tokens=payload["output_tokens"],
            total_tokens=payload["total_tokens"],
        )


def _cache_miss_input_tokens(
    input_tokens: int | None,
    cached_input_tokens: int | None,
) -> int | None:
    if input_tokens is None:
        return None
    if cached_input_tokens is None:
        return None
    return input_tokens - cached_input_tokens


def _input_cache_hit_rate(
    cached_input_tokens: int | None,
    cache_miss_input_tokens: int | None,
) -> float | None:
    if cached_input_tokens is None:
        return None
    if cache_miss_input_tokens is None:
        return None
    cache_accounted_tokens = cached_input_tokens + cache_miss_input_tokens
    if cache_accounted_tokens == 0:
        return None
    return cached_input_tokens / cache_accounted_tokens


class LLMCitationRef(Record):
    """Minimal citation/ref grounding metadata attached to a segment."""

    id: str | None = None
    """Optional because many providers don't emit stable citation IDs."""

    source: str | None = None
    """Optional label; absent when the provider only returns free text."""

    start: int | None = None
    """Optional offset because not all providers expose span data."""

    end: int | None = None
    """Optional offset because not all providers expose span data."""

    url: str | None = None
    """Optional public URL; only present when the provider shares resolved links."""

    provider_name: str | None = None
    """Optional provider identifier for this citation payload."""


class LLMSafetyAnnotation(Record):
    """Minimal safety label metadata surfaced by providers."""

    category: str
    """Provider-reported safety category (e.g., hate, violence)."""

    level: Literal["none", "low", "medium", "high"] = "none"
    """Relative severity ranking emitted by the provider."""

    blocked: bool = False
    """Indicates whether the provider blocked or truncated the content."""

    reasons: list[str] = field(default_factory=list[str])
    """Optional human-readable reasons or evidence strings."""


class LLMProviderMeta(Record):
    """Per-provider attempt metadata captured during completion."""

    provider_name: str
    """Human-friendly provider identifier (e.g., openai, anthropic)."""

    model: str
    """Model identifier used for this attempt."""

    latency_ms: float | None = None
    """Optional because some adapters cannot observe provider latency."""

    response_id: str | None = None
    """Optional provider-native response identifier for this attempt."""


class LLMGeneratedMedia(Record, kw_only=True):
    """Descriptor for media emitted by the model."""

    type: Literal["image", "audio", "file"]
    """Modality of the generated asset."""

    mime_type: str
    """MIME type of the generated asset."""

    url: str | None = None
    """Optional URL referencing the media."""

    data: bytes | None = None
    """Inline media bytes when immediately consumable."""


class LLMToolCallSegmentMeta(Record, kw_only=True):
    """Metadata for a tool call segment."""

    call_id: str
    tool_name: str | None = None
    is_delta: bool = False


class LLMImageSegmentMeta(Record, kw_only=True):
    """Metadata for an image segment."""

    item_id: str
    status: Literal["in_progress", "completed", "generating", "failed"] | None = None
    output_index: int | None = None
    partial_index: int | None = None
    sequence_number: int | None = None


class LLMReasoningSegmentMeta(Record, kw_only=True):
    """Metadata for a reasoning segment."""

    item_id: str
    kind: Literal["summary", "content"]
    index: int
    status: Literal["in_progress", "completed", "incomplete"] | None = None
    is_delta: bool = False


class LLMHostedToolSegmentMeta(Record, kw_only=True):
    """Metadata for provider-hosted tool activity."""

    provider_name: str
    tool_name: str
    item_id: str
    status: Literal["in_progress", "searching", "completed", "failed"] | None = None
    output_index: int | None = None
    sequence_number: int | None = None
    action: Unset[LLMHostedToolAction] = UNSET

    def __post_init__(self) -> None:
        if self.action is not UNSET and isinstance(self.action, dict):
            object.__setattr__(
                self,
                "action",
                LLMHostedToolAction.from_payload(self.action),
            )


type LLMSegmentMeta = (
    LLMToolCallSegmentMeta
    | LLMImageSegmentMeta
    | LLMReasoningSegmentMeta
    | LLMHostedToolSegmentMeta
)


class LLMSegment(Record):
    """Structured slice of provider output with annotations."""

    type: Literal[
        "text",
        "reasoning",
        "tool_call",
        "hosted_tool",
        "citation",
        "error",
        "image",
        "audio",
        "file",
        "other",
    ]
    """Semantic segment type to help consumers interpret content."""

    content: str
    """Raw content payload for this segment."""

    citations: list[LLMCitationRef] = field(default_factory=list[LLMCitationRef])
    """Structured citation metadata attached to this segment."""

    safety: list[LLMSafetyAnnotation] = field(default_factory=list[LLMSafetyAnnotation])
    """Structured safety annotations attached to this segment."""

    meta: LLMSegmentMeta | None = None
    """Optional typed metadata for this segment."""

    media: LLMGeneratedMedia | None = None
    """Optional generated media payload for non-text segments."""


class LLMResponse(Record):
    """Provider-agnostic completion payload."""

    id: str | None = None
    """Optional because several providers do not return stable response IDs."""

    model: str | None = None
    """Optional; adapters fall back to request defaults when providers omit."""

    text: str = ""
    """Aggregated assistant text returned by the provider."""

    tool_calls: list[LLMToolCall] = field(default_factory=list[LLMToolCall])
    """Tool/function invocations extracted from the provider output."""

    usage: Unset[LLMUsage] = UNSET
    """Unset when the provider omits token accounting details."""

    segments: list[LLMSegment] = field(default_factory=list[LLMSegment])
    """Structured segments capturing provider output with annotations."""

    provider_meta: list[LLMProviderMeta] = field(default_factory=list[LLMProviderMeta])
    """All provider attempts (including retries/failovers) for this request."""

    status: str | None = None
    """Optional provider status label (e.g., completed, incomplete)."""

    reasoning: LLMReasoningMeta | None = None
    """Optional reasoning-related metadata (usage/config), separate from segments."""

    reasoning_content: str | None = None
    """Provider-specific thinking content required by some tool-call APIs."""


class LLMStreamChunk(Record):
    """Streaming delta emitted by adapters."""

    text_delta: Unset[str] = UNSET
    """Partial text token(s) produced during streaming."""

    tool_call_delta: Unset[LLMToolCallDelta] = UNSET
    """Incremental tool call arguments emitted during streaming."""

    media: Unset[LLMGeneratedMedia] = UNSET
    """Media payload emitted during streaming (e.g., partial image)."""

    response: Unset[LLMResponse] = UNSET
    """Optional final response snapshot delivered when the stream completes."""

    error: Unset[str] = UNSET
    """Error message surfaced by the provider mid-stream."""


class LLMStreamEvent(Record):
    """Provider-agnostic streaming event with structured metadata."""

    event_type: Literal[
        "response.output_text.delta",
        "response.reasoning.delta",
        "response.function_call_arguments.delta",
        "response.completed",
        "response.error",
        "response.media",
        "response.hosted_tool",
        "response.retrying",
    ]
    """Explicit event type mirroring the provider's semantic intent."""

    text_delta: Unset[str] = UNSET
    """Partial text token(s) produced during streaming."""

    tool_call_delta: Unset[LLMToolCallDelta] = UNSET
    """Incremental tool call arguments emitted during streaming."""

    response: Unset[LLMResponse] = UNSET
    """Optional final response snapshot delivered when the stream completes."""

    error: Unset[str] = UNSET
    """Provider-reported error message surfaced mid-stream."""

    segments: list[LLMSegment] = field(default_factory=list[LLMSegment])
    """Structured segments associated with this event."""

    provider_meta: list[LLMProviderMeta] = field(default_factory=list[LLMProviderMeta])
    """Provider attempt metadata as of this event emission."""

    retry_count: int = 0
    """1-based retry counter for response.retrying events."""

    retry_max: int = 0
    """Maximum retry count for response.retrying events."""

    retry_delay_seconds: float = 0.0
    """Delay before the next provider attempt for response.retrying events."""


class LLMProviderBase(ABC):
    """Interface for LLM providers that bridge vendor SDKs to this contract."""

    def build_usage(
        self,
        *,
        input_tokens: int | None,
        cached_input_tokens: int | None,
        output_tokens: int | None,
        total_tokens: int | None,
        cache_miss_input_tokens: int | None = None,
    ) -> LLMUsage:
        cache_miss_tokens = cache_miss_input_tokens
        if cache_miss_tokens is None:
            cache_miss_tokens = _cache_miss_input_tokens(
                input_tokens,
                cached_input_tokens,
            )
        return LLMUsage(
            input_tokens=input_tokens,
            cached_input_tokens=cached_input_tokens,
            cache_miss_input_tokens=cache_miss_tokens,
            input_cache_hit_rate=self.input_cache_hit_rate(
                cached_input_tokens=cached_input_tokens,
                cache_miss_input_tokens=cache_miss_tokens,
            ),
            output_tokens=output_tokens,
            total_tokens=total_tokens,
        )

    def input_cache_hit_rate(
        self,
        *,
        cached_input_tokens: int | None,
        cache_miss_input_tokens: int | None,
    ) -> float | None:
        return _input_cache_hit_rate(cached_input_tokens, cache_miss_input_tokens)

    @abstractmethod
    async def complete(self, request: LLMInput) -> LLMResponse:
        """Complete a chat conversation with a unified input."""
        raise AceAIImplementationError(
            f"{self.__class__.__name__} must implement complete()"
        )

    @abstractmethod
    def stream(self, request: LLMInput) -> AsyncGenerator[LLMStreamEvent, None]:
        """Preferred streaming API returning structured events."""
        raise AceAIImplementationError(
            f"{self.__class__.__name__} must implement stream()"
        )

    @property
    @abstractmethod
    def modality(self) -> LLMProviderModality:
        """Return the provider's supported modalities (defaults to text-only)."""
        raise NotImplementedError

    @abstractmethod
    async def stt(
        self,
        filename: str,
        file: BinaryIO,
        *,
        model: str,
        prompt: str | None = None,
    ) -> str:
        """Speech-to-text for an audio file. Default impl not provided."""
        raise NotImplementedError
