Main

#导入库和参数配置
from hello_agents import SimpleAgent, HelloAgentsLLM, ReflectionAgent, ToolRegistry
from hello_agents.tools import Tool, ToolParameter
from typing import Dict, Any, List
from tavily import TavilyClient
import os
import json
import re
import numpy as np 
from dotenv import load_dotenv
import asyncio
import nest_asyncio
from mcp.client.sse import sse_client
from mcp.client.session import ClientSession

load_dotenv()

#LLM参数
LLM_MODEL_ID = os.getenv("LLM_MODEL_ID")
LLM_API_KEY = os.getenv("LLM_API_KEY")
LLM_BASE_URL = os.getenv("LLM_BASE_URL")
#Tavily参数
TAVILY_API_KEY = os.getenv("TAVILY_API_KEY","")
#百度MCP参数
BAIDU_TOKEN = os.getenv("BAIDU_MCP_TOKEN","")
#输入json路径配置
INPUT_FILENAME = "data/test_cases.json"

# 搜索源配置
# 可选值: "tavily" (通用/海外) 或 "baidu" (电商/国内)
os.environ["SEARCH_PROVIDER"] = "baidu" 

print("✅ 环境配置完成")

# [Cell 2 终极版] 定义统一搜索工具 (兼容 Tavily 和 Baidu)
# 允许 Jupyter 运行异步
nest_asyncio.apply()

class BatchSearchTool(Tool):
    def __init__(self):
        super().__init__(
            name="batch_search",
            description="统一搜索工具，支持 Tavily 和 Baidu 切换。"
        )
        self.provider = os.environ.get("SEARCH_PROVIDER", "tavily").lower()

    def run(self, parameters: Any) -> str:
        return "请使用 Python 代码直接调用 search_raw 方法获取数据。"

    def search_raw(self, query: str) -> List[Dict]:
        if self.provider == "baidu":
            return self._search_baidu(query)
        else:
            return self._search_tavily(query)

    # --- 引擎 A: Tavily ---
    def _search_tavily(self, query: str) -> List[Dict]:
        api_key = os.environ.get("TAVILY_API_KEY")
        if not api_key: return []
        print(f"    🚀 [Tavily] 正在搜索: {query} ...")
        try:
            tavily = TavilyClient(api_key=api_key)
            response = tavily.search(query, max_results=5, include_images=True)
            results = []
            if 'results' in response:
                for r in response['results']:
                    results.append({
                        "title": r['title'], "url": r['url'], "content": r['content'], 
                        "type": "text", "img": "" # Tavily 文本通常不带图
                    })
            if 'images' in response and response['images']:
                results.append({"images": response['images'][:3], "type": "image"})
            return results
        except Exception as e:
            print(f"      ⚠️ Tavily 异常: {e}")
            return []

    # --- 引擎 B: Baidu MCP ---
    def _search_baidu(self, query: str) -> List[Dict]:
        token = os.environ.get("BAIDU_MCP_TOKEN")
        if not token: return []
        print(f"    🐼 [百度优选] 正在搜索: {query} ...")
        try:
            raw_json_str = asyncio.run(self._async_baidu_call(query, token))
            print(f"      🔍 原始 JSON 响应: {raw_json_str}")
            return self._parse_baidu_response(raw_json_str)
        except Exception as e:
            print(f"      ⚠️ 百度 MCP 异常: {e}")
            return []

    async def _async_baidu_call(self, query: str, token: str) -> str:
        sse_url = f"https://mcp-youxuan.baidu.com/mcp/sse?key={token}"
        async with sse_client(sse_url) as (read, write):
            async with ClientSession(read, write) as session:
                await session.initialize()
                result = await session.call_tool("goods_search", arguments={"query": query})
                return result.content[0].text if result.content else ""

    def _parse_baidu_response(self, json_str: str) -> List[Dict]:
        results = []; images = []
        try:
            data = json.loads(json_str)
            items = data if isinstance(data, list) else []
            
            for item in items[:5]:
                title = item.get("goodsName") or item.get("title") or "未知商品"
                price = item.get("price") or item.get("minPrice") or ""
                shop = item.get("shopName") or item.get("mall") or ""
                url = item.get("detailUrl") or item.get("url") or item.get("ori_url") or "#"
                img = item.get("imgUrl") or item.get("picUrl") or item.get("img")
                
                content = f"价格: {price}元。店铺: {shop}。商品详情: {title}"
                
                # 📝【修复点】直接在 text 类型结果里绑定 img
                results.append({
                    "title": title, "url": url, "content": content, 
                    "type": "text", "img": img 
                })
                if img: images.append(img)
            
            if images: results.append({"images": images[:3], "type": "image"})
                
        except json.JSONDecodeError:
            print("      ⚠️ 百度返回非 JSON 数据")
        return results

    def get_parameters(self):
        return [ToolParameter(name="query", type="string", description="关键词")]

