"""
DeepSeek.v.something Encoding
High‑performance encoding/decoding for chat, tool calls, thinking mode, and tasks.
"""

from __future__ import annotations

import copy
import json
import re
from enum import Enum
from typing import Any, Dict, List, Optional, Tuple, Union

# ============================================================
# Custom Exceptions
# ============================================================

class DeepSeekV4Error(Exception):
    """Base exception for all DeepSeek‑V4 errors."""

class InvalidMessageError(DeepSeekV4Error):
    """Raised when a message is malformed or missing required fields."""

class InvalidThinkingModeError(DeepSeekV4Error):
    """Raised when an invalid thinking mode is provided."""

class InvalidReasoningEffortError(DeepSeekV4Error):
    """Raised when an invalid reasoning effort level is given."""

class ParsingError(DeepSeekV4Error):
    """Raised when model output parsing fails due to malformed syntax."""

# ============================================================
# Constants & Enums
# ============================================================

BOS_TOKEN: str = "<｜begin▁of▁sentence｜>"
EOS_TOKEN: str = "<｜end▁of▁sentence｜>"
THINKING_START_TOKEN: str = "<think>"
THINKING_END_TOKEN: str = "</think>"
DSML_TOKEN: str = "｜DSML｜"

USER_SP_TOKEN: str = "<｜User｜>"
ASSISTANT_SP_TOKEN: str = "<｜Assistant｜>"
LATEST_REMINDER_SP_TOKEN: str = "<｜latest_reminder｜>"

DS_TASK_SP_TOKENS: Dict[str, str] = {
    "action": "<｜action｜>",
    "query": "<｜query｜>",
    "authority": "<｜authority｜>",
    "domain": "<｜domain｜>",
    "title": "<｜title｜>",
    "read_url": "<｜read_url｜>",
}
VALID_TASKS = set(DS_TASK_SP_TOKENS.keys())

class ReasoningEffort(str, Enum):
    LOW = "low"
    HIGH = "high"
    MAX = "max"

REASONING_EFFORT_PROMPTS: Dict[ReasoningEffort, str] = {
    ReasoningEffort.LOW: "",
    ReasoningEffort.HIGH: (
        "Reasoning Effort: Absolute maximum with no shortcuts permitted.\n"
        "You MUST be very thorough in your thinking and comprehensively decompose the problem to resolve the root cause, rigorously stress-testing your logic against all potential paths, edge cases, and adversarial scenarios.\n"
        "Explicitly write out your entire deliberation process, documenting every intermediate step, considered alternative, and rejected hypothesis to ensure absolutely no assumption is left unchecked.\n\n"
    ),
    ReasoningEffort.MAX: (
        "Reasoning Effort: Beyond maximum — exhaustive, relentless, and uncompromising.\n"
        "You MUST reason with the utmost depth and rigor, leaving absolutely nothing to chance: exhaustively decompose the problem into its most fundamental components, trace every causal chain to its root, and resolve the underlying cause rather than any surface symptom.\n"
        "Do not stop reasoning until you have independently verified the solution from multiple angles and are certain that no assumption remains unchecked and no error remains undiscovered.\n\n"
    ),
}
DEFAULT_REASONING_EFFORT = ReasoningEffort.LOW

# Templates
SYSTEM_MSG_TEMPLATE: str = "{content}"
USER_MSG_TEMPLATE: str = "{content}"
LATEST_REMINDER_MSG_TEMPLATE: str = "{content}"
ASSISTANT_MSG_TEMPLATE: str = "{reasoning}{content}{tool_calls}" + EOS_TOKEN
ASSISTANT_MSG_WO_EOS_TEMPLATE: str = "{reasoning}{content}{tool_calls}"
THINKING_TEMPLATE: str = "{reasoning_content}"
RESPONSE_FORMAT_TEMPLATE: str = (
    "## Response Format:\n\nYou MUST strictly adhere to the following schema to reply:\n{schema}"
)
TOOL_CALL_TEMPLATE: str = (
    "<{dsml_token}invoke name=\"{name}\">\n{arguments}\n</{dsml_token}invoke>"
)
TOOL_CALLS_TEMPLATE: str = (
    "<{dsml_token}{tc_block_name}>\n{tool_calls}\n</{dsml_token}{tc_block_name}>"
)
TOOL_CALLS_BLOCK_NAME: str = "tool_calls"
TOOL_OUTPUT_TEMPLATE: str = "<tool_result>{content}</tool_result>"

