# coding=utf-8
# Copyright (c) Huawei Technologies Co., Ltd. 2025. All rights reserved.

import ast
import re
from dataclasses import dataclass, field
from typing import Optional, Union, Callable, List, Literal

from pydantic import BaseModel, ConfigDict, Field, ValidationError

from openjiuwen.core.common.exception.codes import StatusCode
from openjiuwen.core.common.exception.errors import build_error
from openjiuwen.core.common.logging import workflow_logger, LogEventType
from openjiuwen.core.common.security.exception_utils import ExceptionUtils
from openjiuwen.core.workflow.components.base import ComponentConfig
from openjiuwen.core.workflow.components.component import ComponentComposable, ComponentExecutable
from openjiuwen.core.workflow.components.flow.branch_router import BranchRouter
from openjiuwen.core.workflow.components.condition.condition import Condition
from openjiuwen.core.context_engine import ModelContext
from openjiuwen.core.graph.base import Graph
from openjiuwen.core.graph.executable import Output, Input
from openjiuwen.core.session.node import Session
from openjiuwen.core.foundation.llm import (
    BaseMessage, UserMessage, SystemMessage, AssistantMessage, ModelRequestConfig, ModelClientConfig, Model
)
from openjiuwen.core.foundation.prompt import PromptTemplate
from openjiuwen.core.common.security.user_config import UserConfig

LUI = "llm"
NAME = "name"
MODEL = "model"
CLASS = "class"
REASON = "reason"
INPUT = "input"
USER_PROMPT = "user_prompt"
CATEGORY_INFO = "category_info"
CATEGORY_LIST = "category_list"
CATEGORY_NAME_LIST = "category_name_list"
DEFAULT_CLASS = "default_class"
CHAT_HISTORY = "chat_history"
EXAMPLE_CONTENT = "example_content"
ENABLE_HISTORY = "enable_history"
ENABLE_INPUT = "enable_input"
LLM_INPUTS = "llm_inputs"
LLM_OUTPUTS = "llm_outputs"
MODEL_SOURCE = "modelType"
MODEL_NAME = "modelName"
EXTENSION = "extension"
CHAT_HISTORY_MAX_TURN = "chat_history_max_turn"
INTENT_DETECTION_TEMPLATE = "intent_detection_template"
ROLE = "role"
CONTENT = "content"
ROLE_MAP_ZH = {"user": '用户', 'assistant': '助手', 'system': '系统', 'tool': '工具'}
ROLE_MAP_EN = {"user": 'User', 'assistant': 'Assistant', 'system': 'System', 'tool': 'Tool'}
JSON_PARSE_FAIL_REASON_ZH = "当前意图识别的输出:'{result}'格式不符合有效的JSON规范，导致解析失败，因此返回默认分类。"
JSON_PARSE_FAIL_REASON_EN = ("The intent detection output '{result}' does not conform to valid JSON format, "
                             "parsing failed, therefore default category is returned.")
CLASS_KEY_MISSING_REASON_ZH = "当前意图识别的输出 '{result}' 缺少必要的输出'class'分类信息，因此返回默认分类。"
CLASS_KEY_MISSING_REASON_EN = ("The intent detection output '{result}' is missing "
                               "the required 'class' classification field, therefore default category is returned.")
VALIDATION_FAIL_REASON_ZH = "当前意图识别的输出类别 '{intent_class}' 不在预定义的分类列表: '{category_list}'中，因此系统返回默认分类。"
VALIDATION_FAIL_REASON_EN = ("The intent detection output class '{intent_class}' is not in the "
                             "predefined category list: '{category_list}', therefore default category is returned.")
WORKFLOW_CHAT_HISTORY = "workflow_chat_history"

