import json
import asyncio
from typing import AsyncGenerator
import httpx
from langchain_core.messages import (
    AIMessage,
    AIMessageChunk,
    BaseMessage,
    HumanMessage,
    SystemMessage,
    ToolMessage,
)

from ...config import settings

KIA_BASE_URL = "https://api.kie.ai"

# Google Search tool shorthand — mutually exclusive with function calling
KIA_GOOGLE_SEARCH_TOOL = {"type": "function", "function": {"name": "googleSearch"}}


def _to_kia_messages(messages: list[BaseMessage]) -> list[dict]:
    """
    Convert LangChain messages to KIA message format.

    Supported content formats per KIA docs:
      - Plain text   → str
      - Text-only list → joined into plain string
      - Multimodal (image/video/audio/pdf) → list with unified format:
          [{"type": "image_url", "image_url": {"url": "..."}}]
        All media types (image, video, audio, PDF) use the same image_url structure.
    """
    result = []
    for msg in messages:
        if isinstance(msg, SystemMessage):
            result.append({"role": "system", "content": _normalize_content(msg.content)})
        elif isinstance(msg, HumanMessage):
            result.append({"role": "user", "content": _normalize_content(msg.content)})
        elif isinstance(msg, AIMessage):
            result.append({"role": "assistant", "content": _normalize_content(msg.content)})
        elif isinstance(msg, ToolMessage):
            result.append({
                "role": "tool",
                "content": msg.content,
                "tool_call_id": msg.tool_call_id,
            })
        else:
            result.append({"role": "user", "content": str(msg.content)})
    return result


def _normalize_content(content) -> str | list:
    """
    - str                           → returned as-is
    - list of text-only dicts       → joined into plain string
    - list with image_url/media     → returned as-is (multimodal)
    """
    if isinstance(content, str):
        return content

    if isinstance(content, list):
        is_text_only = all(
            isinstance(item, dict) and item.get("type") == "text"
            for item in content
        )
        if is_text_only:
            return " ".join(item.get("text", "") for item in content)
        return content  # multimodal — keep as list

    return str(content)