TOOLS_TEMPLATE: str = """## Tools
You have access to a set of tools to help answer the user's question. You can invoke tools by writing a "<{dsml_token}tool_calls>" block like the following:
<{dsml_token}tool_calls>
<{dsml_token}invoke name="$TOOL_NAME">
<{dsml_token}parameter name="$PARAMETER_NAME" string="true|false">$PARAMETER_VALUE</{dsml_token}parameter>
...
</{dsml_token}invoke>
<{dsml_token}invoke name="$TOOL_NAME2">
...
</{dsml_token}invoke>
</{dsml_token}tool_calls>
String parameters should be specified as is and set `string="true"`. For all other types (numbers, booleans, arrays, objects), pass the value in JSON format and set `string="false"`.
If thinking_mode is enabled (triggered by {thinking_start_token}), you MUST output your complete reasoning inside {thinking_start_token}...{thinking_end_token} BEFORE any tool calls or final response.
Otherwise, output directly after {thinking_end_token} with tool calls or final response.
### Available Tool Schemas
{tool_schemas}
You MUST strictly follow the above defined tool name and parameter schemas to invoke tool calls.
"""

# ============================================================
# Utility Functions
# ============================================================

def to_json(value: Any) -> str:
    """Serialize any value to a JSON string with non‑ASCII preserved."""
    try:
        return json.dumps(value, ensure_ascii=False)
    except (TypeError, ValueError):
        return json.dumps(value, ensure_ascii=True)

def tools_from_openai_format(tools: List[Dict[str, Any]]) -> List[Dict[str, Any]]:
    return [tool["function"] for tool in tools]

def tool_calls_from_openai_format(tool_calls: List[Dict[str, Any]]) -> List[Dict[str, str]]:
    return [
        {"name": tc["function"]["name"], "arguments": tc["function"]["arguments"]}
        for tc in tool_calls
    ]

def tool_calls_to_openai_format(tool_calls: List[Dict[str, str]]) -> List[Dict[str, Any]]:
    return [
        {"type": "function", "function": {"name": tc["name"], "arguments": tc["arguments"]}}
        for tc in tool_calls
    ]

def encode_arguments_to_dsml(tool_call: Dict[str, str]) -> str:
    param_template = '<{dsml_token}parameter name="{key}" string="{is_str}">{value}</{dsml_token}parameter>'
    params = []
    try:
        arguments = json.loads(tool_call["arguments"])
    except json.JSONDecodeError:
        arguments = {"arguments": tool_call["arguments"]}
    for key, value in arguments.items():
        is_str = isinstance(value, str)
        params.append(
            param_template.format(
                dsml_token=DSML_TOKEN,
                key=key,
                is_str="true" if is_str else "false",
                value=value if is_str else to_json(value),
            )
        )
    return "\n".join(params)

def decode_dsml_to_arguments(tool_name: str, tool_args: Dict[str, Tuple[str, str]]) -> Dict[str, str]:
    def _fmt(key: str, value: str, is_str: str) -> str:
        if is_str == "true":
            value = to_json(value)
        return f"{to_json(key)}: {value}"
    args_json = "{" + ", ".join(_fmt(k, v, s) for k, (v, s) in tool_args.items()) + "}"
    return {"name": tool_name, "arguments": args_json}

def render_tools(tools: List[Dict[str, Any]]) -> str:
    tools_json = [to_json(t) for t in tools]
    return TOOLS_TEMPLATE.format(
        tool_schemas="\n".join(tools_json),
        dsml_token=DSML_TOKEN,
        thinking_start_token=THINKING_START_TOKEN,
        thinking_end_token=THINKING_END_TOKEN,
    )

def find_last_user_index(messages: List[Dict[str, Any]]) -> int:
    for idx in range(len(messages) - 1, -1, -1):
        if messages[idx].get("role") in ("user", "developer"):
            return idx
    return -1

# ============================================================
# Message Rendering
# ============================================================

