from __future__ import annotations

import json
import uuid
from typing import List, Union, Dict, Any, Optional, Literal, TypedDict, cast, Mapping, Sequence
from loguru import logger

from turbo_agent_core.schema.states import Message, Action, ReActRound, MessageKnowledgeResourceContent
from turbo_agent_core.schema.enums import MessageRole, ActionStatus, KnowledgeType


# Linkage removed


class KnowledgeResourcePayload(TypedDict, total=False):
    id: str
    uri: str
    type: str
    name: str
    mimetype: str
    mime_type: str
    fileType: Optional[str]


class OpenAIImageUrlPayload(TypedDict):
    url: str


class OpenAITextPart(TypedDict):
    type: Literal["text"]
    text: str


class OpenAIImagePart(TypedDict):
    type: Literal["image_url"]
    image_url: OpenAIImageUrlPayload


OpenAIContentPart = Union[OpenAITextPart, OpenAIImagePart]
OpenAIMessageContent = Union[str, List[OpenAIContentPart]]


class OpenAIFunctionPayload(TypedDict):
    name: str
    arguments: str


class OpenAIToolCallPayload(TypedDict):
    id: str
    type: Literal["function"]
    function: OpenAIFunctionPayload


class OpenAIMessagePayload(TypedDict, total=False):
    role: Literal["user", "system", "assistant", "tool"]
    content: Optional[OpenAIMessageContent]
    tool_calls: List[OpenAIToolCallPayload]
    tool_call_id: str
    reasoning_content: str


class OpenAIIncomingFunctionPayload(TypedDict, total=False):
    name: str
    arguments: Union[str, Dict[str, Any]]


class OpenAIIncomingToolCallPayload(TypedDict, total=False):
    id: str
    function: OpenAIIncomingFunctionPayload


OpenAIIncomingContentPart = Union[OpenAITextPart, OpenAIImagePart]
OpenAIIncomingContent = Union[str, Sequence[OpenAIIncomingContentPart]]


class OpenAIIncomingMessagePayload(TypedDict, total=False):
    role: str
    content: Optional[OpenAIIncomingContent]
    tool_calls: Sequence[OpenAIIncomingToolCallPayload]
    reasoning_content: str


StateContentItem = Union[str, MessageKnowledgeResourceContent, KnowledgeResourcePayload]
StateMessageContent = Optional[Union[str, List[StateContentItem]]]
NormalizedStateContentItem = Union[str, KnowledgeResourcePayload]
NormalizedStateMessageContent = Optional[Union[str, List[NormalizedStateContentItem]]]


def _normalize_content_item(content: StateContentItem) -> NormalizedStateContentItem:
    if isinstance(content, MessageKnowledgeResourceContent):
        return cast(
            KnowledgeResourcePayload,
            content.model_dump(mode="json", by_alias=True, exclude_none=True),
        )
    return content


def _normalize_message_content(content: StateMessageContent) -> NormalizedStateMessageContent:
    if content is None:
        return None
    if isinstance(content, str):
        return content
    if isinstance(content, list):
        return [_normalize_content_item(item) for item in content]
    return str(content)


def _resource_to_openai_part(resource: KnowledgeResourcePayload) -> OpenAIContentPart:
    uri = str(resource.get("uri") or "").strip()
    name = str(resource.get("name") or "未命名附件").strip() or "未命名附件"
    mimetype = str(resource.get("mimetype") or resource.get("mime_type") or "").strip().lower()
    resource_type = str(resource.get("type") or "").strip()

    if uri and (mimetype.startswith("image/") or resource_type == KnowledgeType.Picture.value):
        return {
            "type": "image_url",
            "image_url": {
                "url": uri,
            },
        }

    description = f"[附件] {name}"
    if resource_type:
        description += f" ({resource_type})"
    if uri:
        description += f"\n资源地址: {uri}"

    return {
        "type": "text",
        "text": description,
    }


def _is_openai_content_part(item: object) -> bool:
    if not isinstance(item, dict):
        return False
    part_type = str(item.get("type") or "").strip()
    if part_type == "text":
        return "text" in item
    if part_type == "image_url":
        return isinstance(item.get("image_url"), dict) and "url" in item.get("image_url", {})
    return False


def _content_to_openai_content(content: StateMessageContent) -> OpenAIMessageContent:
    normalized = _normalize_message_content(content)
    if normalized is None:
        return ""
    if isinstance(normalized, str):
        return normalized
    if not isinstance(normalized, list):
        return str(normalized)

    parts: List[OpenAIContentPart] = []
    for item in normalized:
        normalized_item = _normalize_content_item(item)
        if isinstance(normalized_item, str):
            parts.append({"type": "text", "text": normalized_item})
            continue
        if isinstance(normalized_item, dict):
            if _is_openai_content_part(normalized_item):
                parts.append(cast(OpenAIContentPart, normalized_item))
                continue
            parts.append(_resource_to_openai_part(normalized_item))
            continue
        parts.append({"type": "text", "text": str(normalized_item)})

    if not parts:
        return ""
    if len(parts) == 1 and parts[0].get("type") == "text":
        return parts[0].get("text", "")
    return parts


