# planning_agent/langchain_pipeline.py
import os
import sys
import json
from typing import Optional
from langchain_openai import ChatOpenAI
from planner.planning_agent.rag_pipeline import LocalKnowledgeBase

def generate_natural_plan(requirement_text: str, top_k: int = 4,) -> str:
    """
    使用 Qwen3-32B 结合本地知识库（RAG）生成标准化、可执行的测试任务规划。
    requirement_text: 待生成规划的需求文本（通常来自 MinerU 合并文本）。

    """
    api_key = os.getenv("QWEN_API_KEY")
    AI_BASE_URL = os.getenv("AI_BASE_URL")
    AI_MODEL = os.getenv("AI_MODEL")
    AI_API_KEY = os.getenv("AI_API_KEY")
    if not api_key:
        raise ValueError("请先在环境变量中设置 QWEN_API_KEY")


    kb = LocalKnowledgeBase()
    try:
        context = kb.retrieve_context(requirement_text, top_k=top_k)
    except FileNotFoundError:
        # no knowledge files, set empty context
        context = ""


    # instantiate llm (OpenAI-compatible wrapper used for Qwen via dashscope)
    # 关闭流式，直接获取完整响应
    # llm = ChatOpenAI(
    #     model="Qwen/Qwen3-32B",
    #     openai_api_base="https://api.siliconflow.cn/v1",
    #     openai_api_key=api_key.strip(),
    #     streaming=False
    # )
    llm = ChatOpenAI(
    model=AI_MODEL,
    openai_api_base=AI_BASE_URL,
    openai_api_key=AI_API_KEY.strip(),
    streaming=False
    )

    # system prompt: inject retrieved context and strict output template
    system_prompt = (
        "Role / 角色设定"
        "你是一位深耕实验室自动化测试领域、精通各种电子测量仪器（频谱仪、信号源、网络分析仪等）的测试任务规划专家。你的核心能力是将用户不完整、口语化的测试需求，转化为逻辑严密、参数完整、可供自动化系统直接解析的分步执行计划。"
        "Task Description / 任务描述"
        " 用户将提供一段包含测试器件（DUT）、仪器清单和粗略测试项的文本。你必须："
        "1. 需求解构：识别核心测试项、测试频点、仪器参数及判定准则。"
        "2. 内容补全：基于专业知识及接入的历史规划知识库，对模糊描述进行标准化（例如：补全上电时的 OCP/OVP 保护逻辑，细化开关矩阵的切换动作）。"
        "3. 分步规划：将测试项拆解为“步骤级”动作。每一个步骤必须包含：设置参数、执行动作、读取数据。"
        "Constraints / 约束条件"
        " 严禁无关内容：输出必须紧扣测试步骤，不要凭空编造参数或操作，不要输出开场白或解释性废话。"
        " 步骤独立性：每个步骤必须是一个闭环的动作，每个步骤都详细说明参数动作都是什么。"
        " 硬件逻辑：必须考虑开关矩阵（Switch Matrix）的切换逻辑，确保信号路径在每个频段下是正确的。"

        "Output Format / 输出模板"
        "请始终按照以下结构进行输出："
        "[任务概览]"
        " 测试器件：{DUT 描述}"
        " 仪器环境：{仪器列表及连接关系概述}"
        "[详细测试规程]"
        "步骤 X：{步骤名称}"
        "目的：简述本步骤的目标。"
        "仪器配置："
        " {仪器 A}：{参数 1}={值}, {参数 2}={值}..."
        " {开关矩阵}：切换至 {路径 X}"
        " 执行动作："
        "1. {动作 1}"
        "2. {动作 2}"
        " 数据采集/处理："
        " 读取项：{例如 Max Peak 功率}"
        " 计算逻辑：{如有，写出公式，不要写计算结果}"
        " 判定标准：{如有，写出阈值及结果处理}"
        "Execution Logic / 执行逻辑参考（知识库补全指导）"
        "1. 上电项：必须包含电源初始化、设置 OCP（过流保护）/OVP（过压保护）、逐步升压、监测静态电流。"
        "2. 射频测试项：必须先执行路径校准/路径切换，再设置信号源频率/功率，最后设置频谱仪（Span, RBW, VBW, Detector）。"

        "【参考知识库片段】可以若有相同的测试方法可以作为模板参考参，若需求给的参数不全可以参考知识库对对应流程的不全参数进行补充，若需求明确提出，则使用需求中提到的参数。\n"
        f"{context}\n\n"
    )
    
    # ---------------------------
    # 直接生成完整内容（非流式）
    # ---------------------------

    messages = [
        ("system", system_prompt),
        ("user", requirement_text)
    ]
    
    try:
        # 直接调用 invoke 获取完整响应
        response = llm.invoke(messages)
        result = response.content.strip() if response.content else ""
        return result
    except Exception as e:
        print(f"[ERROR] 生成过程发生错误: {e}", flush=True)
        return ""