def _render_system_message(msg: Dict[str, Any], tools: Optional[List[Dict]] = None) -> str:
    parts = [SYSTEM_MSG_TEMPLATE.format(content=msg.get("content", ""))]
    if tools:
        parts.append("\n\n" + render_tools(tools))
    if response_format := msg.get("response_format"):
        parts.append("\n\n" + RESPONSE_FORMAT_TEMPLATE.format(schema=to_json(response_format)))
    return "".join(parts)

def _render_developer_message(msg: Dict[str, Any], tools: Optional[List[Dict]] = None) -> str:
    content = USER_SP_TOKEN + msg.get("content", "")
    parts = [USER_MSG_TEMPLATE.format(content=content)]
    if tools:
        parts.append("\n\n" + render_tools(tools))
    if response_format := msg.get("response_format"):
        parts.append("\n\n" + RESPONSE_FORMAT_TEMPLATE.format(schema=to_json(response_format)))
    return "".join(parts)

def _render_user_message(msg: Dict[str, Any]) -> str:
    parts = [USER_SP_TOKEN]
    if content_blocks := msg.get("content_blocks"):
        block_parts = []
        for block in content_blocks:
            if block.get("type") == "text":
                block_parts.append(block.get("text", ""))
            elif block.get("type") == "tool_result":
                tool_content = block.get("content", "")
                if isinstance(tool_content, list):
                    text_parts = [
                        b.get("text", "") if b.get("type") == "text" else f"[Unsupported {b.get('type')}]"
                        for b in tool_content
                    ]
                    tool_content = "\n\n".join(text_parts)
                block_parts.append(TOOL_OUTPUT_TEMPLATE.format(content=tool_content))
            else:
                block_parts.append(f"[Unsupported {block.get('type')}]")
        parts.append("\n\n".join(block_parts))
    else:
        parts.append(msg.get("content", ""))
    return "".join(parts)

def _render_latest_reminder_message(msg: Dict[str, Any]) -> str:
    return LATEST_REMINDER_SP_TOKEN + LATEST_REMINDER_MSG_TEMPLATE.format(content=msg.get("content", ""))

def _render_assistant_message(
    msg: Dict[str, Any],
    thinking_mode: str,
    drop_thinking: bool,
    last_user_idx: int,
    current_idx: int,
    messages: List[Dict[str, Any]],
) -> str:
    tool_calls = msg.get("tool_calls")
    reasoning_content = msg.get("reasoning_content", "")
    content = msg.get("content", "")
    wo_eos = msg.get("wo_eos", False)

    tc_content = ""
    if tool_calls:
        tc_list = [
            TOOL_CALL_TEMPLATE.format(
                dsml_token=DSML_TOKEN,
                name=tc["name"],
                arguments=encode_arguments_to_dsml(tc),
            )
            for tc in tool_calls
        ]
        tc_content = "\n\n" + TOOL_CALLS_TEMPLATE.format(
            dsml_token=DSML_TOKEN,
            tool_calls="\n".join(tc_list),
            tc_block_name=TOOL_CALLS_BLOCK_NAME,
        )

    prev_has_task = current_idx > 0 and messages[current_idx - 1].get("task") is not None
    thinking_part = ""
    if thinking_mode == "thinking" and not prev_has_task:
        if not drop_thinking or current_idx > last_user_idx:
            thinking_part = THINKING_TEMPLATE.format(reasoning_content=reasoning_content) + THINKING_END_TOKEN

    template = ASSISTANT_MSG_WO_EOS_TEMPLATE if wo_eos else ASSISTANT_MSG_TEMPLATE
    return template.format(reasoning=thinking_part, content=content, tool_calls=tc_content)

