"""Prompt-first hosted model facade for lightweight external use.

This mirrors the ergonomic parts of ``ember.api.models`` without pulling in
the full Ember runtime. The facade uses the hosted Keiro gateway underneath.
"""

from __future__ import annotations

import copy
import json
from collections.abc import Iterator, Mapping, Sequence
from dataclasses import dataclass
from typing import Any

from keiro.client import Client


def _stringify_json(value: object) -> str:
    if isinstance(value, str):
        return value
    if value is None:
        return ""
    try:
        return json.dumps(value, ensure_ascii=False)
    except TypeError:
        return str(value)


def _content_to_text(content: object) -> str:
    if isinstance(content, str):
        return content
    if isinstance(content, bytes):
        return content.decode("utf-8", errors="ignore")
    if isinstance(content, Mapping):
        if content.get("text") is not None:
            return _stringify_json(content.get("text"))
        if content.get("content") is not None:
            return _content_to_text(content.get("content"))
        if content.get("output") is not None:
            return _stringify_json(content.get("output"))
        return _stringify_json(content)
    if isinstance(content, Sequence) and not isinstance(content, (str, bytes, bytearray)):
        return "".join(_content_to_text(part) for part in content)
    if content is None:
        return ""
    return str(content)


def _part_to_text(part: Mapping[str, Any]) -> str:
    if part.get("text") is not None:
        return _stringify_json(part.get("text"))
    if part.get("output") is not None:
        return _stringify_json(part.get("output"))
    if part.get("content") is not None:
        return _content_to_text(part.get("content"))
    return ""


def _extract_response_text(raw_response: Mapping[str, Any]) -> str:
    output = raw_response.get("output")
    if isinstance(output, Sequence) and not isinstance(output, (str, bytes, bytearray)):
        text_parts: list[str] = []
        for item in output:
            if not isinstance(item, Mapping):
                continue
            item_type = str(item.get("type") or "")
            if item_type == "message":
                text_parts.append(_content_to_text(item.get("content")))
            elif item_type in {"output_text", "text"}:
                text_parts.append(_part_to_text(item))
        if text_parts:
            return "".join(text_parts)

    choices = raw_response.get("choices")
    if isinstance(choices, Sequence) and not isinstance(choices, (str, bytes, bytearray)):
        for choice in choices:
            if not isinstance(choice, Mapping):
                continue
            message = choice.get("message")
            if isinstance(message, Mapping):
                text = _content_to_text(message.get("content"))
                if text:
                    return text
    return ""


def _usage_dict(raw_response: Mapping[str, Any]) -> dict[str, Any]:
    usage = raw_response.get("usage")
    if not isinstance(usage, Mapping):
        return {
            "prompt_tokens": 0,
            "completion_tokens": 0,
            "total_tokens": 0,
            "cost": 0.0,
            "actual_cost_usd": None,
            "details": None,
        }

    normalized = dict(usage)
    normalized.setdefault("prompt_tokens", 0)
    normalized.setdefault("completion_tokens", 0)
    normalized.setdefault("total_tokens", 0)
    if "cost" not in normalized:
        normalized["cost"] = normalized.get("cost_usd", 0.0) or 0.0
    normalized.setdefault("actual_cost_usd", None)
    normalized.setdefault("details", None)
    return normalized


def _normalized_context(context: Sequence[Mapping[str, Any]] | None) -> list[dict[str, Any]]:
    if context is None:
        return []
    if isinstance(context, (str, bytes, bytearray)):
        raise TypeError("context must be a sequence of message mappings")

    normalized: list[dict[str, Any]] = []
    for index, item in enumerate(context):
        if not isinstance(item, Mapping):
            raise TypeError(f"context entry {index} must be a mapping")
        normalized.append(copy.deepcopy(dict(item)))
    return normalized


def _build_prompt_messages(
    prompt: str,
    *,
    system: str | None,
    context: Sequence[Mapping[str, Any]] | None,
) -> list[dict[str, Any]]:
    messages: list[dict[str, Any]] = []
    if system is not None:
        messages.append({"role": "system", "content": str(system)})

    messages.extend(_normalized_context(context))
    if prompt:
        messages.append({"role": "user", "content": prompt})
    return messages


def _has_chat_completions_tool_format(messages: Sequence[Mapping[str, Any]]) -> bool:
    return any(
        message.get("role") == "tool"
        or message.get("tool_calls") is not None
        or message.get("tool_call_id") is not None
        for message in messages
    )


def _coerce_responses_part(
    part: Mapping[str, Any],
    *,
    default_type: str,
) -> dict[str, Any]:
    converted = copy.deepcopy(dict(part))
    part_type = str(converted.get("type") or default_type)
    if part_type == "text":
        converted["type"] = default_type
    return converted