def _read_message_field(source: Mapping[str, object] | object, key: str) -> object | None:
    if isinstance(source, Mapping):
        return source.get(key)
    return getattr(source, key, None)


def _read_function_payload(source: object | None) -> OpenAIIncomingFunctionPayload | None:
    if source is None:
        return None
    if isinstance(source, Mapping):
        return cast(OpenAIIncomingFunctionPayload, dict(source))

    name = getattr(source, "name", None)
    arguments = getattr(source, "arguments", None)
    payload: OpenAIIncomingFunctionPayload = {}
    if isinstance(name, str):
        payload["name"] = name
    if isinstance(arguments, (str, dict)):
        payload["arguments"] = arguments
    return payload or None


def _parse_tool_call_arguments(raw_arguments: object | None) -> Optional[Union[str, Dict[str, Any]]]:
    if raw_arguments is None:
        return None
    if isinstance(raw_arguments, dict):
        return raw_arguments
    if isinstance(raw_arguments, str):
        try:
            parsed_arguments = json.loads(raw_arguments)
        except json.JSONDecodeError:
            return raw_arguments
        return parsed_arguments if isinstance(parsed_arguments, dict) else raw_arguments
    return str(raw_arguments)


def _openai_content_to_state_content(content: Optional[OpenAIIncomingContent]) -> str:
    if content is None:
        return ""
    if isinstance(content, str):
        return content

    text_fragments: List[str] = []
    for part in content:
        if not isinstance(part, Mapping):
            text_fragments.append(str(part))
            continue

        part_type = str(part.get("type") or "").strip()
        if part_type == "text":
            text_fragments.append(str(part.get("text") or ""))
            continue

        if part_type == "image_url":
            image_payload = part.get("image_url")
            image_url = ""
            if isinstance(image_payload, Mapping):
                image_url = str(image_payload.get("url") or "").strip()
            text_fragments.append(f"[图片] {image_url}".rstrip())
            continue

        text_fragments.append(str(part))

    return "".join(fragment for fragment in text_fragments if fragment)