def render_message(
    index: int,
    messages: List[Dict[str, Any]],
    thinking_mode: str,
    drop_thinking: bool = True,
    reasoning_effort: Optional[Union[str, ReasoningEffort]] = None,
) -> str:
    if thinking_mode not in ("chat", "thinking"):
        raise InvalidThinkingModeError(f"thinking_mode must be 'chat' or 'thinking', got '{thinking_mode}'")

    msg = messages[index]
    role = msg.get("role")
    if role not in ("system", "developer", "user", "latest_reminder", "assistant"):
        raise InvalidMessageError(f"Unsupported role '{role}' at index {index}")

    tools = msg.get("tools")
    if tools:
        tools = tools_from_openai_format(tools)

    # Pre‑process tool calls (if any) for rendering
    tool_calls = msg.get("tool_calls")
    if tool_calls:
        tool_calls = tool_calls_from_openai_format(tool_calls)
        # We'll put them back later in the assistant renderer
        # But we need to pass them to _render_assistant_message, so we'll just
        # keep the original msg and let _render_assistant_message handle it.

    prefix = ""
    if index == 0 and thinking_mode == "thinking":
        effort = ReasoningEffort(reasoning_effort) if reasoning_effort else DEFAULT_REASONING_EFFORT
        if effort not in REASONING_EFFORT_PROMPTS:
            raise InvalidReasoningEffortError(
                f"Invalid reasoning_effort '{effort}', expected one of {list(ReasoningEffort)}"
            )
        prefix = REASONING_EFFORT_PROMPTS[effort]

    if role == "system":
        rendered = _render_system_message(msg, tools)
    elif role == "developer":
        rendered = _render_developer_message(msg, tools)
    elif role == "user":
        rendered = _render_user_message(msg)
    elif role == "latest_reminder":
        rendered = _render_latest_reminder_message(msg)
    elif role == "assistant":
        last_user_idx = find_last_user_index(messages)
        rendered = _render_assistant_message(
            msg, thinking_mode, drop_thinking, last_user_idx, index, messages
        )
    else:
        raise InvalidMessageError(f"Unhandled role '{role}' at index {index}")

    full = prefix + rendered

    # Append transition tokens
    next_role = messages[index + 1].get("role") if index + 1 < len(messages) else None
    if next_role in ("assistant", "latest_reminder"):
        return full

    task = msg.get("task")
    if task is not None:
        if task not in VALID_TASKS:
            raise InvalidMessageError(f"Invalid task '{task}', allowed: {list(VALID_TASKS)}")
        token = DS_TASK_SP_TOKENS[task]
        if task != "action":
            full += token
        else:
            full += ASSISTANT_SP_TOKEN
            full += THINKING_END_TOKEN if thinking_mode == "chat" else THINKING_START_TOKEN
            full += token
    elif role in ("user", "developer"):
        full += ASSISTANT_SP_TOKEN
        if thinking_mode == "thinking" and (
            (not drop_thinking) or (drop_thinking and index >= find_last_user_index(messages))
        ):
            full += THINKING_START_TOKEN
        else:
            full += THINKING_END_TOKEN

    return full

# ============================================================
# Preprocessing
# ============================================================

def merge_tool_messages(messages: List[Dict[str, Any]]) -> List[Dict[str, Any]]:
    merged: List[Dict[str, Any]] = []
    for msg in copy.deepcopy(messages):
        role = msg.get("role")
        if role == "tool":
            tool_block = {
                "type": "tool_result",
                "tool_use_id": msg.get("tool_call_id", ""),
                "content": msg.get("content", ""),
            }
            if merged and merged[-1].get("role") == "user" and "content_blocks" in merged[-1]:
                merged[-1]["content_blocks"].append(tool_block)
            else:
                merged.append({"role": "user", "content_blocks": [tool_block]})
        elif role == "user":
            text_block = {"type": "text", "text": msg.get("content", "")}
            if (
                merged
                and merged[-1].get("role") == "user"
                and "content_blocks" in merged[-1]
                and merged[-1].get("task") is None
            ):
                merged[-1]["content_blocks"].append(text_block)
            else:
                new_msg = {
                    "role": "user",
                    "content": msg.get("content", ""),
                    "content_blocks": [text_block],
                }
                for key in ("task", "wo_eos", "mask"):
                    if key in msg:
                        new_msg[key] = msg[key]
                merged.append(new_msg)
        else:
            merged.append(msg)
    return merged