def _coerce_responses_content(
    content: object,
    *,
    default_type: str,
) -> list[dict[str, Any]]:
    if isinstance(content, str):
        return [{"type": default_type, "text": content}]
    if isinstance(content, bytes):
        return [{"type": default_type, "text": content.decode("utf-8", errors="ignore")}]
    if isinstance(content, Mapping):
        return [_coerce_responses_part(content, default_type=default_type)]
    if isinstance(content, Sequence) and not isinstance(content, (str, bytes, bytearray)):
        parts: list[dict[str, Any]] = []
        for item in content:
            if isinstance(item, Mapping):
                parts.append(_coerce_responses_part(item, default_type=default_type))
            else:
                parts.append({"type": default_type, "text": _stringify_json(item)})
        return parts
    if content is None:
        return [{"type": default_type, "text": ""}]
    return [{"type": default_type, "text": str(content)}]


def _as_responses_turns(messages: Sequence[Mapping[str, Any]]) -> list[dict[str, Any]]:
    items: list[dict[str, Any]] = []

    for index, message in enumerate(messages):
        if not isinstance(message, Mapping):
            raise TypeError(f"messages entry {index} must be a mapping")

        role = str(message.get("role") or "")
        content = message.get("content", "")

        if role == "tool":
            items.append(
                {
                    "type": "function_call_output",
                    "call_id": str(message.get("tool_call_id", "")),
                    "output": _stringify_json(content),
                }
            )
            continue

        if role == "assistant":
            if content not in ("", None) and content != []:
                items.append(
                    {
                        "type": "message",
                        "role": "assistant",
                        "content": _coerce_responses_content(
                            content,
                            default_type="output_text",
                        ),
                    }
                )

            tool_calls = message.get("tool_calls")
            if isinstance(tool_calls, Sequence) and not isinstance(
                tool_calls, (str, bytes, bytearray)
            ):
                for tool_call in tool_calls:
                    if not isinstance(tool_call, Mapping):
                        continue
                    function = tool_call.get("function")
                    function_name = ""
                    function_arguments: object = "{}"
                    if isinstance(function, Mapping):
                        function_name = str(function.get("name", ""))
                        function_arguments = function.get("arguments", "{}")
                    items.append(
                        {
                            "type": "function_call",
                            "call_id": str(tool_call.get("id", "")),
                            "name": function_name,
                            "arguments": _stringify_json(function_arguments),
                        }
                    )
            continue

        items.append(
            {
                "type": "message",
                "role": role,
                "content": _coerce_responses_content(content, default_type="input_text"),
            }
        )

    return items


def _extract_stream_text(event: Mapping[str, Any]) -> str:
    event_type = event.get("event") or event.get("type")
    if event_type == "delta":
        return _content_to_text(event.get("text"))
    if event_type == "response.output_text.delta":
        return _content_to_text(event.get("delta"))

    choices = event.get("choices")
    if isinstance(choices, Sequence) and not isinstance(choices, (str, bytes, bytearray)):
        for choice in choices:
            if not isinstance(choice, Mapping):
                continue
            delta = choice.get("delta")
            if isinstance(delta, Mapping):
                text = _content_to_text(delta.get("content"))
                if text:
                    return text
    return ""


class Response:
    """Structured wrapper around hosted gateway responses."""

    def __init__(self, raw_response: Mapping[str, Any]):
        self._raw = dict(raw_response)

    @property
    def text(self) -> str:
        return _extract_response_text(self._raw)

    @property
    def usage(self) -> dict[str, Any]:
        return _usage_dict(self._raw)

    @property
    def model_id(self) -> str | None:
        model_id = self._raw.get("model")
        return str(model_id) if isinstance(model_id, str) and model_id else None

    @property
    def metadata(self) -> dict[str, Any]:
        metadata: dict[str, Any] = {}
        for key in ("id", "object", "created"):
            value = self._raw.get(key)
            if value is not None:
                metadata[key] = value
        if self.model_id is not None:
            metadata["model_id"] = self.model_id
        return metadata

    @property
    def output(self) -> list[dict[str, Any]]:
        output = self._raw.get("output")
        if isinstance(output, list):
            return [dict(item) for item in output if isinstance(item, Mapping)]
        return []

    @property
    def raw(self) -> dict[str, Any]:
        return dict(self._raw)

    def __str__(self) -> str:
        return self.text

    def __repr__(self) -> str:
        preview = self.text[:50]
        return f"Response(text={preview!r}, tokens={self.usage['total_tokens']})"


@dataclass(frozen=True)
class ModelInvocation:
    """Composite result exposing text alongside the structured response."""

    text: str
    response: Response

    @property
    def usage(self) -> dict[str, Any]:
        return self.response.usage

    @property
    def model_id(self) -> str | None:
        return self.response.model_id

    @property
    def metadata(self) -> dict[str, Any]:
        return self.response.metadata


