from __future__ import annotations

from typing import Any, Dict, List, Optional
from ...core.types import ChatResponse, Message

from ...client import AielClient
from .remote import RemoteOpenAI
from .types import OpenAIChatCompletionsCreateRequest, OpenAIEmbeddingsCreateRequest

Json = Dict[str, Any]


def _to_float_list(value: Any) -> List[float]:
    if not isinstance(value, list):
        return []
    out: List[float] = []
    for item in value:
        try:
            out.append(float(item))
        except (TypeError, ValueError):
            return []
    return out


def _extract_embeddings(raw: Json) -> List[List[float]]:
    # OpenAI native shape: {"data": [{"embedding": [...]}, ...]}
    data = raw.get("data")
    if isinstance(data, list):
        out: List[List[float]] = []
        for item in data:
            if isinstance(item, dict):
                emb = _to_float_list(item.get("embedding"))
                if emb:
                    out.append(emb)
        if out:
            return out

    # Wrapped variants: {"data": {"data": [...]}} or {"result": {"data": [...]}}
    for key in ("data", "result"):
        nested = raw.get(key)
        if isinstance(nested, dict):
            nested_data = nested.get("data")
            if isinstance(nested_data, list):
                out = []
                for item in nested_data:
                    if isinstance(item, dict):
                        emb = _to_float_list(item.get("embedding"))
                        if emb:
                            out.append(emb)
                if out:
                    return out

    # Connector-friendly variants:
    # {"embedding":[...]} for single, {"embeddings":[...]} or {"vectors":[...]} for multi
    single = _to_float_list(raw.get("embedding"))
    if single:
        return [single]

    for key in ("embeddings", "vectors"):
        value = raw.get(key)
        if isinstance(value, list):
            out = []
            for item in value:
                if isinstance(item, dict):
                    emb = _to_float_list(item.get("embedding"))
                else:
                    emb = _to_float_list(item)
                if emb:
                    out.append(emb)
            if out:
                return out

    return []


class ChatOpenAI:
    """
    LangChain-style ChatOpenAI facade backed by AIEL RemoteOpenAI integration.
    """

    def __init__(
        self,
        *,
        client: AielClient,
        workspace_id: str,
        project_id: str,
        connection_id: str,
        model: str,
        temperature: Optional[float] = None,
        max_tokens: Optional[int] = None,
        **kwargs: Any,
    ) -> None:
        self._remote = RemoteOpenAI(
            client=client,
            workspace_id=workspace_id,
            project_id=project_id,
            connection_id=connection_id,
        )
        self.model = model
        self.temperature = temperature
        self.max_tokens = max_tokens
        self._extra_defaults: Json = dict(kwargs)

    def _to_openai_message(self, message: Message) -> Json:
        out: Json = {
            "role": message.role,
            "content": message.content,
        }
        if message.name:
            out["name"] = message.name
        if message.tool_call_id:
            out["tool_call_id"] = message.tool_call_id
        return out

    def invoke(
        self,
        messages: List[Json],
        *,
        request_id: str | None = None,
        **kwargs: Any,
    ) -> Json:
        extra = dict(self._extra_defaults)
        extra.update(kwargs)
        req = OpenAIChatCompletionsCreateRequest(
            model=self.model,
            temperature=self.temperature,
            max_tokens=self.max_tokens,
            messages=messages,
            extra=extra,
        )
        return self._remote.chat_completions_create(req, request_id=request_id).raw

    def invoke_core(
        self,
        messages: List[Message],
        **kwargs: Any,
    ) -> ChatResponse:
        return self.invoke_messages(messages, **kwargs)

    def _message_content_to_text(self, content: Any) -> str:
        if isinstance(content, str):
            return content
        if isinstance(content, list):
            parts: List[str] = []
            for item in content:
                if isinstance(item, dict):
                    # OpenAI multi-part message content commonly uses {"type":"text","text":"..."}.
                    if isinstance(item.get("text"), str):
                        parts.append(item["text"])
                    elif isinstance(item.get("content"), str):
                        parts.append(item["content"])
            return "".join(parts)
        return ""

    def _to_chat_response(self, raw: Json) -> ChatResponse:
        choices = raw.get("choices")
        if isinstance(choices, list) and choices and isinstance(choices[0], dict):
            first = choices[0]
            msg = first.get("message")
            if isinstance(msg, dict):
                metadata: Json = {"finish_reason": first.get("finish_reason")}
                if isinstance(msg.get("tool_calls"), list):
                    metadata["tool_calls"] = msg["tool_calls"]
                return ChatResponse(
                    message=Message(
                        role="assistant",
                        content=self._message_content_to_text(msg.get("content", "")),
                        metadata=metadata,
                    ),
                    raw=raw,
                )

        return ChatResponse(
            message=Message(role="assistant", content=""),
            raw=raw,
        )

    def invoke_messages(
        self,
        messages: List[Message],
        *,
        request_id: str | None = None,
        **kwargs: Any,
    ) -> ChatResponse:
        raw = self.invoke(
            [self._to_openai_message(m) for m in messages],
            request_id=request_id,
            **kwargs,
        )
        return self._to_chat_response(raw)


class OpenAIEmbeddings:
    """
    LangChain-style OpenAIEmbeddings facade backed by AIEL RemoteOpenAI integration.
    """

    def __init__(
        self,
        *,
        client: AielClient,
        workspace_id: str,
        project_id: str,
        connection_id: str,
        model: str = "text-embedding-3-small",
        dimensions: Optional[int] = None,
        **kwargs: Any,
    ) -> None:
        self._remote = RemoteOpenAI(
            client=client,
            workspace_id=workspace_id,
            project_id=project_id,
            connection_id=connection_id,
        )
        self.model = model
        self.dimensions = dimensions
        self._extra_defaults: Json = dict(kwargs)

    def invoke(
        self,
        input: Any,
        *,
        request_id: str | None = None,
        **kwargs: Any,
    ) -> Json:
        extra = dict(self._extra_defaults)
        extra.update(kwargs)
        req = OpenAIEmbeddingsCreateRequest(
            input=input,
            model=self.model,
            dimensions=self.dimensions,
            extra=extra,
        )
        return self._remote.embeddings_create(req, request_id=request_id).raw

    def embed_query(self, text: str, *, request_id: str | None = None, **kwargs: Any) -> List[float]:
        raw = self.invoke(text, request_id=request_id, **kwargs)
        vectors = _extract_embeddings(raw)
        return vectors[0] if vectors else []

    def embed_documents(
        self,
        texts: List[str],
        *,
        request_id: str | None = None,
        **kwargs: Any,
    ) -> List[List[float]]:
        raw = self.invoke(texts, request_id=request_id, **kwargs)
        return _extract_embeddings(raw)