tool_registry = ToolRegistry()
tool_registry.register_tool(BatchSearchTool())

print("✅ 统一搜索工具已加载！")
print(f"当前模式: {'百度优选 (电商)' if os.environ.get('SEARCH_PROVIDER') == 'baidu' else 'Tavily (通用)'}")

# 初始化大模型
llm = HelloAgentsLLM()

# --- 1. 军师 (Profiler) - 已升级支持多维度画像 ---
PROFILER_PROMPT = """
你是一个精通 MBTI 人格分析与消费市场趋势的 "送礼军师"。
你的任务是根据用户提供的多维度画像，制定 3 个**极度精准**的搜索关键词。

【⚠️ 时效性死命令 (CRITICAL)】
当前时间视作 **2025年11月**。
1. **严禁过时**：绝对不要推荐 2024 年或更早的旧款（除非是经典恒久款如黑胶唱片）。
2. **价格尺度**：给出的价格单位是人民币元。请严格遵照范围进行联想，禁止超出预算范围。

为了确保推荐质量，请参考以下的【优秀思考范例】：

### 范例 1
**用户画像**: 
- 女, 26岁, ISFP (探险家), 金牛座
- 预算: 500-1000元
- 场景: 情人节
- 自定义: 喜欢有质感的生活小物
**军师分析**: 
ISFP 重视审美和感官体验，金牛座喜欢实实在在的质感。情人节需要浪漫。
**生成策略**:
1. 观夏 (To Summer) 昆仑煮雪 晶石香薰 (符合质感与审美)
2. 野兽派 2025 情人节限定 睡衣礼盒 (金牛座喜欢的舒适)
3. 富士 Instax mini Evo 拍立得 (记录生活瞬间)

### 范例 2
**用户画像**: 
- 男, 30岁, INTJ (建筑师), 处女座
- 预算: 1000元以上
- 场景: 生日
- 自定义: 程序员，喜欢整洁
**军师分析**: 
INTJ 追求极致的逻辑和效率，处女座有洁癖，喜欢桌面整洁。
**生成策略**:
1. Keychron Q1 Pro 机械键盘 铝坨坨 (符合极客对工具的追求)
2. 明基 (BenQ) ScreenBar Halo 屏幕挂灯 (极致护眼与桌面美学)
3. 赫曼米勒 (Herman Miller) 显示器支架 (人体工学)

### 范例 3
**用户画像**: 
- 女, 20岁, ENFP (竞选者), 狮子座
- 预算: 300元以内
- 场景: 圣诞节
- 自定义: 喜欢二次元，痛包
**军师分析**: 
ENFP 热情奔放，狮子座喜欢张扬、闪亮的东西。预算有限但要素多。
**生成策略**:
1. 泡泡玛特 圣诞系列 盲盒整端 (符合节日气氛和二次元)
2. WEGO 痛包 镭射款 (符合自定义需求，狮子座喜欢的亮眼)
3.  Chiikawa 吉伊卡哇 圣诞公仔 (当下顶流二次元IP)

---

**现在的任务**：
请根据以下【当前用户画像】进行分析，模仿上述范例的深度，制定搜索策略。

【当前用户画像】
{user_profile_text}

【关键词生成要求】
1. **必须具体**：格式为 `[品牌] + [产品名/系列] + [限定/属性]`。
2. **拒绝大词**：严禁搜索 "礼物"、"口红"、"玩具" 这种泛词。
3. **必须包含品牌**：根据预算推断合适的品牌（如：预算低选名创优品/泡泡玛特，预算高选Dior/索尼）。

【输出格式】
只输出 3 行搜索关键词，每行一个。不要输出分析过程，不要序号。
"""

profiler_agent = SimpleAgent(
    llm=llm,
    name="Agent_Profiler",
    system_prompt=PROFILER_PROMPT
)

