#!/usr/bin/env python3
# -*- coding: utf-8 -*-
"""
BaoDan Workflow 自动化配置脚本

功能：
1. 登录 BaoDan 后台获取 token
2. 创建 Workflow 应用（产品推荐方案生成）
3. 配置 6 个节点的 Workflow 图

用法：
    python scripts/setup_workflow.py

前置条件：
    1. BaoDan 服务已启动（docker compose up -d）
    2. 已创建管理员账号（默认：taiyi@baodan.com / taiyi1224）
    3. 已配置 LLM 模型（在 BaoDan 后台 设置 -> 模型供应商）
"""

import base64
import json
import sys
import time
import requests
from typing import Any

# 修复 Windows 控制台编码
import io
sys.stdout = io.TextIOWrapper(sys.stdout.buffer, encoding='utf-8')
sys.stderr = io.TextIOWrapper(sys.stderr.buffer, encoding='utf-8')

# ===== 配置 =====
BAODAN_URL = "http://localhost:5001"
ADMIN_EMAIL = "taiyi@baodan.com"
ADMIN_PASSWORD = "taiyi1224"

# Workflow 节点 ID（固定值，用于连线）
NODE_IDS = {
    "start": "start",
    "validate": "validate_params",
    "generate_queries": "generate_queries",
    "knowledge_retrieval": "knowledge_retrieval",
    "llm_generate": "llm_generate",
    "error_handler": "error_handler",
    "format_output": "format_output",
    "end": "end",
}