class ModelResult(str):
    """String-like result that retains the structured response object."""

    __slots__ = ("_response",)
    _response: Response

    def __new__(cls, response: Response) -> "ModelResult":
        instance = super().__new__(cls, response.text)
        instance._response = response
        return instance

    @property
    def text(self) -> str:
        return str(self)

    @property
    def response(self) -> Response:
        return self._response

    @property
    def usage(self) -> dict[str, Any]:
        return self._response.usage

    @property
    def model_id(self) -> str | None:
        return self._response.model_id

    @property
    def metadata(self) -> dict[str, Any]:
        return self._response.metadata

    @property
    def output(self) -> list[dict[str, Any]]:
        return self._response.output

    @property
    def raw(self) -> dict[str, Any]:
        return self._response.raw

    def __repr__(self) -> str:
        return f"ModelResult(text={str(self)!r})"


class ModelBinding:
    """Reusable callable with a fixed model id and default parameters."""

    def __init__(self, model_id: str, api: ModelsAPI, **params: Any) -> None:
        self.model_id = model_id
        self._api = api
        self._params = dict(params)

    def __call__(self, prompt: str, **override_params: Any) -> str:
        return str(self._api(self.model_id, prompt, **{**self._params, **override_params}))

    def response(self, prompt: str, **override_params: Any) -> Response:
        return self._api.response(self.model_id, prompt, **{**self._params, **override_params})

    def with_metadata(self, prompt: str, **override_params: Any) -> ModelInvocation:
        return self._api.with_metadata(
            self.model_id,
            prompt,
            **{**self._params, **override_params},
        )

    def stream(self, prompt: str, **override_params: Any) -> Iterator[str]:
        return self._api.stream(self.model_id, prompt, **{**self._params, **override_params})

    def __repr__(self) -> str:
        return f"ModelBinding(model_id={self.model_id!r}, params={self._params!r})"