# ==============================================================================
# 2. 种草达人 (Pitcher) - 文案创作 (加入风格指导)
# ==============================================================================
PITCHER_PROMPT = """
你是一个 **金牌种草文案**。
用户会给你一个 **【商品名称】**。

### 🎯 关键要点 (Few Points)
1.  **痛点直击**: 一句话说清楚为什么买它（限定？显白？绝美？）。
2.  **情绪价值**: 使用 "绝绝子", "氛围感", "心动" 等高频热词。
3.  **字数限制**: 严格控制在 **40字以内**，短小精悍。
4.  **Emoji**: 必须包含 1-2 个 emoji。

### 🌟 创作范例 (Few-Shot)
**输入**: Dior 999 烈艳蓝金
**输出**: 💄本宫不死终是妃！Dior 999 传奇正红，显白更有气场，送女友绝对没错！

**输入**: 泡泡玛特 Labubu 坐坐派对
**输出**: ✨太可爱了吧！Labubu 坐坐派对系列，每一个都丑萌到心巴上，摆在桌上超治愈~

**输入**: 罗技 MX Master 3S
**输出**: 🖱️打工人本命！罗技 Master 3S 静音又顺滑，人体工学设计，手腕再也不累了。

---

**当前任务**：
请为【{input}】写一句朋友圈风格种草语。
"""
pitcher_agent = SimpleAgent(
    llm=llm,
    name="Agent_Pitcher",
    system_prompt=PITCHER_PROMPT
)

print("✅ 智能体初始化完成！")

def load_user_profile(filename):
    # 1. 检查文件是否存在
    if not os.path.exists(filename):
        print(f"⚠️ 未找到配置文件: {filename}")
        # 如果没有文件，将默认数据写入文件
        default_data = {
            "性别": "女",
            "年龄": "24岁",
            "MBTI": "ENFP",
            "星座": "天秤座",
            "预算": "500元以内",
            "节日": "恋爱一周年纪念日",
            "自定义": "喜欢二次元，平时喜欢喝咖啡，不要送太实用的家电"
        }
        with open(filename, "w", encoding="utf-8") as f:
            json.dump(default_data, f, ensure_ascii=False, indent=4)
        print(f"✅ 已自动生成默认配置文件，请修改 {filename} 后再次运行。")
        return default_data

    # 2. 读取文件内容
    try:
        with open(filename, "r", encoding="utf-8") as f:
            data = json.load(f)
        print(f"✅ 成功加载用户画像: {filename}")
        print(f"📋 内容预览: {json.dumps(data, ensure_ascii=False)}")
        return data
    except Exception as e:
        print(f"❌ 读取 JSON 失败: {e}")
        return {}

# 加载数据
user_input_data = load_user_profile(INPUT_FILENAME)

def parse_budget_range(budget_str):
    """解析用户预算字符串，返回 (min, max)"""
    nums = [float(x) for x in re.findall(r'\d+', str(budget_str).replace(',', ''))]
    if not nums: return 0, 999999 
    if "以内" in budget_str or "以下" in budget_str: return 0, nums[0]
    if "以上" in budget_str: return nums[0], 999999
    if len(nums) >= 2: return min(nums), max(nums)
    return 0, nums[0]

def extract_all_prices(raw_results):
    """从搜索结果列表中提取所有有效的价格"""
    prices = []
    for res in raw_results:
        # 只处理文本类型的结果
        if res.get('type') == 'text':
            text = res.get('title', '') + " " + res.get('content', '')
            # 匹配 ¥, $, 元 等格式
            matches = re.findall(r'(?:¥|￥|\$|HK\$|NT\$)\s*(\d+(?:,\d{3})*(?:\.\d+)?)', text)
            for m in matches:
                val = float(m.replace(',', ''))
                # 过滤掉像年份(2025)或过小/过大的异常值
                if 10 < val < 100000 and val not in [2024, 2025, 2026]:
                    prices.append(val)
            # 备用正则：匹配 "xxx元"
            matches_yuan = re.findall(r'(\d+(?:,\d{3})*(?:\.\d+)?)\s*元', text)
            for m in matches_yuan:
                val = float(m.replace(',', ''))
                if 10 < val < 100000 and val not in [2024, 2025, 2026]:
                    prices.append(val)
    return prices


