"""Recover DeepSeek DSML tool calls leaked through assistant text.

Some intranet OpenAI-compatible gateways render DeepSeek V4/V4.1 prompts
correctly but do not run the matching output parser.  The model consequently
returns its native DSML envelope in ``message.content`` and leaves
``message.tool_calls`` empty.  This module is a deliberately opt-in fallback
that converts only complete, schema-advertised calls back to OpenAI tool-call
chunks.

The parser is fail-closed:

* only tools advertised on the current provider request are accepted;
* names are matched exactly (misspellings are never corrected);
* an incomplete/malformed DSML envelope raises instead of executing a partial
  call or exposing the protocol as ordinary assistant text;
* ordinary text is streamed without buffering except for the short prefix
  needed to detect a marker split across provider chunks.
"""

from __future__ import annotations

import json
import re
import uuid
from collections.abc import Mapping
from dataclasses import dataclass, field
from typing import Any

from langchain_core.messages import AIMessageChunk
from langchain_core.outputs import ChatGenerationChunk


class DSMLToolCallProtocolError(RuntimeError):
    """A configured DSML provider returned an unsafe or incomplete call."""


@dataclass(frozen=True, slots=True)
class RecoveredToolCall:
    name: str
    arguments: dict[str, Any]
    call_id: str


@dataclass(frozen=True, slots=True)
class ParsedDSMLContent:
    content: str
    calls: tuple[RecoveredToolCall, ...]
    found: bool


_DSML_TOKEN = r"(?:｜DSML｜|\|DSML\|)"
_OUTER_NAMES = r"(?:tool_calls|toolcalls|function_calls|calls|tool)"
_START_RE = re.compile(
    rf"<\s*{_DSML_TOKEN}\s*{_OUTER_NAMES}\s*>",
    re.IGNORECASE,
)
_END_RE = re.compile(
    rf"</\s*{_DSML_TOKEN}\s*{_OUTER_NAMES}\s*>",
    re.IGNORECASE,
)
_ANY_DSML_RE = re.compile(rf"<\s*/?\s*{_DSML_TOKEN}", re.IGNORECASE)
_INVOKE_RE = re.compile(
    rf"<\s*{_DSML_TOKEN}\s*invoke\s+name\s*=\s*"
    rf"(?P<quote>[\"'])(?P<name>[^\"']+)(?P=quote)\s*>"
    rf"(?P<body>.*?)"
    rf"</\s*{_DSML_TOKEN}\s*invoke\s*>",
    re.IGNORECASE | re.DOTALL,
)
_PARAM_RE = re.compile(
    rf"<\s*{_DSML_TOKEN}\s*parameter\s+name\s*=\s*"
    rf"(?P<quote>[\"'])(?P<name>[^\"']+)(?P=quote)"
    rf"(?:\s+string\s*=\s*(?P<squote>[\"'])(?P<string>true|false)(?P=squote))?\s*>"
    rf"(?P<value>.*?)"
    rf"</\s*{_DSML_TOKEN}\s*parameter\s*>",
    re.IGNORECASE | re.DOTALL,
)


def _marker_candidates() -> tuple[str, ...]:
    names = ("tool_calls", "toolcalls", "function_calls", "calls", "tool")
    tokens = ("｜DSML｜", "|DSML|")
    return tuple(
        f"<{token}{space}{name}>"
        for token in tokens
        for space in ("", " ")
        for name in names
    )


_START_CANDIDATES = _marker_candidates()


def _longest_possible_start_suffix(text: str) -> int:
    """Length of the suffix that may be a split DSML opening marker."""
    max_len = min(len(text), max(map(len, _START_CANDIDATES)))
    for size in range(max_len, 0, -1):
        suffix = text[-size:]
        if any(candidate.startswith(suffix) for candidate in _START_CANDIDATES):
            return size
    return 0


def _parse_parameter_value(raw: str, string_flag: str | None) -> Any:
    if (string_flag or "").lower() == "true":
        return raw
    try:
        return json.loads(raw)
    except (TypeError, ValueError) as exc:
        if string_flag is None:
            # DeepSeek occasionally omits the flag for an otherwise ordinary
            # string parameter. Match the provider parser's safe literal
            # fallback, while keeping explicit string=false strict.
            return raw
        raise DSMLToolCallProtocolError(
            "DeepSeek DSML contained invalid JSON for a non-string parameter"
        ) from exc