def convert_state_message_to_openai_message(core_message: Message) -> List[OpenAIMessagePayload]:
    """
    将 Turbo Agent 的 Core Message (原 State Message) 转换为 OpenAI 格式的消息列表。
    """
    if core_message.role == MessageRole.user:
        return [{"role": "user", "content": _content_to_openai_content(core_message.content)}]
    
    elif core_message.role == MessageRole.system:
        return [{"role": "system", "content": _content_to_openai_content(core_message.content)}]
    
    elif core_message.role == MessageRole.assistant:
        results: List[OpenAIMessagePayload] = []
        
        # 1. 优先处理 reactRounds (多轮对话历史复原)
        if hasattr(core_message, "reactRounds") and core_message.reactRounds:
            for round_data in core_message.reactRounds:
                
                # 1.1 构造 Assistant 消息 (可能包含 Tool Calls)
                tool_calls: List[OpenAIToolCallPayload] = []
                for action in round_data.actions:
                    args = {}
                    if isinstance(action.input, dict):
                        args = action.input
                    elif isinstance(action.input, str):
                        try:
                            args = json.loads(action.input)
                        except json.JSONDecodeError:
                            args = {"__raw__": action.input}
                    
                    tool_calls.append({
                        "id": action.id,
                        "type": "function",
                        "function": {
                            "name": action.name or "unknown_tool",
                            "arguments": json.dumps(args, ensure_ascii=False)
                        }
                    })

                assistant_content = round_data.content
                if isinstance(assistant_content, dict):
                    assistant_content = str(assistant_content)

                msg: OpenAIMessagePayload = {
                    "role": "assistant",
                    "content": cast(Optional[OpenAIMessageContent], assistant_content)
                }
                
                if tool_calls:
                    msg["tool_calls"] = tool_calls
                
                if round_data.thought:
                    msg["reasoning_content"] = round_data.thought

                if msg["content"] or tool_calls or msg.get("reasoning_content"):
                    results.append(msg)
                
                # 1.2 构造 Tool 结果消息
                for action in round_data.actions:
                    raw_content = action.observation
                    if raw_content is None:
                        tool_content = "Tool output not available"
                    elif isinstance(raw_content, str):
                        tool_content = raw_content
                    else:
                        try:
                            tool_content = json.dumps(raw_content, ensure_ascii=False)
                        except Exception:
                            tool_content = str(raw_content)

                    results.append({
                        "role": "tool",
                        "tool_call_id": action.id,
                        "content": tool_content
                    })
            
            return results

        # 2. Fallback: 普通 actions 处理 (无 reactRounds)
        if core_message.actions:
            tool_calls: List[OpenAIToolCallPayload] = []
            
            # 使用第一个 action 的 intent/summary 作为 assistant content
            first_action = core_message.actions[0]
            ai_content = first_action.intent or first_action.summary or ""
            if core_message.content and not ai_content:
                ai_content = core_message.content

            for action in core_message.actions:
                args = {}
                if isinstance(action.input, dict):
                    args = action.input
                elif isinstance(action.input, str):
                    try:
                        args = json.loads(action.input)
                    except json.JSONDecodeError:
                        args = {"__raw__": action.input}
                
                tool_calls.append({
                    "id": action.id,
                    "type": "function",
                    "function": {
                        "name": action.name or "unknown_tool",
                        "arguments": json.dumps(args, ensure_ascii=False)
                    }
                })
            
            msg: OpenAIMessagePayload = {
                "role": "assistant",
                "content": cast(Optional[OpenAIMessageContent], ai_content),
                "tool_calls": tool_calls
            }
            if core_message.reasoning_content:
                msg["reasoning_content"] = core_message.reasoning_content
            
            results.append(msg)
            
            # 添加 Tool Results
            for action in core_message.actions:
                raw_content = action.observation
                if raw_content is None:
                    tool_content = "Tool output not available"
                elif isinstance(raw_content, str):
                    tool_content = raw_content
                else:
                    try:
                        tool_content = json.dumps(raw_content, ensure_ascii=False)
                    except Exception:
                        tool_content = str(raw_content)

                results.append({
                    "role": "tool",
                    "tool_call_id": action.id,
                    "content": tool_content
                })

            # 如果还有额外的 content 且未被 consummed
            if core_message.content and core_message.content != ai_content:
                 results.append({
                     "role": "assistant",
                     "content": _content_to_openai_content(core_message.content)
                 })
            
            return results
        
        # 3. 纯文本 Assistant 回复
        msg: OpenAIMessagePayload = {
            "role": "assistant",
            "content": _content_to_openai_content(core_message.content)
        }
        if core_message.reasoning_content:
            msg["reasoning_content"] = core_message.reasoning_content
        return [msg]
    # Fallback
    return [{"role": core_message.role.value, "content": _content_to_openai_content(core_message.content)}]

def convert_openai_message_to_state_message(message: Mapping[str, object] | object) -> Optional[Message]:
    """
    将 OpenAI 格式的消息 (dict 或 ChatCompletionMessage 对象) 转换为 Core Message。
    """
    if not message:
        return None
    
    role_raw = _read_message_field(message, "role")
    content_raw = _read_message_field(message, "content")
    tool_calls_raw = _read_message_field(message, "tool_calls")
    reasoning_raw = _read_message_field(message, "reasoning_content")

    role_str = str(role_raw or "")
    content = cast(Optional[OpenAIIncomingContent], content_raw)
    reasoning_content = reasoning_raw if isinstance(reasoning_raw, str) else None
    
    role = MessageRole.user
    if role_str == "assistant":
        role = MessageRole.assistant
    elif role_str == "system":
        role = MessageRole.system
    elif role_str == "tool":
        # Tool message usually doesn't convert directly to a standalone State Message 
        # But if we must:
        return Message(
            id=str(uuid.uuid4()),
            role=MessageRole.assistant,
            content=_openai_content_to_state_content(content)
        )

    msg_id = str(uuid.uuid4())
    actions = []

    if isinstance(tool_calls_raw, Sequence) and not isinstance(tool_calls_raw, (str, bytes, bytearray)):
        for tc in tool_calls_raw:
            tc_id = _read_message_field(tc, "id")
            func = _read_function_payload(_read_message_field(tc, "function"))
            fname = func.get("name") if func else None
            parsed_arguments = _parse_tool_call_arguments(func.get("arguments") if func else None)

            action = Action(
                id=str(tc_id or uuid.uuid4()),
                name=fname if isinstance(fname, str) else None,
                input=parsed_arguments or {},
                status=ActionStatus.Pending
            )
            actions.append(action)

    # Construct ReActRound to preserve structure
    react_round = ReActRound(
        id=str(uuid.uuid4()),
        thought=reasoning_content,
        content=_openai_content_to_state_content(content) if content is not None else None,
        belongToMessageId=msg_id,
        actions=actions
    )

    msg = Message(
        id=msg_id,
        role=role,
        content=_openai_content_to_state_content(content),
        actions=actions,
        reactRounds=[react_round]
    )
    
    if reasoning_content:
        msg.reasoning_content = reasoning_content
            
    return msg