def find_best_product(hunter, profiler_agent, keyword, budget_min, budget_max):
    limit_upper = budget_max * 1.2
    limit_lower = budget_min * 0.8
    
    all_candidates = []
    current_kw = keyword
    
    # --- Round 1: 首次搜索 ---
    print(f"       🕵️ 第1次搜索: {current_kw} 价格")
    results_1 = hunter.search_raw(f"{current_kw} 价格")
    
    fallback_img = ""
    for r in results_1:
        if r.get('images'): 
            fallback_img = r['images'][0]
            break

    has_valid_info = False
    for res in results_1:
        if res.get('type') == 'text':
            has_valid_info = True
            p_vals = extract_all_prices([res])
            if p_vals:
                res['price_val'] = p_vals[0]
                # 📝【核心修复点1】记录当前结果所属的关键词
                res['source_kw'] = current_kw 
                all_candidates.append(res)
                
                if limit_lower <= p_vals[0] <= limit_upper:
                    if not res.get('img') and fallback_img: res['img'] = fallback_img
                    return res, f"约 {p_vals[0]}元", current_kw

    # --- 机制 3: 无数据防御 ---
    if not has_valid_info:
        print(f"       ⚠️ [机制3触发] 首次搜索无有效信息。")
        correction_prompt = f"原策略 '{current_kw}' 搜索结果为空，请推荐一个同品类但更热门的具体商品型号。只输出关键词。"
        new_kw = profiler_agent.run(correction_prompt).strip()
        print(f"       🔄 军师换词: {new_kw}")
        current_kw = new_kw
        
        results = hunter.search_raw(f"{current_kw} 价格")
        
        # 更新 fallback_img
        fallback_img = "" 
        for r in results:
            if r.get('images'): 
                fallback_img = r['images'][0]
                break
                
        for res in results:
            if res.get('type') == 'text':
                p_vals = extract_all_prices([res])
                if p_vals:
                    res['price_val'] = p_vals[0]
                    # 📝【核心修复点1】记录关键词
                    res['source_kw'] = current_kw
                    all_candidates.append(res)

    # --- 机制 1 & 2: 价格修正 ---
    avg_price = np.mean([c['price_val'] for c in all_candidates]) if all_candidates else 0
    
    if avg_price > 0:
        correction_prompt = ""
        if avg_price > limit_upper:
            print(f"       💸 [机制1触发] 均价 {int(avg_price)} > 上限 {int(limit_upper)}，找平替...")
            correction_prompt = f"原策略 '{current_kw}' 均价约 {int(avg_price)}元，超预算 ({budget_max}元)。请推荐一个同品类更便宜的具体型号（平替）。只输出关键词。"
        elif avg_price < limit_lower:
            print(f"       📉 [机制2触发] 均价 {int(avg_price)} < 下限 {int(limit_lower)}，找升级款...")
            correction_prompt = f"原策略 '{current_kw}' 均价约 {int(avg_price)}元，低于预算下限 ({budget_min}元)。请推荐一个同品类更高端的型号。只输出关键词。"
            
        if correction_prompt:
            new_kw = profiler_agent.run(correction_prompt).strip()
            print(f"       🔄 军师修正: {new_kw}")
            current_kw = new_kw
            
            results_2 = hunter.search_raw(f"{new_kw} 价格")
            
            # 更新 fallback_img
            fallback_img = "" 
            for r in results_2:
                if r.get('images'): 
                    fallback_img = r['images'][0]
                    break
            
            for res in results_2:
                if res.get('type') == 'text':
                    p_vals = extract_all_prices([res])
                    if p_vals:
                        res['price_val'] = p_vals[0]
                        # 📝【核心修复点1】记录关键词
                        res['source_kw'] = current_kw
                        all_candidates.append(res) 
                        
                        if limit_lower <= p_vals[0] <= limit_upper:
                            if not res.get('img') and fallback_img: res['img'] = fallback_img
                            tag = "(平替)" if avg_price > limit_upper else "(升级)"
                            return res, f"约 {p_vals[0]}元 {tag}", current_kw

    # --- 机制 4: 兜底防御 ---
    print("       ⚠️ [机制4触发] 启用强制兜底模式...")
    best_fallback = None
    status_msg = "暂无报价"
    
    if all_candidates:
        # 选离预算最近的
        target = (budget_min + budget_max) / 2
        best_fallback = sorted(all_candidates, key=lambda x: abs(x['price_val'] - target))[0]
        p = best_fallback['price_val']
        
        if p > limit_upper: status_msg = f"约 {p}元 (⚠️超预算)"
        elif p < limit_lower: status_msg = f"约 {p}元 (📉低于预算)"
        else: status_msg = f"约 {p}元"
        
    elif results_1:
        # 实在没数据，硬取第一条
        for res in results_1:
            if res.get('type') == 'text': 
                best_fallback = res
                # 兜底时如果也没价格，就用原始关键词
                best_fallback['source_kw'] = keyword 
                break
    
    if best_fallback:
        if not best_fallback.get('img') and fallback_img:
            best_fallback['img'] = fallback_img
        
        # 📝【核心修复点2】返回结果里记录的那个 source_kw，而不是当前的 current_kw
        final_name_to_use = best_fallback.get('source_kw', current_kw)
        
        return best_fallback, status_msg, final_name_to_use
        
    return None, "搜索失败", keyword