def _parse_invoke_body(body: str) -> dict[str, Any]:
    params: dict[str, Any] = {}
    spans: list[tuple[int, int]] = []
    for match in _PARAM_RE.finditer(body):
        name = match.group("name").strip()
        if not name or name in params:
            raise DSMLToolCallProtocolError(
                "DeepSeek DSML contained an empty or duplicate parameter name"
            )
        params[name] = _parse_parameter_value(
            match.group("value"),
            match.group("string"),
        )
        spans.append(match.span())

    if spans:
        remainder_parts: list[str] = []
        cursor = 0
        for start, end in spans:
            remainder_parts.append(body[cursor:start])
            cursor = end
        remainder_parts.append(body[cursor:])
        if "".join(remainder_parts).strip():
            raise DSMLToolCallProtocolError(
                "DeepSeek DSML invoke mixed parameter tags with unexpected text"
            )
        return params

    stripped = body.strip()
    if not stripped:
        return {}
    try:
        direct = json.loads(stripped)
    except (TypeError, ValueError) as exc:
        raise DSMLToolCallProtocolError(
            "DeepSeek DSML invoke contained neither parameters nor a JSON object"
        ) from exc
    if not isinstance(direct, dict):
        raise DSMLToolCallProtocolError(
            "DeepSeek DSML direct arguments must be a JSON object"
        )
    return direct


def parse_dsml_content(
    text: str,
    *,
    allowed_tools: set[str],
) -> ParsedDSMLContent:
    """Parse complete V4/V4.1 DSML blocks and remove them from content."""
    if not text or not _ANY_DSML_RE.search(text):
        return ParsedDSMLContent(content=text, calls=(), found=False)

    calls: list[RecoveredToolCall] = []
    visible_parts: list[str] = []
    cursor = 0
    found = False

    while True:
        start = _START_RE.search(text, cursor)
        if start is None:
            visible_parts.append(text[cursor:])
            break
        found = True
        visible_parts.append(text[cursor : start.start()])
        end = _END_RE.search(text, start.end())
        if end is None:
            raise DSMLToolCallProtocolError(
                "DeepSeek DSML tool-call block was not closed"
            )
        inner = text[start.end() : end.start()]
        invoke_spans: list[tuple[int, int]] = []
        for invoke in _INVOKE_RE.finditer(inner):
            name = invoke.group("name").strip()
            if name not in allowed_tools:
                raise DSMLToolCallProtocolError(
                    f"DeepSeek DSML requested an unadvertised tool: {name!r}"
                )
            calls.append(
                RecoveredToolCall(
                    name=name,
                    arguments=_parse_invoke_body(invoke.group("body")),
                    call_id=f"call_deerflow_dsml_{uuid.uuid4().hex}",
                )
            )
            invoke_spans.append(invoke.span())
        if not invoke_spans:
            raise DSMLToolCallProtocolError(
                "DeepSeek DSML tool-call block contained no complete invoke"
            )
        remainder_parts: list[str] = []
        inner_cursor = 0
        for invoke_start, invoke_end in invoke_spans:
            remainder_parts.append(inner[inner_cursor:invoke_start])
            inner_cursor = invoke_end
        remainder_parts.append(inner[inner_cursor:])
        if "".join(remainder_parts).strip():
            raise DSMLToolCallProtocolError(
                "DeepSeek DSML tool-call block contained unexpected text"
            )
        cursor = end.end()

    if not found:
        # A DSML sentinel was present but did not match any supported complete
        # envelope. Do not leak it as chat text or guess a repair.
        raise DSMLToolCallProtocolError(
            "DeepSeek DSML marker used an unsupported tool-call envelope"
        )
    return ParsedDSMLContent(
        content="".join(visible_parts),
        calls=tuple(calls),
        found=True,
    )


def _payload_tool_names(payload: Mapping[str, Any]) -> set[str]:
    names: set[str] = set()
    tools = payload.get("tools")
    if not isinstance(tools, list):
        return names
    for tool in tools:
        function = tool.get("function") if isinstance(tool, Mapping) else None
        name = function.get("name") if isinstance(function, Mapping) else None
        if isinstance(name, str) and name:
            names.add(name)
    return names


