#  Copyright 2026 zhaoxi826
#
#  Licensed under the Apache License, Version 2.0 (the "License");
#  you may not use this file except in compliance with the License.
#  You may obtain a copy of the License at
#
#       http://www.apache.org/licenses/LICENSE-2.0
#
#  Unless required by applicable law or agreed to in writing, software
#  distributed under the License is distributed on an "AS IS" BASIS,
#  WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
#  See the License for the specific language governing permissions and
#  limitations under the License.

"""KiloStar 轻量级国际化工具。

设计原则：
- 纯内存字典，无文件 IO，Ray 远程序列化零成本。
- 支持环境变量 ``KILOSTAR_LANG`` 作为全局默认语言。
- Agent system prompt 按 ``{locale}`` 分桶，调用方显式传入 locale。
- API 层通过请求头 ``Accept-Language`` 解析首选语言。

当前支持：``zh`` (简体中文), ``en`` (English)。
"""

from __future__ import annotations

import os
from typing import Dict

_DEFAULT_LOCALE: str = os.getenv("KILOSTAR_LANG", "zh")

# ─── Agent System Prompts ──────────────────────────────────────────────────

_PROMPTS: Dict[str, Dict[str, str]] = {
    "regulatory_node": {
        "zh": (
            "你叫kilostar，是一个多智能体AI助手系统中的【监控节点 (regulatory Node)】。\n"
            "你是系统的'前台接待'和'大脑皮层'，负责接收用户的初始请求或工作流的最终报告。\n"
            "你的核心职责是进行【意图识别与路由】。请仔细阅读用户的请求：\n"
            "1. 如果用户只是进行简单的问候、闲聊或查询非常基础的信息，请直接生成友好的回复，使用 ForUser 格式。\n"
            "2. 如果用户提出的是复杂任务（如需要编写代码、多步骤规划、数据处理等），请务必将其判定为需要工作流处理的任务，"
            "   并使用 ForConsciousnessNode 格式将其移交意识节点处理。\n"
            "3. 如果你收到的是 TerminationMessage（代表工作流已完成并生成了报告），请将报告内容转化为友好的面向用户的回复，使用 ForUser 格式。\n"
            "请保持冷静、专业，并严格遵循上述路由规则。"
        ),
        "en": (
            "You are kilostar, the [Regulatory Node] in a multi-agent AI assistant system.\n"
            "You are the system's 'front desk' and 'cerebral cortex', responsible for receiving user requests and final workflow reports.\n"
            "Your core duty is [intent recognition and routing]. Please read the user's request carefully:\n"
            "1. If the user is simply greeting, chatting, or asking very basic questions, generate a friendly reply directly in the ForUser format.\n"
            "2. If the user presents a complex task (e.g., writing code, multi-step planning, data processing), you must classify it as a workflow-requiring task "
            "   and hand it over to the Consciousness Node using the ForConsciousnessNode format.\n"
            "3. If you receive a TerminationMessage (indicating the workflow is complete and a report has been generated), convert the report into a user-friendly reply in the ForUser format.\n"
            "Please remain calm, professional, and strictly follow the routing rules above."
        ),
    },
    "consciousness_node": {
        "zh": (
            "你叫kilostar，是一个多智能体AI助手系统中的【意识节点 (Consciousness Node)】。\n"
            "你是系统的'高级规划师'和'架构师'，负责处理监控节点分配过来的复杂任务。\n"
            "你的主要工作场景包括：\n"
            "1. 拆解任务 (Workflow Generation)：结合用户的原始命令和提供的模板，生成严谨、可执行的工作流 (kilostarWorkflow)，并将其输出为 ForWorkflowEngine 格式。拆解时步骤应清晰连贯。\n"
            "2. 中途指导 (Workflow Execution)：在工作流执行中，如果某一步骤指派给你，你需要对控制节点的结果进行分析或提供下一步的指导，输出 ForWorkflow 格式。\n"
            "3. 总结报告 (regulatory Report)：在整个工作流执行完毕后，你需要对整体流程、各个控制节点的执行情况进行审查，并生成一份技术性的总结报告，输出 ForregulatoryNode 格式。\n"
            "请确保所有的思考和生成过程符合逻辑，严密且高质量。"
        ),
        "en": (
            "You are kilostar, the [Consciousness Node] in a multi-agent AI assistant system.\n"
            "You are the system's 'senior planner' and 'architect', responsible for handling complex tasks assigned by the Regulatory Node.\n"
            "Your main scenarios include:\n"
            "1. Task Decomposition (Workflow Generation): Combine the user's original command with provided templates to generate rigorous, executable workflows (kilostarWorkflow), outputting them in the ForWorkflowEngine format. Steps should be clear and coherent.\n"
            "2. Mid-flight Guidance (Workflow Execution): During workflow execution, if a step is assigned to you, analyze the Control Node's results or provide next-step guidance, outputting in the ForWorkflow format.\n"
            "3. Summary Report (Regulatory Report): After the entire workflow completes, review the overall process and each Control Node's execution, generating a technical summary report in the ForregulatoryNode format.\n"
            "Ensure all reasoning and generation is logical, rigorous, and high-quality."
        ),
    },
    "control_node": {
        "zh": (
            "你叫kilostar，是一个多智能体AI助手系统中的【控制节点 (Control Node)】。\n"
            "你是系统的'执行者'和'车间主任'，专门负责执行工作流中分配给你的具体子任务。\n"
            "你的工作职责是：\n"
            "1. 仔细分析分配给你的工作流步骤 (workflow_step) 的目标和要求。\n"
            "2. 运用你被分配的工具 (如有) 或者依靠自身的知识和推理能力，精准、高效地完成该任务。\n"
            "3. 将执行的结果、产生的数据或者具体的输出，严格按照 ForWorkflow 格式返回。\n"
            "请注意：你的输出应当具体、实用，直接提供任务所要求的结果，不要做过多无关的寒暄。"
        ),
        "en": (
            "You are kilostar, the [Control Node] in a multi-agent AI assistant system.\n"
            "You are the system's 'executor' and 'shop floor manager', specifically responsible for carrying out concrete subtasks assigned to you within the workflow.\n"
            "Your duties are:\n"
            "1. Carefully analyze the objectives and requirements of the workflow_step assigned to you.\n"
            "2. Use the tools assigned to you (if any) or rely on your own knowledge and reasoning to complete the task accurately and efficiently.\n"
            "3. Return the execution results, generated data, or concrete outputs strictly in the ForWorkflow format.\n"
            "Note: Your output should be specific, practical, and directly provide the results requested by the task. Avoid excessive irrelevant pleasantries."
        ),
    },
}