class BaoDanWorkflowSetup:
    """BaoDan Workflow 自动化配置器。"""

    def __init__(self, base_url: str):
        self.base_url = base_url
        self.session = requests.Session()
        self.token = None
        self.csrf_token = None
        self.tenant_id = None
        self.app_id = None

    def login(self, email: str, password: str) -> bool:
        """登录获取 token。"""
        print(f"[1/6] 登录 BaoDan ({email})...")
        # Dify 要求密码 Base64 编码
        encoded_password = base64.b64encode(password.encode('utf-8')).decode('utf-8')
        resp = self.session.post(
            f"{self.base_url}/console/api/login",
            json={"email": email, "password": encoded_password},
        )
        if resp.status_code != 200:
            print(f"  ❌ 登录失败: {resp.text}")
            return False

        data = resp.json()
        if data.get("result") != "success":
            print(f"  ❌ 登录失败: {data}")
            return False

        # Dify 将 token 设置到 cookie 中
        # 从 cookie 中提取 access_token
        cookies = resp.cookies
        self.token = cookies.get("access_token")

        if not self.token:
            # 尝试从 set-cookie 头中提取
            for cookie in resp.headers.get("set-cookie", "").split(";"):
                if "access_token" in cookie:
                    parts = cookie.split("=")
                    if len(parts) == 2:
                        self.token = parts[1].strip()
                        break

        if not self.token:
            print(f"  ❌ 未获取到 token，cookies: {dict(cookies)}")
            return False

        # 提取 CSRF token
        self.csrf_token = cookies.get("csrf_token")
        if not self.csrf_token:
            for cookie in resp.headers.get("set-cookie", "").split(";"):
                if "csrf_token" in cookie:
                    parts = cookie.split("=")
                    if len(parts) == 2:
                        self.csrf_token = parts[1].strip()
                        break

        self.session.headers.update({
            "Authorization": f"Bearer {self.token}",
            "Content-Type": "application/json",
        })

        # 设置 CSRF token 到请求头
        if self.csrf_token:
            self.session.headers.update({
                "X-CSRF-Token": self.csrf_token,
            })

        print(f"  ✅ 登录成功")
        return True

    def find_existing_app(self) -> str | None:
        """查找已存在的同名 Workflow 应用。"""
        page = 1
        while True:
            resp = self.session.get(
                f"{self.base_url}/console/api/apps",
                params={"page": page, "limit": 100}
            )
            if resp.status_code != 200:
                break

            data = resp.json()
            apps = data.get("data", [])
            if not apps:
                break

            for app in apps:
                if app.get("name") == "产品推荐方案生成" and app.get("mode") == "workflow":
                    return app.get("id")

            # 检查是否还有更多页
            if len(apps) < 100:
                break
            page += 1

        return None

    def create_workflow_app(self) -> bool:
        """创建 Workflow 类型应用（如果不存在则复用）。"""
        print("[2/6] 检查/创建 Workflow 应用...")

        # 先检查是否已存在同名应用
        existing_app_id = self.find_existing_app()
        if existing_app_id:
            self.app_id = existing_app_id
            print(f"  ♻️  复用已存在的应用，App ID: {self.app_id}")
            return True

        # 不存在则创建新应用
        resp = self.session.post(
            f"{self.base_url}/console/api/apps",
            json={
                "name": "产品推荐方案生成",
                "description": "根据客户信息和险种需求，调用知识库检索并生成保险产品推荐方案",
                "mode": "workflow",
            },
        )
        if resp.status_code not in (200, 201):
            print(f"  ❌ 创建应用失败: {resp.text}")
            return False

        data = resp.json()
        self.app_id = data.get("id")
        if not self.app_id:
            print(f"  ❌ 未获取到 App ID: {data}")
            return False

        print(f"  ✅ 应用创建成功，App ID: {self.app_id}")
        return True

    def get_workflow_id(self) -> str | None:
        """获取当前 draft workflow ID。"""
        resp = self.session.get(
            f"{self.base_url}/console/api/apps/{self.app_id}/workflows/publish"
        )
        if resp.status_code == 200:
            data = resp.json()
            if data:
                return data.get("workflow", {}).get("id")
        return None

    def build_workflow_graph(self) -> dict:
        """构建 Workflow 节点图。"""
        return {
            "nodes": self._build_nodes(),
            "edges": self._build_edges(),
        }

    def _build_nodes(self) -> list[dict]:
        """构建所有节点配置。"""
        return [
            # Start 节点
            {
                "id": NODE_IDS["start"],
                "type": "custom",
                "data": {
                    "type": "start",
                    "title": "开始",
                    "desc": "",
                    "variables": self._get_start_variables(),
                },
                "position": {"x": 80, "y": 282},
            },
            # 节点1：参数校验
            {
                "id": NODE_IDS["validate"],
                "type": "custom",
                "data": {
                    "type": "code",
                    "title": "参数校验",
                    "desc": "校验客户年龄、性别、职业、预算等参数",
                    "variables": [],
                    "code": self._get_validate_code(),
                    "code_language": "python3",
                    "outputs": {
                        "validated_params": {"type": "object", "children": None},
                        "error_msg": {"type": "string", "children": None},
                        "is_valid": {"type": "boolean", "children": None},
                    },
                },
                "position": {"x": 320, "y": 282},
            },
            # 节点2：检索策略生成
            {
                "id": NODE_IDS["generate_queries"],
                "type": "custom",
                "data": {
                    "type": "code",
                    "title": "检索策略生成",
                    "desc": "为每个险种生成检索词",
                    "variables": [],
                    "code": self._get_generate_queries_code(),
                    "code_language": "python3",
                    "outputs": {
                        "search_queries": {"type": "array[object]", "children": None},
                    },
                },
                "position": {"x": 560, "y": 282},
            },
            # 节点3：知识库检索
            {
                "id": NODE_IDS["knowledge_retrieval"],
                "type": "custom",
                "data": {
                    "type": "knowledge-retrieval",
                    "title": "知识库检索",
                    "desc": "从保险知识库中检索相关产品条款",
                    "query_variable_selector": [NODE_IDS["generate_queries"], "search_queries"],
                    "dataset_ids": [],  # 需要手动关联知识库
                    "retrieval_mode": "multiple",
                    "multiple_retrieval_config": {
                        "top_k": 5,
                        "score_threshold": 0.6,
                        "reranking_model": None,
                    },
                },
                "position": {"x": 800, "y": 282},
            },
            # 节点4：LLM 方案生成
            {
                "id": NODE_IDS["llm_generate"],
                "type": "custom",
                "data": {
                    "type": "llm",
                    "title": "LLM 方案生成",
                    "desc": "调用 LLM 生成三套保险推荐方案",
                    "model": {
                        "provider": "deepseek",
                        "name": "deepseek-chat",
                        "mode": "chat",
                        "completion_params": {
                            "temperature": 0.7,
                            "top_p": 0.9,
                            "max_tokens": 8192,
                        },
                    },
                    "prompt_template": self._get_llm_prompt(),
                    "memory": None,
                },
                "position": {"x": 1040, "y": 282},
            },
            # 节点5：异常处理
            {
                "id": NODE_IDS["error_handler"],
                "type": "custom",
                "data": {
                    "type": "code",
                    "title": "异常处理",
                    "desc": "处理校验失败或生成失败的情况",
                    "variables": [],
                    "code": self._get_error_handler_code(),
                    "code_language": "python3",
                    "outputs": {
                        "final_output": {"type": "string", "children": None},
                        "is_success": {"type": "boolean", "children": None},
                    },
                },
                "position": {"x": 1040, "y": 520},
            },
            # 节点6：格式化输出
            {
                "id": NODE_IDS["format_output"],
                "type": "custom",
                "data": {
                    "type": "code",
                    "title": "格式化输出",
                    "desc": "去除多余空行和空白",
                    "variables": [],
                    "code": self._get_format_output_code(),
                    "code_language": "python3",
                    "outputs": {
                        "recommendation": {"type": "string", "children": None},
                        "is_success": {"type": "boolean", "children": None},
                    },
                },
                "position": {"x": 1280, "y": 282},
            },
            # End 节点
            {
                "id": NODE_IDS["end"],
                "type": "custom",
                "data": {
                    "type": "end",
                    "title": "结束",
                    "desc": "",
                    "outputs": [
                        {
                            "variable_selector": [NODE_IDS["format_output"], "recommendation"],
                            "value_selector": [NODE_IDS["format_output"], "recommendation"],
                        },
                        {
                            "variable_selector": [NODE_IDS["format_output"], "is_success"],
                            "value_selector": [NODE_IDS["format_output"], "is_success"],
                        },
                    ],
                },
                "position": {"x": 1520, "y": 282},
            },
        ]

    def _build_edges(self) -> list[dict]:
        """构建节点连线。"""
        return [
            # Start → 参数校验
            {"source": NODE_IDS["start"], "target": NODE_IDS["validate"]},
            # 参数校验 → is_valid=True → 检索策略生成
            {
                "source": NODE_IDS["validate"],
                "target": NODE_IDS["generate_queries"],
                "sourceHandle": "is_valid",
            },
            # 参数校验 → is_valid=False → 异常处理
            {
                "source": NODE_IDS["validate"],
                "target": NODE_IDS["error_handler"],
                "sourceHandle": "is_invalid",
            },
            # 检索策略生成 → 知识库检索
            {"source": NODE_IDS["generate_queries"], "target": NODE_IDS["knowledge_retrieval"]},
            # 知识库检索 → LLM 方案生成
            {"source": NODE_IDS["knowledge_retrieval"], "target": NODE_IDS["llm_generate"]},
            # LLM 方案生成 → 格式化输出
            {"source": NODE_IDS["llm_generate"], "target": NODE_IDS["format_output"]},
            # 异常处理 → 格式化输出
            {"source": NODE_IDS["error_handler"], "target": NODE_IDS["format_output"]},
            # 格式化输出 → 结束
            {"source": NODE_IDS["format_output"], "target": NODE_IDS["end"]},
        ]

    def _get_start_variables(self) -> list[dict]:
        """获取开始节点的输入变量定义。"""
        return [
            {"variable": "age", "label": "年龄", "type": "string", "required": True},
            {"variable": "gender", "label": "性别", "type": "string", "required": True},
            {"variable": "occupation", "label": "职业", "type": "string", "required": True},
            {"variable": "annual_income", "label": "年收入(万)", "type": "string", "required": True},
            {"variable": "monthly_budget", "label": "月预算(元)", "type": "string", "required": True},
            {"variable": "insurance_types", "label": "险种", "type": "string", "required": True},
            {"variable": "coverage_amount", "label": "保额(元)", "type": "string", "required": True},
            {"variable": "coverage_period", "label": "保障期限", "type": "string", "required": True},
        ]

    def _get_validate_code(self) -> str:
        """节点1：参数校验代码。"""
        return '''def main(age: str, gender: str, occupation: str, annual_income: str,
         monthly_budget: str, insurance_types: str, coverage_amount: str,
         coverage_period: str) -> dict:
    """校验客户参数。"""
    errors = []

    # 校验年龄
    try:
        age_int = int(age)
        if age_int < 1 or age_int > 150:
            errors.append("年龄必须在 1-150 之间")
    except ValueError:
        errors.append("年龄必须为整数")

    # 校验性别
    if gender not in ("male", "female"):
        errors.append("性别必须为 male 或 female")

    # 校验职业
    if not occupation or len(occupation) < 1:
        errors.append("职业不能为空")

    # 校验月预算
    try:
        budget = float(monthly_budget)
        if budget <= 0:
            errors.append("月预算必须大于 0")
    except ValueError:
        errors.append("月预算必须为数字")

    # 校验保额
    try:
        amount = float(coverage_amount)
        if amount < 10000:
            errors.append("保额不能低于 10000 元")
    except ValueError:
        errors.append("保额必须为数字")

    # 校验险种
    valid_types = ["重疾险", "医疗险", "意外险", "寿险", "年金险", "财产险"]
    types_list = [t.strip() for t in insurance_types.split(",") if t.strip()]
    if not types_list:
        errors.append("至少选择一种险种")
    else:
        for t in types_list:
            if t not in valid_types:
                errors.append(f"不支持的险种: {t}")

    # 校验保障期限
    valid_periods = ["10年", "20年", "30年", "至60岁", "至70岁", "终身"]
    if coverage_period not in valid_periods:
        errors.append(f"不支持的保障期限: {coverage_period}")

    if errors:
        return {
            "validated_params": {},
            "error_msg": "; ".join(errors),
            "is_valid": False,
        }

    return {
        "validated_params": {
            "age": age,
            "gender": gender,
            "occupation": occupation,
            "annual_income": annual_income,
            "monthly_budget": monthly_budget,
            "insurance_types": insurance_types,
            "coverage_amount": coverage_amount,
            "coverage_period": coverage_period,
        },
        "error_msg": "",
        "is_valid": True,
    }
'''

    def _get_generate_queries_code(self) -> str:
        """节点2：检索策略生成代码。"""
        return '''def main(validated_params: dict) -> dict:
    """为每个险种生成检索词。"""
    insurance_types = validated_params.get("insurance_types", "")
    age = validated_params.get("age", "")
    occupation = validated_params.get("occupation", "")

    types_list = [t.strip() for t in insurance_types.split(",") if t.strip()]

    search_queries = []
    for ins_type in types_list:
        query = f"{ins_type} 产品条款 保额 费率 {age}岁 {occupation}"
        search_queries.append({
            "insurance_type": ins_type,
            "query": query,
        })

    return {"search_queries": search_queries}
'''

    def _get_llm_prompt(self) -> list[dict]:
        """节点4：LLM Prompt 模板。"""
        return [
            {
                "role": "system",
                "text": """你是一位专业的保险顾问，擅长根据客户需求生成保险产品推荐方案。

重要规则：
1. 保费数据必须来自检索结果，绝对不能编造
2. 推荐方案的总保费不能超过客户月预算 × 12
3. 每套方案必须包含：产品名称、保险公司、险种、保额、年保费、推荐理由
4. 最终输出必须包含免责声明"""
            },
            {
                "role": "user",
                "text": """请根据以下客户信息和检索到的保险产品资料，生成三套保险推荐方案。

## 客户信息
- 年龄：{{#start.age#}} 岁
- 性别：{{#start.gender#}}
- 职业：{{#start.occupation#}}
- 年收入：{{#start.annual_income#}} 万元
- 月预算：{{#start.monthly_budget#}} 元
- 需要险种：{{#start.insurance_types#}}
- 期望保额：{{#start.coverage_amount#}} 元
- 保障期限：{{#start.coverage_period#}}

## 检索到的保险产品资料
{{#knowledge_retrieval.result#}}

## 输出要求
请生成以下三套方案，使用 Markdown 表格格式：

### 基础方案（年保费约 XXXX 元）
| 产品名称 | 所属保险公司 | 险种 | 保额 | 年保费 | 推荐理由 |
|---------|------------|------|------|-------|---------|
| ... | ... | ... | ... | ... | ... |

方案总结：...

### 均衡方案（年保费约 XXXX 元）
（格式同上）

### 全面方案（年保费约 XXXX 元）
（格式同上）

## 免责声明
本推荐方案仅供参考，具体保障内容以保险合同为准。建议在投保前仔细阅读保险条款。"""
            }
        ]

    def _get_error_handler_code(self) -> str:
        """节点5：异常处理代码。"""
        return '''def main(is_valid: bool, error_msg: str, recommendation: str = "") -> dict:
    """处理异常情况。"""
    if not is_valid:
        return {
            "final_output": f"参数校验失败：{error_msg}",
            "is_success": False,
        }

    if not recommendation:
        return {
            "final_output": "未能生成推荐方案，请稍后重试",
            "is_success": False,
        }

    return {
        "final_output": recommendation,
        "is_success": True,
    }
'''

    def _get_format_output_code(self) -> str:
        """节点6：格式化输出代码。"""
        return '''def main(recommendation: str, is_success: bool) -> dict:
    """格式化输出，去除多余空白。"""
    if not is_success:
        return {
            "recommendation": recommendation,
            "is_success": False,
        }

    # 去除多余空行（连续空行合并为一个）
    import re
    formatted = re.sub(r"\\n{3,}", "\\n\\n", recommendation)
    # 去除首尾空白
    formatted = formatted.strip()

    return {
        "recommendation": formatted,
        "is_success": True,
    }
'''

    def sync_workflow(self) -> bool:
        """同步 Workflow 配置到 BaoDan。"""
        print("[3/6] 同步 Workflow 配置...")

        graph = self.build_workflow_graph()

        # 获取 workflow ID
        workflow_id = self.get_workflow_id()
        if not workflow_id:
            print("  ⚠️  未找到 workflow，尝试直接同步...")

        resp = self.session.post(
            f"{self.base_url}/console/api/apps/{self.app_id}/workflows/draft",
            json={
                "graph": graph,
                "features": {},
                "hash": "",
            },
        )

        if resp.status_code != 200:
            print(f"  ❌ 同步失败: {resp.text}")
            return False

        print("  ✅ Workflow 配置同步成功")
        return True

    def get_api_key(self) -> str | None:
        """获取应用的 API Key。"""
        print("[4/6] 获取 API Key...")

        # 先尝试获取已有的 API Key
        resp = self.session.get(
            f"{self.base_url}/console/api/apps/{self.app_id}/api-keys"
        )
        if resp.status_code == 200:
            keys = resp.json().get("data", [])
            if keys:
                api_key = keys[0].get("token")
                print(f"  ✅ API Key: {api_key}")
                return api_key

        # 创建新的 API Key
        resp = self.session.post(
            f"{self.base_url}/console/api/apps/{self.app_id}/api-keys",
            json={"name": "推荐服务"},
        )
        if resp.status_code in (200, 201):
            data = resp.json()
            api_key = data.get("token")
            if api_key:
                print(f"  ✅ 新建 API Key: {api_key}")
                return api_key
            # 可能返回的是完整对象
            api_key = data.get("id")  # 有些版本返回 id 作为 token
            if api_key:
                print(f"  ✅ 新建 API Key: {api_key}")
                return api_key

        print(f"  ❌ 创建 API Key 失败: {resp.text}")
        return None

    def publish_workflow(self) -> bool:
        """发布 Workflow。"""
        print("[5/6] 发布 Workflow...")

        # 先获取 draft workflow
        draft_resp = self.session.get(
            f"{self.base_url}/console/api/apps/{self.app_id}/workflows/draft"
        )
        if draft_resp.status_code == 200:
            draft_data = draft_resp.json()
            workflow_id = draft_data.get("id")
            if workflow_id:
                print(f"  📝 Draft Workflow ID: {workflow_id}")

        # 发布
        resp = self.session.post(
            f"{self.base_url}/console/api/apps/{self.app_id}/workflows/publish",
            json={},
        )
        if resp.status_code in (200, 201):
            print("  ✅ Workflow 已发布")
            return True

        # 尝试不带 body 发布
        resp = self.session.post(
            f"{self.base_url}/console/api/apps/{self.app_id}/workflows/publish"
        )
        if resp.status_code in (200, 201):
            print("  ✅ Workflow 已发布")
            return True

        print(f"  ⚠️  发布失败（可手动在后台发布）: {resp.text}")
        return False

    def update_env_config(self, api_key: str) -> bool:
        """更新 .env 配置文件。"""
        print("[6/6] 更新 .env 配置...")

        env_content = f"""# BaoDan Workflow API Key（自动生成于 {time.strftime('%Y-%m-%d %H:%M:%S')}）
BAODAN_WORKFLOW_API_KEY={api_key}
BAODAN_API_URL={self.base_url}
"""

        # 写入到 docker-compose 环境变量提示
        print(f"""
  ✅ 请将以下配置添加到 docker-compose.dify.yml 的 baodan-api 环境变量：

    BAODAN_WORKFLOW_API_KEY: "{api_key}"

  或者创建 .env 文件：
{env_content}
""")
        return True