def sort_tool_results_by_call_order(messages: List[Dict[str, Any]]) -> List[Dict[str, Any]]:
    last_tool_call_order: Dict[str, int] = {}
    for msg in messages:
        role = msg.get("role")
        if role == "assistant" and msg.get("tool_calls"):
            last_tool_call_order = {}
            for idx, tc in enumerate(msg["tool_calls"]):
                tc_id = tc.get("id") or tc.get("function", {}).get("id", "")
                if tc_id:
                    last_tool_call_order[tc_id] = idx
        elif role == "user" and msg.get("content_blocks"):
            tool_blocks = [b for b in msg["content_blocks"] if b.get("type") == "tool_result"]
            if len(tool_blocks) > 1 and last_tool_call_order:
                sorted_blocks = sorted(
                    tool_blocks,
                    key=lambda b: last_tool_call_order.get(b.get("tool_use_id", ""), 0)
                )
                sorted_idx = 0
                new_blocks = []
                for block in msg["content_blocks"]:
                    if block.get("type") == "tool_result":
                        new_blocks.append(sorted_blocks[sorted_idx])
                        sorted_idx += 1
                    else:
                        new_blocks.append(block)
                msg["content_blocks"] = new_blocks
    return messages

def _drop_thinking_messages(messages: List[Dict[str, Any]]) -> List[Dict[str, Any]]:
    last_user_idx = find_last_user_index(messages)
    keep_roles = {"user", "system", "tool", "latest_reminder", "direct_search_results"}
    result = []
    for idx, msg in enumerate(messages):
        role = msg.get("role")
        if role in keep_roles or idx >= last_user_idx:
            result.append(msg)
        elif role == "assistant":
            new_msg = copy.copy(msg)
            new_msg.pop("reasoning_content", None)
            result.append(new_msg)
    return result

# ============================================================
# Main Encoding Function
# ============================================================

def encode_messages(
    messages: List[Dict[str, Any]],
    thinking_mode: str,
    context: Optional[List[Dict[str, Any]]] = None,
    drop_thinking: bool = True,
    add_default_bos_token: bool = True,
    reasoning_effort: Optional[Union[str, ReasoningEffort]] = None,
) -> str:
    context = context or []

    messages = merge_tool_messages(messages)
    messages = sort_tool_results_by_call_order(context + messages)[len(context):]
    if context:
        context = merge_tool_messages(context)
        context = sort_tool_results_by_call_order(context)

    full_messages = context + messages

    prompt = BOS_TOKEN if add_default_bos_token and len(context) == 0 else ""

    effective_drop_thinking = drop_thinking
    if any(m.get("tools") for m in full_messages):
        effective_drop_thinking = False

    if thinking_mode == "thinking" and effective_drop_thinking:
        full_messages = _drop_thinking_messages(full_messages)
        context_len = len(_drop_thinking_messages(context))
        num_to_render = len(full_messages) - context_len
    else:
        context_len = len(context)
        num_to_render = len(messages)

    for idx in range(num_to_render):
        prompt += render_message(
            idx + context_len,
            full_messages,
            thinking_mode=thinking_mode,
            drop_thinking=effective_drop_thinking,
            reasoning_effort=reasoning_effort,
        )

    return prompt

# ============================================================
# Parsing (Decoding Model Output)
# ============================================================

def _read_until_stop(index: int, text: str, stop: List[str]) -> Tuple[int, str, Optional[str]]:
    min_pos = len(text)
    matched = None
    for s in stop:
        pos = text.find(s, index)
        if pos != -1 and pos < min_pos:
            min_pos = pos
            matched = s
    if matched is not None:
        return min_pos + len(matched), text[index:min_pos], matched
    return len(text), text[index:], None

