from __future__ import annotations

import json
from enum import IntEnum
from typing import Any, Dict, Optional

from langchain_core.messages import SystemMessage, HumanMessage, AIMessage, BaseMessage
from langchain_core.runnables import Runnable

from alayaflow.component.model import ModelManager
from alayaflow.utils.logger import AlayaFlowLogger


logger = AlayaFlowLogger()


class ResponseFormat(IntEnum):
    TEXT = 0
    JSON = 1

class LLMComponent:
    def __init__(
        self,
        *,
        model_id: str,
        system_prompt: str,
        prompt: str,

        temperature: Optional[float] = None,
        top_p: Optional[float] = None,
        max_tokens: Optional[int] = None,

        response_format: ResponseFormat = ResponseFormat.TEXT,
        json_schema: Optional[Dict[str, Any]] = None,
        outputs: Optional[Dict[str, str]] = None,
        retry_json_once: bool = True,
    ):
        # —— 配置即成员变量（= Spec）——
        self.model_id = model_id
        self.system_prompt = system_prompt
        self.prompt = prompt

        self.temperature = temperature
        self.top_p = top_p
        self.max_tokens = max_tokens

        self.response_format = response_format
        self.json_schema = json_schema
        self.outputs = outputs or {}
        self.retry_json_once = retry_json_once
        
        # —— 依赖注入（获取全局单例 ModelManager）——
        self._model_manager = ModelManager()
    
    def _get_llm(self) -> Runnable:
        bind_kwargs: Dict[str, Any] = {}

        if self.temperature is not None:
            bind_kwargs["temperature"] = self.temperature
        if self.top_p is not None:
            bind_kwargs["top_p"] = self.top_p
        if self.max_tokens is not None:
            bind_kwargs["max_tokens"] = self.max_tokens

        llm = self._model_manager.get_model(self.model_id, runtime_config=bind_kwargs)

        return llm
    
    def _invoke(self, system_prompt: str) -> AIMessage:
        llm = self._get_llm()
        resp = llm.invoke(
            [
                SystemMessage(content=system_prompt),
                HumanMessage(content=self.prompt),
            ]
        )
        if isinstance(resp, AIMessage):
            ret = resp
        elif isinstance(resp, BaseMessage):
            ret = AIMessage(
                content=getattr(resp, "content", str(resp)),
                additional_kwargs=getattr(resp, "additional_kwargs", {}) or {},
                response_metadata=getattr(resp, "response_metadata", {}) or {},
            )
        else:
            ret = AIMessage(content=str(resp))

        logger.debug(f"Invoke LLM.\nllm: {self.model_id}\n"
                f"=== system_prompt ===\n{system_prompt}\n"
                f"=== prompt ===\n{self.prompt}\n"
                f"=== response ===\n{ret.content}")
        
        return ret

    def __call__(self) -> AIMessage:
        if self.response_format == ResponseFormat.TEXT:
            return self._invoke(self.system_prompt)

        json_system = (
            self.system_prompt
            + "\n\n【输出格式要求】\n"
              "你必须只输出严格合法的 JSON。\n"
              "不要输出任何解释、前后缀、Markdown 代码块或多余文本。\n"
        )

        msg = self._invoke(json_system)

        try:
            json.loads(msg.content)
            return msg
        except Exception:
            if not self.retry_json_once:
                return msg

            strict_system = (
                self.system_prompt
                + "\n\n【严格输出格式要求】\n"
                  "只输出严格合法的 JSON（json.loads 必须可解析）。\n"
                  "禁止输出任何非 JSON 字符。\n"
            )
            return self._invoke(strict_system)