@dataclass(slots=True)
class DSMLToolCallRecovery:
    """Per-provider-invocation DSML compatibility state."""

    enabled: bool
    allowed_tools: set[str] = field(default_factory=set)
    pending: str = ""
    capture: str = ""
    capture_template: ChatGenerationChunk | None = None
    recovered: bool = False

    def configure_from_payload(self, payload: Mapping[str, Any]) -> None:
        if self.enabled:
            self.allowed_tools = _payload_tool_names(payload)

    def recover_nonstream_response(self, response: dict[str, Any]) -> None:
        if not self.enabled or not self.allowed_tools:
            return
        choices = response.get("choices")
        if not isinstance(choices, list):
            return
        for choice in choices:
            if not isinstance(choice, dict):
                continue
            message = choice.get("message")
            if not isinstance(message, dict) or message.get("tool_calls"):
                continue
            content = message.get("content")
            if not isinstance(content, str):
                continue
            parsed = parse_dsml_content(content, allowed_tools=self.allowed_tools)
            if not parsed.found:
                continue
            message["content"] = parsed.content
            message["tool_calls"] = [
                {
                    "id": call.call_id,
                    "type": "function",
                    "function": {
                        "name": call.name,
                        "arguments": json.dumps(
                            call.arguments,
                            ensure_ascii=False,
                            separators=(",", ":"),
                        ),
                    },
                }
                for call in parsed.calls
            ]
            choice["finish_reason"] = "tool_calls"
            self.recovered = True

    def feed(self, chunk: ChatGenerationChunk) -> list[ChatGenerationChunk]:
        if not self.enabled or not self.allowed_tools:
            return [chunk]
        message = chunk.message
        if not isinstance(message, AIMessageChunk):
            return [chunk]
        if message.tool_call_chunks or message.tool_calls:
            if self.capture:
                raise DSMLToolCallProtocolError(
                    "Provider mixed native tool calls with an incomplete DSML block"
                )
            return [chunk]

        content = message.content
        if not isinstance(content, str):
            return [chunk]

        if self.recovered:
            if _ANY_DSML_RE.search(content):
                raise DSMLToolCallProtocolError(
                    "Provider emitted another DSML block after a recovered tool call"
                )
            return [self._force_tool_finish(chunk)]

        if self.capture:
            self.capture += content
            self.capture_template = chunk
            return self._complete_capture_if_ready()

        combined = self.pending + content
        had_pending = bool(self.pending)
        self.pending = ""
        start = _START_RE.search(combined)
        if start is not None:
            before = combined[: start.start()]
            self.capture = combined[start.start() :]
            self.capture_template = chunk
            output: list[ChatGenerationChunk] = []
            if before or self._has_noncontent_signal(message):
                output.append(self._replace_content(chunk, before))
            output.extend(self._complete_capture_if_ready())
            return output

        hold = _longest_possible_start_suffix(combined)
        if hold:
            safe = combined[:-hold]
            self.pending = combined[-hold:]
            self.capture_template = chunk
            if safe or self._has_noncontent_signal(message):
                return [self._replace_content(chunk, safe)]
            return []
        return [self._replace_content(chunk, combined)] if had_pending else [chunk]

    def finish(self) -> list[ChatGenerationChunk]:
        if not self.enabled:
            return []
        if self.capture:
            raise DSMLToolCallProtocolError(
                "Provider stream ended inside a DeepSeek DSML tool-call block"
            )
        if not self.pending:
            return []
        pending = self.pending
        self.pending = ""
        if self.capture_template is None:
            return []
        return [self._replace_content(self.capture_template, pending)]

    def _complete_capture_if_ready(self) -> list[ChatGenerationChunk]:
        if _END_RE.search(self.capture) is None:
            return []
        parsed = parse_dsml_content(
            self.capture,
            allowed_tools=self.allowed_tools,
        )
        if not parsed.calls or self.capture_template is None:
            raise DSMLToolCallProtocolError(
                "DeepSeek DSML recovery produced no complete tool call"
            )
        template = self.capture_template
        self.capture = ""
        self.recovered = True
        message = template.message
        assert isinstance(message, AIMessageChunk)
        tool_chunks = [
            {
                "name": call.name,
                "args": json.dumps(
                    call.arguments,
                    ensure_ascii=False,
                    separators=(",", ":"),
                ),
                "id": call.call_id,
                "index": index,
                "type": "tool_call_chunk",
            }
            for index, call in enumerate(parsed.calls)
        ]
        response_metadata = dict(message.response_metadata)
        if response_metadata.get("finish_reason"):
            response_metadata["finish_reason"] = "tool_calls"
        recovered_message = AIMessageChunk(
            content=parsed.content,
            additional_kwargs=dict(message.additional_kwargs),
            response_metadata=response_metadata,
            id=message.id,
            name=message.name,
            usage_metadata=message.usage_metadata,
            tool_call_chunks=tool_chunks,  # type: ignore[arg-type]
        )
        generation_info = dict(template.generation_info or {})
        if generation_info.get("finish_reason"):
            generation_info["finish_reason"] = "tool_calls"
        return [
            ChatGenerationChunk(
                message=recovered_message,
                generation_info=generation_info,
            )
        ]

    @staticmethod
    def _replace_content(
        chunk: ChatGenerationChunk,
        content: str,
    ) -> ChatGenerationChunk:
        return ChatGenerationChunk(
            message=chunk.message.model_copy(update={"content": content}),
            generation_info=chunk.generation_info,
        )

    @staticmethod
    def _has_noncontent_signal(message: AIMessageChunk) -> bool:
        return bool(
            message.additional_kwargs
            or message.response_metadata
            or message.usage_metadata
        )

    @staticmethod
    def _force_tool_finish(chunk: ChatGenerationChunk) -> ChatGenerationChunk:
        message = chunk.message
        if not isinstance(message, AIMessageChunk):
            return chunk
        response_metadata = dict(message.response_metadata)
        if response_metadata.get("finish_reason"):
            response_metadata["finish_reason"] = "tool_calls"
        generation_info = dict(chunk.generation_info or {})
        if generation_info.get("finish_reason"):
            generation_info["finish_reason"] = "tool_calls"
        return ChatGenerationChunk(
            message=message.model_copy(update={"response_metadata": response_metadata}),
            generation_info=generation_info or chunk.generation_info,
        )


__all__ = [
    "DSMLToolCallProtocolError",
    "DSMLToolCallRecovery",
    "ParsedDSMLContent",
    "RecoveredToolCall",
    "parse_dsml_content",
]