def parse_tool_calls(index: int, text: str) -> Tuple[int, Optional[str], List[Dict[str, str]]]:
    tool_calls: List[Dict[str, Any]] = []
    stop_token = None
    end_block = f"</{DSML_TOKEN}{TOOL_CALLS_BLOCK_NAME}>"

    while index < len(text):
        index, content_before, stop_token = _read_until_stop(
            index, text, [f"<{DSML_TOKEN}invoke", end_block]
        )
        if content_before != ">\n":
            raise ParsingError(f"Expected '>\\n' after invoke, got '{content_before}'")

        if stop_token == end_block:
            break

        if stop_token is None:
            raise ParsingError("Unexpected end of text while parsing tool calls")

        # Read tool name
        index, tool_name_content, stop_token = _read_until_stop(
            index, text, [f"<{DSML_TOKEN}parameter", f"</{DSML_TOKEN}invoke"]
        )
        name_match = re.search(r'^\s*name="(.*?)">\n$', tool_name_content, re.DOTALL)
        if not name_match:
            raise ParsingError(f"Could not parse tool name from: '{tool_name_content}'")
        tool_name = name_match.group(1)

        tool_args: Dict[str, Tuple[str, str]] = {}

        while stop_token == f"<{DSML_TOKEN}parameter":
            index, param_content, stop_token = _read_until_stop(
                index, text, [f"/{DSML_TOKEN}parameter"]
            )
            param_match = re.search(
                r'^ name="(.*?)" string="(true|false)">(.*?)<$',
                param_content,
                re.DOTALL
            )
            if not param_match:
                raise ParsingError(f"Could not parse parameter from: '{param_content}'")
            param_name, is_str, param_value = param_match.groups()
            if param_name in tool_args:
                raise ParsingError(f"Duplicate parameter name '{param_name}'")
            tool_args[param_name] = (param_value, is_str)

            index, content_after, stop_token = _read_until_stop(
                index, text, [f"<{DSML_TOKEN}parameter", f"</{DSML_TOKEN}invoke"]
            )
            if content_after != ">\n":
                raise ParsingError(f"Expected '>\\n' after parameter, got '{content_after}'")

        tool_call = decode_dsml_to_arguments(tool_name, tool_args)
        tool_calls.append(tool_call)

    return index, stop_token, tool_calls

def parse_message_from_completion_text(text: str, thinking_mode: str) -> Dict[str, Any]:
    if thinking_mode not in ("chat", "thinking"):
        raise InvalidThinkingModeError(f"thinking_mode must be 'chat' or 'thinking', got '{thinking_mode}'")

    summary_content = ""
    reasoning_content = ""
    tool_calls = []
    index = 0
    start_tool_calls = f"\n\n<{DSML_TOKEN}{TOOL_CALLS_BLOCK_NAME}"

    if thinking_mode == "thinking":
        index, content, stop_token = _read_until_stop(
            index, text, [THINKING_END_TOKEN, start_tool_calls]
        )
        reasoning_content = content
        if stop_token != THINKING_END_TOKEN:
            raise ParsingError("Missing closing </think> token in thinking mode output")

    index, content, stop_token = _read_until_stop(
        index, text, [EOS_TOKEN, start_tool_calls]
    )
    summary_content = content
    is_tool_calling = (stop_token == start_tool_calls)

    if is_tool_calling:
        index, stop_token, tool_calls = parse_tool_calls(index, text)
        index, leftover, stop_token = _read_until_stop(index, text, [EOS_TOKEN])
        if leftover:
            raise ParsingError(f"Unexpected content after tool calls: '{leftover}'")
        if stop_token not in (EOS_TOKEN, None):
            raise ParsingError(f"Expected EOS after tool calls, got '{stop_token}'")

    if index != len(text):
        rest = text[index:].strip()
        if rest:
            raise ParsingError(f"Unexpected trailing content: '{rest}'")

    for token in (BOS_TOKEN, EOS_TOKEN, THINKING_START_TOKEN, THINKING_END_TOKEN, DSML_TOKEN):
        if token in summary_content or token in reasoning_content:
            raise ParsingError(f"Special token '{token}' found in content after parsing")

    return {
        "role": "assistant",
        "content": summary_content,
        "reasoning_content": reasoning_content,
        "tool_calls": tool_calls_to_openai_format(tool_calls),
    }

# ============================================================
# Example Usage
# ============================================================

if __name__ == "__main__":
    # Minimal example: a simple user‑assistant conversation
    messages = [
        {"role": "user", "content": "What is the capital of France?"}
    ]
    prompt = encode_messages(messages, thinking_mode="chat")
    print("Encoded prompt:\n", prompt)

    dummy_response = "The capital of France is Paris." + EOS_TOKEN
    parsed = parse_message_from_completion_text(dummy_response, thinking_mode="chat")
    print("Parsed assistant message:", parsed)