class KiaAIService:
    """
    Calls KIA AI API (https://api.kie.ai) directly via httpx SSE streaming.
    """

    # ------------------------------------------------------------------
    # Internal helpers
    # ------------------------------------------------------------------

    def _build_payload(
        self,
        model: str,
        messages: list[BaseMessage],
        reasoning_effort: str | None,
        include_thoughts: bool,
        google_search: bool,
        tools: list[dict] | None,
        response_format: dict | None,
    ) -> tuple[str, dict]:
        """Build the KIA API URL and request payload."""
        model_name = model.removeprefix("kia/")
        url = f"{KIA_BASE_URL}/{model_name}/v1/chat/completions"
        kia_messages = _to_kia_messages(messages)

        print(f"[KiaAI] POST {url}")
        print(f"[KiaAI] msgs={len(kia_messages)} | reasoning={reasoning_effort} | thoughts={include_thoughts} | gsearch={google_search}")
        for i, m in enumerate(kia_messages):
            preview = m["content"][:120] if isinstance(m["content"], str) else str(m["content"])[:120]
            print(f"[KiaAI] msg[{i}] role={m['role']} | {preview!r}")

        payload: dict = {"messages": kia_messages, "stream": True}
        if reasoning_effort:
            payload["reasoning_effort"] = reasoning_effort
        if include_thoughts:
            payload["include_thoughts"] = True
        if google_search:
            payload["tools"] = [KIA_GOOGLE_SEARCH_TOOL]
        elif tools:
            payload["tools"] = tools
        elif response_format:
            payload["response_format"] = response_format

        print(f"[KiaAI] payload keys: {[k for k in payload if k != 'messages']}")
        return url, payload

    async def _iter_sse(
        self, url: str, payload: dict
    ) -> AsyncGenerator[AIMessageChunk, None]:
        """
        Core SSE parser — yields one AIMessageChunk per content delta token.
        Final chunk: AIMessageChunk(content="", usage_metadata={...}).
        Raises RuntimeError on KIA server errors (triggers retry in stream()).

        Uses aiter_bytes(chunk_size=32) instead of aiter_lines() so every
        byte received from the network is processed immediately without
        waiting for httpx's internal 64 KB read buffer to fill up.
        """
        usage: dict = {}

        async with httpx.AsyncClient(timeout=180.0) as client:
            async with client.stream(
                "POST",
                url,
                content=json.dumps(payload),
                headers={
                    "Content-Type": "application/json",
                    "Authorization": f"Bearer {settings.kia_api_key}",
                },
            ) as resp:
                print(f"[KiaAI] HTTP {resp.status_code} | {url}")
                resp.raise_for_status()

                line_count = 0
                buf = ""
                done = False

                async for raw in resp.aiter_bytes(chunk_size=32):
                    buf += raw.decode("utf-8", errors="replace")

                    while "\n" in buf:
                        line, buf = buf.split("\n", 1)
                        line = line.rstrip("\r")
                        line_count += 1
                        if line_count <= 3:
                            print(f"[KiaAI] raw[{line_count}]: {line!r}")

                        if not line:
                            continue

                        # Non-data lines may carry KIA error JSON
                        if not line.startswith("data:"):
                            try:
                                err = json.loads(line)
                                code = err.get("code")
                                if code and int(code) >= 400:
                                    msg = err.get("msg", "Unknown error")
                                    print(f"[KiaAI] Error line: code={code} msg={msg!r}")
                                    raise RuntimeError(f"KIA error {code}: {msg}")
                            except (json.JSONDecodeError, ValueError):
                                pass
                            print(f"[KiaAI] non-data line: {line!r}")
                            continue

                        raw_data = line[6:] if line.startswith("data: ") else line[5:]

                        if raw_data.strip() == "[DONE]":
                            print("[KiaAI] [DONE]")
                            done = True
                            break

                        try:
                            data = json.loads(raw_data)
                        except json.JSONDecodeError:
                            print(f"[KiaAI] JSON decode error: {raw_data!r}")
                            continue

                        # KIA server error embedded in stream body
                        if data.get("code") and int(data["code"]) >= 500:
                            msg = data.get("msg", "Unknown error")
                            print(f"[KiaAI] Error chunk: {data['code']} {msg!r}")
                            raise RuntimeError(f"KIA error {data['code']}: {msg}")

                        # Yield content tokens
                        for choice in (data.get("choices") or []):
                            delta = choice.get("delta") or {}
                            token = delta.get("content") or ""
                            if token:
                                yield AIMessageChunk(content=token)

                        # Capture usage from final summary chunk
                        if data.get("usage"):
                            usage = data["usage"]

                    if done:
                        break

        # Final chunk carries usage metadata
        print(f"[KiaAI] Stream done | total_lines={line_count} | usage={usage}")
        yield AIMessageChunk(
            content="",
            usage_metadata={
                "input_tokens": usage.get("prompt_tokens", 0),
                "output_tokens": usage.get("completion_tokens", 0),
                "total_tokens": usage.get("total_tokens", 0),
            },
        )

    # ------------------------------------------------------------------
    # Public API
    # ------------------------------------------------------------------

    async def stream(
        self,
        model: str,
        messages: list[BaseMessage],
        reasoning_effort: str | None = None,
        include_thoughts: bool = False,
        google_search: bool = False,
        tools: list[dict] | None = None,
        response_format: dict | None = None,
        max_retries: int = 3,
        retry_delay: float = 2.0,
    ) -> AsyncGenerator[AIMessageChunk, None]:
        """
        Token-by-token async generator.
        Yields AIMessageChunk(content=token) per SSE delta.
        Last chunk: AIMessageChunk(content="", usage_metadata={...}).
        Retries on server errors only if no tokens were yielded yet.
        """
        url, payload = self._build_payload(
            model, messages, reasoning_effort, include_thoughts,
            google_search, tools, response_format,
        )

        last_error: Exception | None = None
        for attempt in range(1, max_retries + 1):
            if attempt > 1:
                print(f"[KiaAI] Stream retry {attempt}/{max_retries} in {retry_delay}s ...")
                await asyncio.sleep(retry_delay)

            yielded_any = False
            try:
                async for chunk in self._iter_sse(url, payload):
                    yielded_any = True
                    yield chunk
                return  # success
            except RuntimeError as e:
                if yielded_any:
                    raise  # can't retry after content was already sent
                last_error = e
                print(f"[KiaAI] Attempt {attempt} failed (no content yet): {e}")

        raise last_error

    async def invoke(
        self,
        model: str,
        messages: list[BaseMessage],
        reasoning_effort: str | None = None,
        include_thoughts: bool = False,
        google_search: bool = False,
        tools: list[dict] | None = None,
        response_format: dict | None = None,
        max_retries: int = 3,
        retry_delay: float = 2.0,
    ) -> AIMessage:
        """Collects all streaming chunks into a single AIMessage."""
        full_content = ""
        usage: dict = {}
        async for chunk in self.stream(
            model=model,
            messages=messages,
            reasoning_effort=reasoning_effort,
            include_thoughts=include_thoughts,
            google_search=google_search,
            tools=tools,
            response_format=response_format,
            max_retries=max_retries,
            retry_delay=retry_delay,
        ):
            full_content += chunk.content or ""
            if chunk.usage_metadata:
                usage = chunk.usage_metadata

        print(f"[KiaAI] invoke() done | content_len={len(full_content)} | usage={usage}")
        ai_msg = AIMessage(content=full_content)
        ai_msg.usage_metadata = usage
        return ai_msg


# Compatibility stub — KIA models bypass LangChain and use KiaAIService directly
def get_kia_model(model: str):
    return None