RESULT = "result"
CATEGORY_NAME_ITS1 = "category_name_list"
FEW_SHOT_NUM = 5
ENABLE_Q2L = 'enableKnowledge'
DEFAULT_QUERY_CATE = 'title'
DEFAULT_CLASS_CATE = 'content'
DEFAULT_INT = "不确定，其他的意图"
SEARCH_TYPE = "faq"
SEARCH_NUM = 5
CLASSIFICATION_ID = "classificationId"
CLASSIFICATION_DEFAULT_ID = 0
CLASSIFICATION_NAME = "name"
CLASSIFICATION_DEFAULT_NAME_ZH = "默认意图"
CLASSIFICATION_DEFAULT_NAME_EN = "Default intent"
KG_FILTER_KEY = "filter_string"
KG_FILTER_PREFIX = "category:"
KG_SCOPE = "scope"

_PROVIDER_NAME_MAP = {
    "openai": "OpenAI",
    "siliconflow": "SiliconFlow",
}

DEFAULT_SYSTEM_PROMPT_ZH = "你是一个识别用户输入意图的AI助手。"
DEFAULT_SYSTEM_PROMPT_EN = "You are an AI assistant that identifies user input intent."

DEFAULT_USER_PROMPT_ZH = """
{{user_prompt}}

当前可供选择的功能分类如下：
{{category_info}}

用户与助手的对话历史：
{{chat_history}}

当前输入：
{{input}}

请根据当前输入和对话历史分析并输出最适合的功能分类。输出格式为 JSON，包含以下两个字段：
class: 代表分类结果
reason: 说明为何选择该分类
例如: {"class": "分类xx", "reason": "当前输入xxx"}
请参考以下示例：
{{example_content}}
如果没有合适的分类，请输出 {{default_class}}。
"""

DEFAULT_USER_PROMPT_EN = """
{{user_prompt}}

Available function categories:
{{category_info}}

Conversation history between user and assistant:
{{chat_history}}

Current input:
{{input}}

Please analyze the current input and conversation history, and output the most suitable function category. Output format is JSON with the following two fields:
class: represents the classification result
reason: explains why this classification was chosen
Example: {"class": "Category1", "reason": "Current input xxx"}
Please refer to the following examples:
{{example_content}}
If no suitable category exists, output {{default_class}}.
"""


def get_default_template(accept_language: Literal['zh', 'en'] = 'zh'):
    """Return prompt template based on accept_language."""
    if accept_language == 'en':
        return PromptTemplate(
            content=[
                SystemMessage(content=DEFAULT_SYSTEM_PROMPT_EN),
                UserMessage(content=DEFAULT_USER_PROMPT_EN),
            ]
        )
    return PromptTemplate(
        content=[
            SystemMessage(content=DEFAULT_SYSTEM_PROMPT_ZH),
            UserMessage(content=DEFAULT_USER_PROMPT_ZH),
        ]
    )


@dataclass
class IntentDetectionCompConfig(ComponentConfig):
    model_id: Optional[str] = None
    model_client_config: Optional[ModelClientConfig] = field(default=None)
    model_config: Optional[ModelRequestConfig] = field(default=None)
    category_name_list: list[str] = field(default_factory=list)
    user_prompt: str = ""
    example_content: list[str] = field(default_factory=list)
    enable_history: bool = False
    chat_history_max_turn: int = 3
    accept_language: Literal['zh', 'en'] = field(default='zh')


@dataclass
class IntentDetectionDefaultConfig:
    category_list: list[str] = field(default_factory=list)
    intent_detection_template: PromptTemplate = field(default_factory=lambda: get_default_template('zh'))
    default_class: str = '分类0'
    enable_input: bool = True


class IntentDetectionInput(BaseModel):
    query: str
    model_config = ConfigDict(extra='allow')  # Allow any extra fields


class IntentDetectionOutput(BaseModel):
    classification_id: int = Field(default=-1)
    reason: str = Field(default="")
    category_name: str = Field(default="")


