# -*- coding: utf-8 -*-
"""OpenAI 插件 chat.py / llm.py 补丁共用逻辑（专题四、五、十）。"""
from __future__ import annotations

import re
import uuid

ENSURE_SYSTEM_FIRST = '''def _ensure_system_first(messages: list[PromptMessage]) -> list[PromptMessage]:
    """Qwen3/vLLM chat templates accept ONLY one system message, at the start.

    Dify (esp. Agent apps) stacks several system/developer prompts: KB assistant
    prompt + agent tool prompt + an empty system + config manifest... Merging them
    into a single system message (joined by blank lines, empties dropped) keeps it
    at the front; every other message keeps its original relative order.
    """
    system_parts: list[str] = []
    rest: list[PromptMessage] = []
    for m in messages:
        if isinstance(m, (SystemPromptMessage, DeveloperPromptMessage)):
            text = _text_content(m)
            if text.strip():
                system_parts.append(text)
        else:
            rest.append(m)
    if system_parts:
        rest.insert(0, SystemPromptMessage(content="\\n\\n".join(system_parts)))
    return rest
'''

DEBUG_DUMP_BLOCK = '''    # DEBUG: dump messages before sending to vLLM/Qwen
    try:
        with open('/tmp/dify_openai_messages.log', 'a', encoding='utf-8') as _df:
            _df.write('--- model=%s stream=%s ---\\n' % (model, stream))
            _df.write(json.dumps(
                [message(item) for item in _ensure_system_first(prompt_messages)],
                ensure_ascii=False, indent=2))
            _df.write('\\n\\n')
    except Exception as _e:
        pass
'''

GET_CUSTOMIZABLE_MODEL_SCHEMA = '''    def get_customizable_model_schema(
        self, model: str, credentials: dict
    ) -> AIModelEntity:
        schema = {item.model: item for item in self.predefined_models()}.get(
            _base_model(model)
        )
        if schema is None:
            # Fallback for custom / non-OpenAI base models served via an
            # OpenAI-compatible endpoint (e.g. Qwen, DeepSeek, local vLLM):
            # clone a chat-capable predefined schema so the model is listable
            # and can be set as the default model.
            template = next(
                (m for m in self.predefined_models()
                 if m.model_properties.get("mode") == "chat"),
                None,
            )
            if template is None:
                template = self.predefined_models()[0]
            props = dict(template.model_properties)
            mp = (credentials or {}).get("model_properties") or {}
            VLLM_MAX_CTX = 64000
            try:
                props["context_size"] = int(mp["context_size"]) if "context_size" in mp else VLLM_MAX_CTX
            except (ValueError, TypeError):
                props["context_size"] = VLLM_MAX_CTX
            # NOTE: do NOT put max_tokens into model_properties.
            props.pop("max_tokens", None)
            schema = template.model_copy(deep=True, update={"model_properties": props})
        return schema.model_copy(
            deep=True,
            update={
                "model": model,
                "label": I18nObject(zh_hans=model, en_us=model),
                "fetch_from": FetchFrom.CUSTOMIZABLE_MODEL,
            },
        )
'''


def find_plugin_file(d, filename: str) -> str:
    import _config as C

    out, err = d.sh(
        "docker exec %s sh -c 'find /app/storage/cwd -path \"*langgenius/openai*\" "
        "-name %s 2>/dev/null | head -1'" % (C.PLUGIN_CONTAINER, filename)
    )
    path = out.strip()
    if not path:
        raise RuntimeError("找不到 openai 插件 %s: %s" % (filename, err))
    return path


def read_remote_file(d, path: str) -> str:
    out, err = d.sh("docker exec docker-plugin_daemon-1 cat %s" % path)
    if not out and err:
        raise RuntimeError(err)
    return out


def replace_function(content: str, func_name: str, new_func: str) -> str:
    pattern = rf"(?ms)^def {re.escape(func_name)}\(.*?\n(?=^def |\Z)"
    if not re.search(pattern, content):
        raise RuntimeError("未找到函数 def %s" % func_name)
    return re.sub(pattern, new_func.rstrip() + "\n\n", content, count=1)


def replace_method(content: str, method_name: str, new_method: str) -> str:
    pattern = rf"(?ms)^    def {re.escape(method_name)}\(.*?\n(?=^    def |\Z)"
    if not re.search(pattern, content):
        raise RuntimeError("未找到方法 def %s" % method_name)
    return re.sub(pattern, new_method.rstrip() + "\n", content, count=1)


def inject_debug_log(content: str) -> str:
    if "# DEBUG: dump messages before sending to vLLM/Qwen" in content:
        return content
    marker = "client.chat.completions.create("
    idx = content.find(marker)
    if idx < 0:
        raise RuntimeError("未找到 client.chat.completions.create(")
    return content[:idx] + DEBUG_DUMP_BLOCK + "\n" + content[idx:]


def remove_debug_log(content: str) -> str:
    pattern = r"(?ms)^    # DEBUG: dump messages before sending to vLLM/Qwen\n.*?^    except Exception as _e:\n        pass\n"
    new_content, n = re.subn(pattern, "", content)
    if n == 0:
        raise RuntimeError("未找到调试日志块（可能已删除）")
    return new_content


def deploy_plugin_file(d, remote_path: str, content: str, *, restart: bool = True) -> None:
    import _config as C

    tag = uuid.uuid4().hex[:8]
    host_path = "/root/plugin_patch_%s.py" % tag
    bak = remote_path + ".bak_" + tag
    with d.sftp.open(host_path, "w") as f:
        f.write(content)
    cmds = [
        "docker exec docker-plugin_daemon-1 cp %s %s" % (remote_path, bak),
        "docker cp %s docker-plugin_daemon-1:%s" % (host_path, remote_path),
        "docker exec docker-plugin_daemon-1 python3 -m py_compile %s" % remote_path,
    ]
    if restart:
        cmds.append("cd %s && docker compose restart plugin_daemon" % C.DIFY_DOCKER_DIR)
    for cmd in cmds:
        out, err = d.sh(cmd, t=180)
        print((out or err).strip() or cmd)
    print("backup:", bak)