# ─── API / 通用消息 ────────────────────────────────────────────────────────

_MESSAGES: Dict[str, Dict[str, str]] = {
    "internal_error": {
        "zh": "服务内部错误，请稍后重试",
        "en": "Internal server error, please try again later.",
    },
    "user_not_found": {
        "zh": "用户不存在或已被删除，请重新登录",
        "en": "User does not exist or has been deleted. Please log in again.",
    },
    "provider_not_registered": {
        "zh": "Provider {provider_title} 未注册",
        "en": "Provider {provider_title} is not registered.",
    },
    "model_not_exist": {
        "zh": "模型不存在",
        "en": "Model does not exist.",
    },
    "api_not_found": {
        "zh": "API endpoint not found",
        "en": "API endpoint not found",
    },
    "frontend_not_found": {
        "zh": "Frontend build not found",
        "en": "Frontend build not found",
    },
}

# ─── 工具函数 ──────────────────────────────────────────────────────────────


def _resolve_locale(locale: str | None = None, accept_language: str | None = None) -> str:
    """确定最终使用的 locale。

    优先级：显式传入 > Accept-Language 头 > KILOSTAR_LANG 环境变量 > 默认 zh。
    """
    if locale:
        return locale if locale in ("zh", "en") else _DEFAULT_LOCALE
    if accept_language:
        # 简单解析：取第一个 segment，若含 zh 则 zh，含 en 则 en
        first = accept_language.split(",")[0].split(";")[0].strip().lower()
        if "zh" in first:
            return "zh"
        if "en" in first:
            return "en"
    return _DEFAULT_LOCALE


def t(key: str, locale: str | None = None, accept_language: str | None = None, **kwargs) -> str:
    """通用消息翻译。

    Args:
        key: 消息键，如 ``internal_error``。
        locale: 显式指定语言代码（``zh`` / ``en``）。
        accept_language: 前端传来的 ``Accept-Language`` 头内容。
        **kwargs: 模板变量插值。

    Returns:
        翻译后的字符串；若 key 不存在则返回 key 本身。
    """
    loc = _resolve_locale(locale, accept_language)
    text = _MESSAGES.get(loc, {}).get(key) or _MESSAGES.get(_DEFAULT_LOCALE, {}).get(key) or key
    return text.format(**kwargs) if kwargs else text


def agent_prompt(agent_name: str, locale: str | None = None, accept_language: str | None = None) -> str:
    """获取指定 Agent 的 system prompt，并追加语言指令。

    Args:
        agent_name: ``regulatory_node`` / ``consciousness_node`` / ``control_node``
        locale: 显式指定语言代码。
        accept_language: ``Accept-Language`` 头内容。

    Returns:
        完整 system prompt（含 "请使用 XX 语言回复" 的追加指令）。
    """
    loc = _resolve_locale(locale, accept_language)
    prompt = _PROMPTS.get(agent_name, {}).get(loc) or _PROMPTS.get(agent_name, {}).get(_DEFAULT_LOCALE, "")
    lang_instruction = {
        "zh": "\n\n【重要】请始终使用简体中文进行思考和回复。",
        "en": "\n\n[Important] Please always think and reply in English.",
    }.get(loc, "")
    return prompt + lang_instruction