class ModelsAPI:
    """Hosted facade shaped like ``ember.api.models`` for quick testing.

    Caches a single :class:`Client` instance for HTTP connection reuse.
    Call :meth:`close` (or use as a context manager) to release the
    underlying connection pool explicitly.
    """

    def __init__(
        self,
        *,
        api_key: str | None = None,
        base_url: str | None = None,
        timeout: float | None = None,
        retries: int = 3,
        pool_maxsize: int = 100,
    ) -> None:
        self._api_key = api_key
        self._base_url = base_url
        self._timeout = timeout
        self._retries = retries
        self._pool_maxsize = pool_maxsize
        self._cached_client: Client | None = None

    def _client(self) -> Client:
        if self._cached_client is not None:
            return self._cached_client
        client_kwargs: dict[str, Any] = {
            "api_key": self._api_key,
            "base_url": self._base_url,
            "retries": self._retries,
            "pool_maxsize": self._pool_maxsize,
        }
        if self._timeout is not None:
            client_kwargs["timeout"] = self._timeout
        self._cached_client = Client(**client_kwargs)
        return self._cached_client

    def close(self) -> None:
        """Close the underlying HTTP client and release connections."""
        if self._cached_client is not None:
            self._cached_client.close()
            self._cached_client = None

    def __enter__(self) -> "ModelsAPI":
        return self

    def __exit__(self, *_args: object) -> None:
        self.close()

    def _prepare_messages(
        self, prompt: str, params: dict[str, Any],
    ) -> list[dict[str, Any]]:
        """Pop system/context from *params* and build chat messages."""
        system = params.pop("system", None)
        if system is not None and not isinstance(system, str):
            system = str(system)
        context = params.pop("context", None)
        if context is not None and (
            isinstance(context, (str, bytes, bytearray))
            or not isinstance(context, Sequence)
        ):
            raise TypeError("context must be a sequence when supplied")
        return _build_prompt_messages(prompt, system=system, context=context)

    def _chat_response(self, model: str, prompt: str, **params: Any) -> Response:
        stream_requested = bool(params.pop("stream", False))
        if stream_requested:
            raise ValueError("Use ModelsAPI.stream() for streaming invocations")

        if "messages" in params:
            raise ValueError(
                "Use Client.chat() or ModelsAPI.responses(messages=...) for raw messages"
            )

        messages = self._prepare_messages(prompt, params)
        return Response(self._client().chat(messages, model=model, **params))

    def _build_responses_payload(
        self,
        payload: Mapping[str, Any] | None,
        *,
        input: object | None,
        instructions: object | None,
        messages: Sequence[Mapping[str, Any]] | None,
        params: dict[str, Any],
    ) -> dict[str, Any]:
        request_payload: dict[str, Any] = {}
        if payload is not None:
            request_payload.update(copy.deepcopy(dict(payload)))

        if input is not None:
            request_payload["input"] = copy.deepcopy(input)

        if instructions is not None:
            request_payload["instructions"] = copy.deepcopy(instructions)

        if messages is not None:
            for index, message in enumerate(messages):
                if not isinstance(message, Mapping):
                    raise TypeError(f"messages entry {index} must be a mapping")

            if _has_chat_completions_tool_format(messages):
                # _as_responses_turns builds fresh dicts; no pre-copy needed.
                request_payload["input"] = _as_responses_turns(messages)
            else:
                request_payload["messages"] = [
                    copy.deepcopy(dict(m)) for m in messages
                ]

        prompt = params.pop("prompt", None)
        if prompt is not None:
            if "input" in request_payload and request_payload["input"] is not None:
                raise ValueError("Cannot supply both 'prompt' and 'input'")
            request_payload["input"] = prompt

        context = params.pop("context", None)
        if context is not None and "messages" not in request_payload:
            if isinstance(context, (str, bytes, bytearray)) or not isinstance(context, Sequence):
                raise TypeError("context must be a sequence when supplied")
            request_payload["messages"] = _normalized_context(context)

        max_tokens = params.pop("max_tokens", None)
        if max_tokens is not None and "max_output_tokens" not in request_payload:
            request_payload["max_output_tokens"] = max_tokens

        params.pop("stream", None)

        for key, value in list(params.items()):
            if value is None:
                continue
            request_payload[key] = copy.deepcopy(value)

        return {key: value for key, value in request_payload.items() if value is not None}

    def __call__(self, model: str, prompt: str, **params: Any) -> ModelResult:
        return ModelResult(self._chat_response(model, prompt, **params))

    def response(self, model: str, prompt: str, **params: Any) -> Response:
        return self._chat_response(model, prompt, **params)

    def with_metadata(self, model: str, prompt: str, **params: Any) -> ModelInvocation:
        response = self._chat_response(model, prompt, **params)
        return ModelInvocation(text=response.text, response=response)

    def instance(self, model: str, **params: Any) -> ModelBinding:
        return ModelBinding(model, self, **params)

    def stream(self, model: str, prompt: str, **params: Any) -> Iterator[str]:
        params.pop("stream", None)
        messages = self._prepare_messages(prompt, params)

        def _iterator() -> Iterator[str]:
            for event in self._client().chat_stream(messages, model=model, **params):
                text = _extract_stream_text(event)
                if text:
                    yield text

        return _iterator()

    def responses(
        self,
        model: str,
        *,
        payload: Mapping[str, Any] | None = None,
        input: object | None = None,
        instructions: object | None = None,
        messages: Sequence[Mapping[str, Any]] | None = None,
        **params: Any,
    ) -> Response:
        stream_requested = bool(params.pop("stream", False))
        if stream_requested:
            raise ValueError("Use ModelsAPI.responses_stream() for streaming invocations")
        timeout = params.pop("timeout", None)

        request_payload = self._build_responses_payload(
            payload,
            input=input,
            instructions=instructions,
            messages=messages,
            params=dict(params),
        )
        request_payload["model"] = model
        return Response(self._client().responses(request_payload, timeout=timeout))

    def responses_stream(
        self,
        model: str,
        *,
        payload: Mapping[str, Any] | None = None,
        input: object | None = None,
        instructions: object | None = None,
        messages: Sequence[Mapping[str, Any]] | None = None,
        **params: Any,
    ) -> Iterator[str]:
        params.pop("stream", None)
        timeout = params.pop("timeout", None)
        request_payload = self._build_responses_payload(
            payload,
            input=input,
            instructions=instructions,
            messages=messages,
            params=dict(params),
        )
        request_payload["model"] = model

        def _iterator() -> Iterator[str]:
            for event in self._client().responses_stream(
                request_payload, timeout=timeout,
            ):
                text = _extract_stream_text(event)
                if text:
                    yield text

        return _iterator()

    def list(self, provider: str | None = None, *, status: str | None = None) -> list[str]:
        return sorted(self.discover(provider=provider, status=status))

    def discover(
        self,
        provider: str | None = None,
        *,
        status: str | None = None,
    ) -> dict[str, dict[str, Any]]:
        payload = self._client().get_models(provider=provider, status=status)

        data = payload.get("data")
        if not isinstance(data, list):
            raise ValueError("Gateway /v1/models response did not contain a 'data' list")

        discovered: dict[str, dict[str, Any]] = {}
        for entry in data:
            if not isinstance(entry, Mapping):
                continue
            model_id = entry.get("id")
            if not isinstance(model_id, str) or not model_id:
                continue
            discovered[model_id] = dict(entry)
        return discovered


__all__ = [
    "ModelBinding",
    "ModelInvocation",
    "ModelResult",
    "ModelsAPI",
    "Response",
]