if not user_input_data:
    print("❌ 未加载用户数据")
else:
    # 0. 解析预算
    b_min, b_max = parse_budget_range(user_input_data.get('预算', ''))
    print(f"\n💰 预算范围: {b_min} - {b_max}元")

    # 1. 军师制定策略
    profile_text = "\n".join([f"- {k}: {v if v else '未知/不限'}" for k, v in user_input_data.items()])
    print(f"\n🚀 任务启动...\n{'-'*40}")
    print("\n🧠 [1/3] 军师正在制定初步策略...")
    search_strategy = profiler_agent.run(f"请根据以下用户画像制定搜索策略：\n\n{profile_text}")
    print(f"📝 策略: \n{search_strategy}")

    # 2. 准备循环
    keywords = [k.strip() for k in search_strategy.replace("，", ",").replace("\n", ",").split(',') if k.strip()]
    final_items = []
    hunter = BatchSearchTool()

    print(f"\n🔄 进入处理流程 (共 {len(keywords)} 个商品)...")

    for index, kw in enumerate(keywords):
        print(f"\n    👉 [商品 {index+1}/{len(keywords)}] 正在处理: {kw}")
        
        # 调用智能搜索函数 (传入 min 和 max)
        valid_result, price_status, final_kw = find_best_product(hunter, profiler_agent, kw, b_min, b_max)
        
        if not valid_result:
            print("       ❌ 彻底无数据，跳过。")
            continue
            
        # === 生成文案 ===
        product_name = valid_result.get('title', final_kw)
        
        print(f"       ✍️ 正在撰写文案: {product_name[:30]}...")
        pitch_prompt = f"""
        商品：{product_name}
        价格：{price_status}
        卖点片段：{valid_result.get('content', '')[:200]}...
        
        请写一句30字以内的种草文案。
        """
        pitch = pitcher_agent.run(pitch_prompt)
        
        final_items.append({
            "name": final_kw, 
            "title_full": product_name,
            "price": price_status,
            "desc": pitch.replace("\n", " ").strip(),
            "img": valid_result.get('img', ''),
            "link": valid_result.get('url', '')
        })
        print(f"       ✅ 已收录 (状态: {price_status})")

# --- 4. 渲染与保存 ---
print(f"\n💾 正在生成最终报告...")

if not final_items:
    final_md = "很抱歉，网络搜索似乎出现了问题，未能获取到任何商品信息。"
else:
    table_header = "| 🎁 礼物名称 | 💰 价格 | ✨ 种草理由 | 🖼️ 图片/链接 |\n| :--- | :--- | :--- | :--- |\n"
    table_rows = []
    
    for item in final_items:
        # 1. 清洗文本字段 (防止 | 破坏表格)
        name = item.get('name', '未知').replace("|", "/")
        price = item.get('price', '暂无').replace("|", "/")
        desc = item.get('desc', '').replace("|", "/")
        
        # 2. 🚨【核心修复】清洗链接中的竖线
        # 百度/京东链接常包含 '|'，必须替换为 '%7C'，否则 Markdown 表格会炸
        raw_link = item.get('link', '#')
        safe_link = raw_link.replace("|", "%7C")
        
        raw_img = item.get('img', '')
        safe_img = raw_img.replace("|", "%7C")
        
        # 3. 构建媒体列
        if safe_img and safe_img.startswith("http"):
            # 图片链接套购买链接
            media = f"[![图]({safe_img})]({safe_link})"
        else:
            media = f"[点击购买]({safe_link})"
        
        # 4. 组装行 (注意名字上的链接也要用 safe_link)
        # 使用 strip() 去除可能的首尾空格
        row = f"| [{name}]({safe_link}) | {price} | {desc} | {media} |"
        table_rows.append(row)
        
    final_md = table_header + "\n".join(table_rows)
filename = "outputs/gift_plan_output.md"
# 确保输出目录存在
os.makedirs(os.path.dirname(filename), exist_ok=True)

with open(filename, "w", encoding="utf-8") as f:
    f.write(final_md)
print(f"🎉 任务完成！文件已保存: {os.path.abspath(filename)}")