# ui/agent/messages.py
from __future__ import annotations

import base64
import json
import mimetypes
import re
import uuid
from pathlib import Path
from typing import Any

from .config import TOOLS


TOOL_NAMES = {tool["function"]["name"] for tool in TOOLS}

TEXT_FILE_SUFFIXES = (".txt", ".md", ".csv", ".json", ".xml", ".html", ".htm")


def _file_path(item: Any) -> str | None:
    if isinstance(item, str):
        return item
    if isinstance(item, dict):
        if "path" in item:
            return item["path"]
        if item.get("type") == "file":
            file_data = item.get("file") or {}
            return file_data.get("path") or file_data.get("url")
    return None


def _text_content(item: Any) -> str | None:
    if isinstance(item, str):
        return item
    if isinstance(item, dict) and item.get("type") == "text":
        return item.get("text") or ""
    return None


def _file_to_api_part(path: str) -> dict[str, Any]:
    mime_type, _ = mimetypes.guess_type(path)
    mime_type = mime_type or "application/octet-stream"
    name = Path(path).name

    if mime_type.startswith("image/"):
        data = Path(path).read_bytes()
        encoded = base64.b64encode(data).decode("utf-8")
        return {
            "type": "image_url",
            "image_url": {"url": f"data:{mime_type};base64,{encoded}"},
        }

    if mime_type.startswith("text/") or path.lower().endswith(TEXT_FILE_SUFFIXES):
        try:
            text = Path(path).read_text(encoding="utf-8", errors="replace")
            return {"type": "text", "text": f"[File: {name}]\n{text}"}
        except OSError:
            pass

    return {"type": "text", "text": f"[Attached file: {name}]"}


def _parts_to_api_content(parts: list[dict[str, Any]]) -> str | list[dict[str, Any]]:
    if not parts:
        return ""
    if len(parts) == 1 and parts[0]["type"] == "text":
        return parts[0]["text"]
    return parts


def multimodal_input_to_api_content(
    message: str | dict[str, Any],
) -> str | list[dict[str, Any]]:
    """Convert a ChatInterface user message to Inference API content."""
    if isinstance(message, str):
        return message

    parts: list[dict[str, Any]] = []
    for file_ref in message.get("files") or []:
        path = _file_path(file_ref) if isinstance(file_ref, dict) else file_ref
        if path:
            parts.append(_file_to_api_part(path))

    text = (message.get("text") or "").strip()
    if text:
        parts.append({"type": "text", "text": text})

    return _parts_to_api_content(parts)


def chat_content_to_api(content: Any) -> str | list[dict[str, Any]] | None:
    """Convert Gradio chat history content to Inference API content."""
    if isinstance(content, str):
        return content or None

    if not isinstance(content, list):
        return None

    parts: list[dict[str, Any]] = []
    for item in content:
        text = _text_content(item)
        if text:
            parts.append({"type": "text", "text": text})
            continue
        path = _file_path(item)
        if path:
            parts.append(_file_to_api_part(path))

    if not parts:
        return None
    return _parts_to_api_content(parts)


def history_to_api_messages(history: list[dict[str, Any]]) -> list[dict[str, Any]]:
    """Convert Gradio chat history to API messages, skipping UI-only tool steps."""
    messages: list[dict[str, Any]] = []
    for msg in history:
        if msg.get("metadata"):
            continue
        role = msg.get("role")
        if role not in ("user", "assistant"):
            continue
        api_content = chat_content_to_api(msg.get("content"))
        if api_content:
            messages.append({"role": role, "content": api_content})
    return messages


def assistant_message_dict(
    content: str,
    tool_calls: list[dict[str, Any]] | None,
) -> dict[str, Any]:
    message: dict[str, Any] = {"role": "assistant", "content": content}
    if tool_calls:
        message["tool_calls"] = tool_calls
    return message


def langchain_messages_to_api(messages: list[Any]) -> list[dict[str, Any]]:
    """Convert LangGraph/LangChain message objects to Inference API messages."""
    from langchain_core.messages import AIMessage, BaseMessage, ToolMessage

    api_messages: list[dict[str, Any]] = []
    for message in messages:
        if isinstance(message, dict):
            role = message.get("role")
            if role not in ("system", "user", "assistant", "tool"):
                continue
            entry: dict[str, Any] = {
                "role": role,
                "content": message.get("content", ""),
            }
            if role == "tool" and message.get("tool_call_id"):
                entry["tool_call_id"] = message["tool_call_id"]
            if role == "assistant" and message.get("tool_calls"):
                entry["tool_calls"] = message["tool_calls"]
            api_messages.append(entry)
            continue

        if isinstance(message, ToolMessage):
            api_messages.append(
                {
                    "role": "tool",
                    "content": message.content,
                    "tool_call_id": message.tool_call_id,
                }
            )
            continue

        if isinstance(message, AIMessage):
            entry = {"role": "assistant", "content": message.content or ""}
            if message.tool_calls:
                entry["tool_calls"] = [
                    {
                        "id": tool_call.get("id")
                        if isinstance(tool_call, dict)
                        else tool_call["id"],
                        "type": "function",
                        "function": {
                            "name": tool_call.get("name")
                            if isinstance(tool_call, dict)
                            else tool_call["name"],
                            "arguments": json.dumps(
                                tool_call.get("args", {})
                                if isinstance(tool_call, dict)
                                else tool_call["args"]
                            ),
                        },
                    }
                    for tool_call in message.tool_calls
                ]
            api_messages.append(entry)
            continue

        if isinstance(message, BaseMessage):
            role = message.type
            if role == "human":
                role = "user"
            elif role == "ai":
                role = "assistant"
            api_messages.append({"role": role, "content": message.content})

    return api_messages