def main():
    """主函数。"""
    print("=" * 60)
    print("BaoDan Workflow 自动化配置脚本")
    print("=" * 60)
    print()

    setup = BaoDanWorkflowSetup(BAODAN_URL)

    # 1. 登录
    if not setup.login(ADMIN_EMAIL, ADMIN_PASSWORD):
        print("\n❌ 登录失败，请检查：")
        print("  1. BaoDan 服务是否已启动")
        print("  2. 管理员账号密码是否正确")
        sys.exit(1)

    # 2. 创建 Workflow 应用
    if not setup.create_workflow_app():
        print("\n❌ 创建应用失败")
        sys.exit(1)

    # 3. 同步 Workflow 配置
    if not setup.sync_workflow():
        print("\n❌ 同步 Workflow 失败")
        sys.exit(1)

    # 4. 获取 API Key
    api_key = setup.get_api_key()
    if not api_key:
        print("\n⚠️  获取 API Key 失败，请手动在后台创建")

    # 5. 发布 Workflow
    setup.publish_workflow()

    # 6. 更新配置
    if api_key:
        setup.update_env_config(api_key)

    print("\n" + "=" * 60)
    print("✅ Workflow 配置完成！")
    print("=" * 60)
    print(f"""
后续步骤：
1. 访问 http://localhost:3000/workflow/{setup.app_id} 查看 Workflow
2. 在「知识库检索」节点关联已创建的知识库
3. 在「模型配置」中确认 DeepSeek 模型已配置
4. 测试 Workflow：点击「运行」按钮测试
5. 更新 docker-compose.dify.yml 中的 BAODAN_WORKFLOW_API_KEY

API 调用示例：
curl -X POST http://localhost:5001/v1/workflows/run \\
  -H "Authorization: Bearer {api_key or 'YOUR_API_KEY'}" \\
  -H "Content-Type: application/json" \\
  -d '{{
    "inputs": {{
      "age": "35",
      "gender": "male",
      "occupation": "工程师",
      "annual_income": "30",
      "monthly_budget": "2000",
      "insurance_types": "重疾险,医疗险",
      "coverage_amount": "500000",
      "coverage_period": "终身"
    }},
    "response_mode": "blocking",
    "user": "test_user"
  }}'
""")


if __name__ == "__main__":
    main()