@dataclass()
class IntentDetectionExecutable(ComponentExecutable):
    def __init__(self, component_config: IntentDetectionCompConfig):
        super().__init__()
        self._session: Union[Session, None] = None
        self._llm: Union[Model, None] = None
        self._initialized: bool = False
        self._config = component_config
        self._init_default_config_category_list(component_config)
        self._append_default_category()
        self._router: Union[BranchRouter, None] = None

    @staticmethod
    def _refix_llm_output(input_str):
        json_path = r'\{.*\}'
        match = re.search(json_path, input_str, re.DOTALL)
        if match:
            res = match.group(0)
            res = res.replace("false", "False").replace("true", "True").replace("null", "None")
            return res
        else:
            return input_str

    async def invoke(self, inputs: Input, session: Session, context: ModelContext) -> Output:
        """Invoke IntentDetection node"""
        self._set_session(session)
        self._router.set_session(session)
        await self._initialize_if_needed()
        await self._write_user_message_to_context(inputs, context)
        chat_history = await self._get_chat_history_from_context(context)
        current_inputs = self._prepare_detection_inputs(inputs, chat_history)
        llm_output = await self._invoke_llm_and_get_result(current_inputs)
        workflow_logger.info(
            "Intent detection completed",
            event_type=LogEventType.WORKFLOW_COMPONENT_END,
            component_id=self._session.get_executable_id(),
            component_type_str="IntentDetectionComponent",
            session_id=self._session.get_session_id(),
            metadata={
                "has_output": bool(llm_output),
                "output_length": len(llm_output) if llm_output else 0,
                "sensitive_mode": UserConfig.is_sensitive()
            }
        )
        await self._write_assistant_message_to_context(llm_output, context)
        intent_res = self._parse_detection_result(llm_output)
        return intent_res

    def set_router(self, router):
        self._router = router
        return self

    def post_commit(self) -> bool:
        return True

    async def _get_chat_history_from_context(self, context) -> List[BaseMessage]:
        chat_history = []
        if self._config.enable_history and context is not None:
            context_window = await context.get_context_window(dialogue_round=self._config.chat_history_max_turn)
            if context_window is not None:
                chat_history = context_window.get_messages()
        return chat_history

    def _get_category_info(self):
        return "\n".join(f"{cid}: {cname}" for cid, cname in
                         zip(self._default_config.category_list, self._config.category_name_list))

    def _set_session(self, session: Session):
        self._session = session

    async def _create_llm_instance(self) -> Model:
        if self._config.model_id is None:
            if self._config.model_client_config is None or self._config.model_config is None:
                raise build_error(
                    StatusCode.COMPONENT_INTENT_DETECTION_INVOKE_CALL_FAILED,
                    error_msg="failed to create llm instance"
                )
            return Model(self._config.model_client_config, self._config.model_config)
        else:
            from openjiuwen.core.runner import Runner
            return await Runner.resource_mgr.get_model(id=self._config.model_id)

    async def _initialize_if_needed(self):
        if not self._initialized:
            try:
                self._llm = await self._create_llm_instance()
                self._initialized = True
            except Exception as e:
                raise build_error(
                    StatusCode.COMPONENT_INTENT_DETECTION_LLM_INIT_FAILED,
                    error_msg="failed to initialize llm if needed",
                    cause=e
                ) from e

    def _prepare_detection_inputs(self, inputs, chat_history):
        current_inputs = {}
        global_intent_map = []

        current_inputs.update({
            USER_PROMPT: self._config.user_prompt,
            CATEGORY_INFO: self._get_category_info(),
            DEFAULT_CLASS: self._default_config.default_class,
            ENABLE_HISTORY: self._config.enable_history,
            ENABLE_INPUT: self._default_config.enable_input,
            EXAMPLE_CONTENT: "\n\n".join(self._config.example_content),
            CHAT_HISTORY_MAX_TURN: self._config.chat_history_max_turn,
            CHAT_HISTORY: ""
        })

        # Process chat history
        if self._config.enable_history:
            chat_history_str = self._format_chat_history(chat_history)
            current_inputs.update({CHAT_HISTORY: chat_history_str})

        # Process current input
        if self._default_config.enable_input:
            try:
                intent_detection_input = IntentDetectionInput.model_validate(inputs)
                current_inputs.update({INPUT: intent_detection_input.query or ""})
            except ValidationError as e:
                raise build_error(
                    StatusCode.COMPONENT_INTENT_DETECTION_INPUT_PARAM_ERROR,
                    error_msg=ExceptionUtils.format_validation_error(e),
                    cause=e
                ) from e

        current_inputs['global_intent_map'] = global_intent_map

        return current_inputs

    def _format_chat_history(self, chat_history):
        role_map = ROLE_MAP_EN if self._config.accept_language == 'en' else ROLE_MAP_ZH
        chat_history_str = ""
        for history in chat_history[-self._config.chat_history_max_turn:]:
            if history.role in role_map:
                chat_history_str += "{}: {}\n".format(
                    role_map.get(history.role), history.content
                )
        return chat_history_str

    def _pre_process(self, inputs: dict):
        """Pre-process inputs for model"""
        final_prompts = self._default_config.intent_detection_template.format(inputs).to_messages()
        return final_prompts

    def _parse_detection_result(self, llm_output):
        intent_class, reason = self._post_process_intent_detection(llm_output)
        intent_id_and_name = self._get_intent_id_and_name(intent_class)
        return IntentDetectionOutput(classification_id=intent_id_and_name.get(CLASSIFICATION_ID, -1), reason=reason,
                                     category_name=intent_id_and_name.get(CLASSIFICATION_NAME, "")
                                     ).model_dump(exclude_defaults=True)

    async def _invoke_llm_and_get_result(self, current_inputs):
        """invoke llm and get result"""
        llm_inputs = self._default_config.intent_detection_template.format(current_inputs).to_messages()
        workflow_logger.info(
            "Intent detection LLM invoke started",
            event_type=LogEventType.WORKFLOW_COMPONENT_START,
            component_id=self._session.get_executable_id(),
            component_type_str="IntentDetectionComponent",
            session_id=self._session.get_session_id(),
            metadata={
                "has_inputs": bool(llm_inputs),
                "input_count": len(llm_inputs) if llm_inputs else 0,
                "sensitive_mode": UserConfig.is_sensitive()
            }
        )
        llm_output_content = ""

        try:
            llm_output = await self._llm.invoke(messages=llm_inputs)
            llm_output_content = llm_output.content
        except Exception as e:
            raise build_error(
                StatusCode.COMPONENT_INTENT_DETECTION_INVOKE_CALL_FAILED,
                error_msg="failed to invoke llm and get result",
                cause=e
            ) from e
        workflow_logger.info(
            "Intent detection LLM invoke completed",
            event_type=LogEventType.WORKFLOW_COMPONENT_END,
            component_id=self._session.get_executable_id(),
            component_type_str="IntentDetectionComponent",
            session_id=self._session.get_session_id(),
            metadata={
                "has_output": bool(llm_output_content),
                "output_length": len(llm_output_content) if llm_output_content else 0,
                "sensitive_mode": UserConfig.is_sensitive()
            }
        )

        return llm_output_content

    def _post_process_intent_detection(self, result):
        """Post-process the result"""
        lang = self._config.accept_language
        json_fail_reason = JSON_PARSE_FAIL_REASON_EN if lang == 'en' else JSON_PARSE_FAIL_REASON_ZH
        class_missing_reason = CLASS_KEY_MISSING_REASON_EN if lang == 'en' else CLASS_KEY_MISSING_REASON_ZH
        validation_fail_reason = VALIDATION_FAIL_REASON_EN if lang == 'en' else VALIDATION_FAIL_REASON_ZH
        category_pattern = r"Category\d+" if lang == 'en' else r"分类\d+"

        try:
            result = self._refix_llm_output(result)
            parsed_dict = ast.literal_eval(result)
        except Exception:
            return self._default_config.default_class, json_fail_reason.format(result=result)

        if not isinstance(parsed_dict, dict):
            return self._default_config.default_class, json_fail_reason.format(result=result)

        # post_process class information
        if not parsed_dict.get(CLASS):
            return self._default_config.default_class, class_missing_reason.format(result=parsed_dict)

        intent_class = parsed_dict.get(CLASS).replace('\n', '').replace(' ', '').replace('"', '').replace("'", '')
        match = re.search(category_pattern, intent_class, re.IGNORECASE)
        if match:
            matched = match.group(0)
            if lang == 'en':
                digit_part = re.search(r'\d+', matched)
                matched = f"Category{digit_part.group(0)}" if digit_part else matched
            parsed_dict.update({CLASS: matched})

        if parsed_dict.get(CLASS) not in self._default_config.category_list:
            reason = validation_fail_reason.format(
                intent_class=parsed_dict.get(CLASS),
                category_list=self._default_config.category_list
            )
            parsed_dict.update({CLASS: self._default_config.default_class, REASON: reason})

        return parsed_dict.get(CLASS), parsed_dict.get(REASON, '')

    def _validate_llm_parsed_result(self, result):
        """Validation of LLM output"""
        return result in self._default_config.category_list

    def _append_default_category(self):
        default_name = (CLASSIFICATION_DEFAULT_NAME_EN if self._config.accept_language == 'en'
                        else CLASSIFICATION_DEFAULT_NAME_ZH)
        self._default_config.category_list = [self._default_config.default_class] + self._default_config.category_list
        self._config.category_name_list = [default_name] + self._config.category_name_list

    def _get_intent_id_and_name(self, intent_class):
        default_name = (CLASSIFICATION_DEFAULT_NAME_EN if self._config.accept_language == 'en'
                        else CLASSIFICATION_DEFAULT_NAME_ZH)
        intent_res = {CLASSIFICATION_ID: CLASSIFICATION_DEFAULT_ID, CLASSIFICATION_NAME: default_name}
        idx = next((i for i, category in enumerate(self._default_config.category_list) if category == intent_class), -1)
        if idx > -1:
            intent_res = {CLASSIFICATION_ID: idx, CLASSIFICATION_NAME: self._config.category_name_list[idx]}
        return intent_res

    def _init_default_config_category_list(self, component_config: IntentDetectionCompConfig):
        lang = getattr(component_config, 'accept_language', 'zh') or 'zh'
        category_prefix = "Category" if lang == 'en' else "分类"
        default_class = f"{category_prefix}0"
        self._default_config = IntentDetectionDefaultConfig(
            intent_detection_template=get_default_template(lang),
            default_class=default_class
        )
        for index, _ in enumerate(component_config.category_name_list, start=1):
            self._default_config.category_list.append(f"{category_prefix}{index}")

    async def _write_user_message_to_context(self, inputs, context):
        if context is None or not self._config.enable_history:
            return
        try:
            query = IntentDetectionInput.model_validate(inputs).query or ""
        except ValidationError:
            return
        if query:
            await context.add_messages([UserMessage(content=query)])

    async def _write_assistant_message_to_context(self, content: str, context):
        if context is None or not self._config.enable_history:
            return
        if content:
            await context.add_messages([AssistantMessage(content=content)])


class IntentDetectionComponent(ComponentComposable):
    def __init__(self, component_config: Optional[IntentDetectionCompConfig] = None):
        super().__init__()
        self._executable = None
        self._config = component_config
        self._router = BranchRouter()

    @property
    def executable(self) -> IntentDetectionExecutable:
        if self._executable is None:
            self._executable = self.to_executable()
        return self._executable

    def add_component(self, graph: Graph, node_id: str, wait_for_all: bool = False) -> None:
        graph.add_node(node_id, self.to_executable(), wait_for_all=wait_for_all)
        graph.add_conditional_edges(node_id, self._router)

    def to_executable(self) -> IntentDetectionExecutable:
        return IntentDetectionExecutable(self._config).set_router(self._router)

    def add_branch(self, condition: Union[str, Callable[[], bool], Condition], target: Union[str, list[str]],
                   branch_id: str = None):
        if isinstance(target, str):
            target = [target]
        self._router.add_branch(condition, target, branch_id=branch_id)

    def router(self) -> BranchRouter:
        return self._router