def api_turn_to_ai_message(
    content: str,
    reasoning: str,
    tool_calls: list[dict[str, Any]] | None,
) -> "AIMessage":
    """Build a LangChain AIMessage from one Inference API turn."""
    from langchain_core.messages import AIMessage

    additional_kwargs: dict[str, Any] = {}
    if reasoning:
        additional_kwargs["reasoning_content"] = reasoning

    lc_tool_calls: list[dict[str, Any]] = []
    if tool_calls:
        for tool_call in tool_calls:
            function = tool_call.get("function") or {}
            raw_args = function.get("arguments", "{}")
            if isinstance(raw_args, str):
                try:
                    parsed_args = json.loads(raw_args)
                except json.JSONDecodeError:
                    parsed_args = {"raw": raw_args}
            else:
                parsed_args = raw_args
            lc_tool_calls.append(
                {
                    "name": str(function.get("name") or ""),
                    "args": parsed_args,
                    "id": str(tool_call.get("id") or f"call_{uuid.uuid4().hex}"),
                    "type": "tool_call",
                }
            )

    return AIMessage(
        content=content or "",
        additional_kwargs=additional_kwargs,
        tool_calls=lc_tool_calls,
    )


def parse_tool_calls(message: Any) -> list[dict[str, Any]] | None:
    if not message.tool_calls:
        return None
    return [
        {
            "id": tool_call.id,
            "type": tool_call.type,
            "function": {
                "name": tool_call.function.name,
                "arguments": tool_call.function.arguments,
            },
        }
        for tool_call in message.tool_calls
    ]


def parse_text_tool_calls(text: str) -> list[dict[str, Any]] | None:
    """Recover tool calls when a model emits them as text instead of tool_calls."""
    if not text:
        return None

    for tool_name in TOOL_NAMES:
        arguments = _parse_json_call(text, tool_name) or _parse_python_call(
            text, tool_name
        )
        if arguments is None and tool_name == "update_globe":
            arguments = _parse_update_globe_plan(text)
        if arguments is None:
            continue
        return [
            {
                "id": f"call_{uuid.uuid4().hex}",
                "type": "function",
                "function": {
                    "name": tool_name,
                    "arguments": json.dumps(arguments),
                },
            }
        ]

    return None


def _parse_json_call(text: str, tool_name: str) -> dict[str, Any] | None:
    match = re.search(rf"(?:default_api:)?{re.escape(tool_name)}\s*\{{", text)
    if not match:
        return None

    start = text.find("{", match.start())
    if start == -1:
        return None

    decoder = json.JSONDecoder()
    try:
        arguments, _ = decoder.raw_decode(text[start:])
    except json.JSONDecodeError:
        return _parse_key_values(text[start:])

    return arguments if isinstance(arguments, dict) else None


def _parse_python_call(text: str, tool_name: str) -> dict[str, Any] | None:
    match = re.search(rf"{re.escape(tool_name)}\s*\((?P<args>[^)]*)\)", text)
    if not match:
        return None
    return _parse_key_values(match.group("args"))


def _parse_key_values(text: str) -> dict[str, str] | None:
    pairs = re.findall(r'(\w+)\s*=\s*["\']([^"\']+)["\']', text)
    if not pairs:
        pairs = re.findall(r'["\'](\w+)["\']\s*:\s*["\']([^"\']+)["\']', text)
    if not pairs:
        return None
    return {key: value for key, value in pairs}


def _parse_update_globe_plan(text: str) -> dict[str, Any] | None:
    """Recover the common prose format: "I will call update_globe... Parameters: ..."."""
    if not re.search(r"\bupdate_globe\b", text):
        return None

    action_match = re.search(r'\baction\s*[:=]\s*["\']?(\w+)["\']?', text)
    if not action_match:
        return None

    countries_match = re.search(r"\bcountries\s*[:=]\s*\[(?P<countries>[^\]]+)\]", text)
    countries = (
        re.findall(r'["\']([A-Z]{2})["\']', countries_match.group("countries"))
        if countries_match
        else []
    )

    labels_match = re.search(r"\blabels\s*[:=]\s*\[(?P<labels>[^\]]+)\]", text)
    labels = (
        re.findall(r'["\']([^"\']+)["\']', labels_match.group("labels"))
        if labels_match
        else []
    )

    arguments: dict[str, Any] = {"action": action_match.group(1)}
    if countries:
        arguments["countries"] = countries
    if labels:
        arguments["labels"] = labels

    zoom_match = re.search(r"\bzoom\s*[:=]\s*(\d+(?:\.\d+)?)", text)
    if zoom_match:
        arguments["zoom"] = float(zoom_match.group(1))

    return arguments