#!/usr/bin/env python3
"""
travelagent67 v3 — Native Function Calling + Real Amap API
Uses OpenAI tools/function_calling protocol (not text parsing).
"""
import json
import os
import re
import sys
import time
import urllib.request
import urllib.parse
from datetime import datetime
from flask import Flask, request, jsonify

app = Flask(__name__)

SILICONFLOW_URL = "https://api.siliconflow.cn/v1/chat/completions"
# 设置: export SILICONFLOW_API_KEY=...  (申请: https://cloud.siliconflow.cn/account/ak)
SILICONFLOW_KEY = os.environ.get("SILICONFLOW_API_KEY", "")
MODEL = "Qwen/Qwen2.5-32B-Instruct"
# 设置: export AMAP_MCP_KEY=...  (申请: https://lbs.amap.com/dev/key/app)
AMAP_KEY = os.environ.get("AMAP_MCP_KEY", "")
AMAP_BASE = "https://restapi.amap.com/v3"
HISTORY_FILE = os.environ.get("HISTORY_FILE", "/app/qademo/travel_history.json")

# ══════════════════════════════════════════
# Amap Tools (real API)
# ══════════════════════════════════════════

def amap_text_search(keywords, city="", types=""):
    params = f"key={AMAP_KEY}&keywords={urllib.parse.quote(keywords)}&output=json"
    if city: params += f"&city={urllib.parse.quote(city)}"
    if types: params += f"&types={urllib.parse.quote(types)}"
    with urllib.request.urlopen(f"{AMAP_BASE}/place/text?{params}", timeout=10) as r:
        data = json.loads(r.read())
        return [{"name": p.get("name",""), "address": p.get("address",""),
                 "type": p.get("type",""), "location": p.get("location",""),
                 "tel": p.get("tel","")} for p in data.get("pois",[])[:5]]

def amap_around_search(location, keywords="", types="", radius=3000):
    params = f"key={AMAP_KEY}&location={location}&radius={radius}&output=json"
    if keywords: params += f"&keywords={urllib.parse.quote(keywords)}"
    if types: params += f"&types={urllib.parse.quote(types)}"
    with urllib.request.urlopen(f"{AMAP_BASE}/place/around?{params}", timeout=10) as r:
        data = json.loads(r.read())
        return [{"name": p.get("name",""), "address": p.get("address",""),
                 "distance": p.get("distance",""), "location": p.get("location","")}
                for p in data.get("pois",[])[:5]]

def amap_geocode(address, city=""):
    params = f"key={AMAP_KEY}&address={urllib.parse.quote(address)}&output=json"
    if city:
        params += f"&city={urllib.parse.quote(city)}"
    with urllib.request.urlopen(f"{AMAP_BASE}/geocode/geo?{params}", timeout=10) as r:
        data = json.loads(r.read())
        g = data.get("geocodes",[{}])[0]
        return {"address": g.get("formatted_address",""), "location": g.get("location",""),
                "city": g.get("city",""), "district": g.get("district","")}

def amap_weather(city):
    params = f"key={AMAP_KEY}&city={urllib.parse.quote(city)}&extensions=all&output=json"
    with urllib.request.urlopen(f"{AMAP_BASE}/weather/weatherInfo?{params}", timeout=10) as r:
        data = json.loads(r.read())
        f = data.get("forecasts",[{}])[0]
        return {"city": f.get("city",""), "forecasts": [
            {"date": c.get("date",""), "day": c.get("dayweather",""),
             "night": c.get("nightweather",""), "daytemp": c.get("daytemp",""),
             "nighttemp": c.get("nighttemp","")} for c in f.get("casts",[])[:4]]}

def amap_driving(origin, destination):
    params = f"key={AMAP_KEY}&origin={origin}&destination={destination}&output=json"
    with urllib.request.urlopen(f"{AMAP_BASE}/direction/driving?{params}", timeout=10) as r:
        data = json.loads(r.read())
        p = data.get("route",{}).get("paths",[{}])[0]
        return {"distance_m": p.get("distance",""), "duration_s": p.get("duration",""),
                "strategy": p.get("strategy","")}

def amap_district(keywords):
    params = f"key={AMAP_KEY}&keywords={urllib.parse.quote(keywords)}&subdistrict=1&output=json"
    with urllib.request.urlopen(f"{AMAP_BASE}/config/district?{params}", timeout=10) as r:
        data = json.loads(r.read())
        d = data.get("districts",[{}])[0]
        return {"name": d.get("name",""), "center": d.get("center",""),
                "level": d.get("level",""),
                "sub": [s.get("name","") for s in d.get("districts",[])[:10]]}

# Tool dispatch — args may be dict or JSON string (LLM sometimes sends string)
def _ensure_dict(a):
    if isinstance(a, str):
        try:
            return json.loads(a)
        except:
            return {"input": a}
    return a

TOOL_FNS = {
    "maps_text_search": lambda a: amap_text_search(_ensure_dict(a).get("keywords",""), _ensure_dict(a).get("city",""), _ensure_dict(a).get("types","")),
    "maps_around_search": lambda a: amap_around_search(_ensure_dict(a).get("location",""), _ensure_dict(a).get("keywords",""), _ensure_dict(a).get("types",""), _ensure_dict(a).get("radius",3000)),
    "maps_geo": lambda a: amap_geocode(_ensure_dict(a).get("address",""), _ensure_dict(a).get("city","")),
    "maps_weather": lambda a: amap_weather(_ensure_dict(a).get("city","")),
    "maps_driving": lambda a: amap_driving(_ensure_dict(a).get("origin",""), _ensure_dict(a).get("destination","")),
    "maps_district_search": lambda a: amap_district(_ensure_dict(a).get("keywords","")),
}

# OpenAI tools schema for function calling
TOOLS_SCHEMA = [
    {"type": "function", "function": {
        "name": "maps_text_search",
        "description": "按关键词搜索POI（景点、餐厅、酒店等），返回名称、地址、经纬度",
        "parameters": {"type": "object", "properties": {
            "keywords": {"type": "string", "description": "搜索关键词，如'北京 景点'"},
            "city": {"type": "string", "description": "城市名称"},
            "types": {"type": "string", "description": "POI类型: 风景名胜|餐饮服务|住宿服务|购物服务|交通设施"},
        }, "required": ["keywords"]}}},
    {"type": "function", "function": {
        "name": "maps_around_search",
        "description": "在指定坐标附近搜索POI",
        "parameters": {"type": "object", "properties": {
            "location": {"type": "string", "description": "中心坐标，格式'经度,纬度'"},
            "keywords": {"type": "string", "description": "搜索关键词"},
            "types": {"type": "string", "description": "POI类型"},
            "radius": {"type": "integer", "description": "搜索半径(米)，默认3000"},
        }, "required": ["location"]}}},
    {"type": "function", "function": {
        "name": "maps_geo",
        "description": "将地址转换为经纬度坐标。务必传入city参数以确保定位到正确城市",
        "parameters": {"type": "object", "properties": {
            "address": {"type": "string", "description": "地址，如'故宫博物院'"},
            "city": {"type": "string", "description": "所在城市，如'北京'。必填，防止同名地点定位到错误城市"},
        }, "required": ["address", "city"]}}},
    {"type": "function", "function": {
        "name": "maps_weather",
        "description": "查询城市天气预报",
        "parameters": {"type": "object", "properties": {
            "city": {"type": "string", "description": "城市名称，如'北京'"},
        }, "required": ["city"]}}},
    {"type": "function", "function": {
        "name": "maps_driving",
        "description": "查询两点间驾车路线、距离和耗时",
        "parameters": {"type": "object", "properties": {
            "origin": {"type": "string", "description": "起点坐标'经度,纬度'"},
            "destination": {"type": "string", "description": "终点坐标'经度,纬度'"},
        }, "required": ["origin", "destination"]}}},
    {"type": "function", "function": {
        "name": "maps_district_search",
        "description": "查询行政区划信息",
        "parameters": {"type": "object", "properties": {
            "keywords": {"type": "string", "description": "城市或区域名称"},
        }, "required": ["keywords"]}}},
]

SYSTEM_PROMPT = """你是一位专业的旅行规划助手。请根据用户请求，使用提供的地图工具查询真实POI和路线，生成详细行程计划。

【规划策略】
1. 先用 maps_weather 查天气，maps_district_search 了解区划
2. 搜索景点: maps_text_search(keywords="故宫 天坛 颐和园", city="北京", types="风景名胜")
   搜索美食: maps_text_search(keywords="北京烤鸭", city="北京", types="餐饮服务")
   搜索酒店: maps_text_search(keywords="酒店", city="北京", types="住宿服务")
   注意: types 必须用以下标准值之一: 风景名胜|餐饮服务|住宿服务|购物服务
   搜索景点时，keywords 要包含具体知名景点名称，不要只写泛泛的"景点"
   多次搜索不同类型的景点，确保覆盖充足：
   - 第一次: 核心景点(故宫、天坛等)
   - 第二次: 特色景点(南锣鼓巷、798等)
   - 第三次: 自然景点(颐和园、北海公园等)
3. 用 maps_geo 验证关键地点坐标，**必须传 city 参数**防止定位到其他城市的同名地点
4. 编排行程后，用 maps_driving 验证相邻景点的通勤时间
5. 所有地点必须来自工具返回的真实数据，严禁编造
6. 每天安排：上午景点→午餐→下午景点→晚餐→住宿，至少5项
7. **严禁重复**：不同天的行程中不能出现相同的地点名称

【最终输出】
规划完成后，直接输出纯JSON（不要代码块），格式：
{"title":"行程标题","estimatedCost":1500.0,"totalDays":3,"dailyPlans":[{"dayNumber":1,"items":[{"type":"景点|餐饮|住宿|购物","name":"地点名","description":"描述","location":{"lat":39.9,"lng":116.4,"name":"位置"}}]}]}"""


# ══════════════════════════════════════════
# ReAct with native function calling
# ══════════════════════════════════════════

def llm_call_with_tools(messages, tools=None, max_tokens=4096):
    """Call LLM with function calling support"""
    body = {
        "model": MODEL,
        "messages": messages,
        "temperature": 0.2,
        "max_tokens": max_tokens,
    }
    if tools:
        body["tools"] = tools
        body["tool_choice"] = "auto"

    data = json.dumps(body, ensure_ascii=False).encode('utf-8')
    req = urllib.request.Request(
        SILICONFLOW_URL, data=data,
        headers={"Content-Type": "application/json", "Authorization": f"Bearer {SILICONFLOW_KEY}"},
        method="POST"
    )
    with urllib.request.urlopen(req, timeout=120) as resp:
        return json.loads(resp.read())


def plan_travel(query, max_steps=25):
    """ReAct loop with native function calling"""
    messages = [
        {"role": "system", "content": SYSTEM_PROMPT},
        {"role": "user", "content": query},
    ]
    tool_log = []

    for step in range(max_steps):
        resp = llm_call_with_tools(messages, tools=TOOLS_SCHEMA)
        msg = resp["choices"][0]["message"]

        # Check for tool calls
        tool_calls = msg.get("tool_calls")

        if tool_calls:
            # Append assistant message with tool_calls
            messages.append(msg)

            # Execute each tool call
            for tc in tool_calls:
                fn_name = tc["function"]["name"]
                fn_args = json.loads(tc["function"]["arguments"])

                if fn_name in TOOL_FNS:
                    try:
                        result = TOOL_FNS[fn_name](fn_args)
                        result_str = json.dumps(result, ensure_ascii=False)
                        if len(result_str) > 3000:
                            result_str = result_str[:3000] + "...[truncated]"
                    except Exception as e:
                        result_str = json.dumps({"error": str(e)})
                else:
                    result_str = json.dumps({"error": f"Unknown tool: {fn_name}"})

                tool_log.append({"step": step+1, "tool": fn_name, "args": fn_args, "result_preview": result_str[:200]})

                # Add tool result as tool message
                messages.append({
                    "role": "tool",
                    "tool_call_id": tc["id"],
                    "content": result_str,
                })
        else:
            # No tool calls — LLM returned final content
            content = msg.get("content", "")
            # Try to extract JSON
            try:
                # Strip code fences
                clean = content.strip()
                if clean.startswith("```"):
                    clean = clean.split("\n", 1)[1] if "\n" in clean else clean[3:]
                    if clean.endswith("```"):
                        clean = clean[:-3]
                    clean = clean.strip()
                    if clean.startswith("json"):
                        clean = clean[4:].strip()
                itinerary = json.loads(clean)
                return itinerary, tool_log, step + 1
            except:
                m = re.search(r'\{[\s\S]*"dailyPlans"[\s\S]*\}', content)
                if m:
                    try:
                        return json.loads(m.group()), tool_log, step + 1
                    except:
                        pass
            # Ask to finalize
            messages.append({"role": "user", "content": "请直接输出最终行程JSON，不要其他文字。"})

    return None, tool_log, max_steps


def _postprocess_itinerary(itinerary):
    """Post-process: deduplicate items, validate structure"""
    if not itinerary or "dailyPlans" not in itinerary:
        return itinerary

    seen_names = set()
    for day in itinerary.get("dailyPlans", []):
        deduped = []
        for item in day.get("items", []):
            name = item.get("name", "")
            if name and name not in seen_names:
                seen_names.add(name)
                deduped.append(item)
        day["items"] = deduped

    return itinerary


# ══════════════════════════════════════════
# API
# ══════════════════════════════════════════

@app.route('/v1/chat/completions', methods=['POST'])
def chat_completions():
    data = request.json or {}
    msgs = data.get('messages', [])
    if not msgs: return jsonify({'error': 'messages required'}), 400
    user_msg = next((m['content'] for m in msgs if m['role'] == 'user'), '')
    t0 = time.time()
    try:
        itinerary, log, steps = plan_travel(user_msg)
        itinerary = _postprocess_itinerary(itinerary)
        content = json.dumps(itinerary, ensure_ascii=False) if itinerary else "Failed"
        return jsonify({"id": f"travel-{int(time.time()*1000)}", "object": "chat.completion",
            "model": "travelagent67-v3-react", "choices": [{"index": 0, "message": {"role": "assistant", "content": content}, "finish_reason": "stop"}],
            "usage": {"steps": steps, "tools_called": len(log)}, "tool_log": log})
    except Exception as e:
        return jsonify({"error": {"message": str(e)}}), 500

@app.route('/api/travel/plan', methods=['POST'])
def api_plan():
    data = request.json or {}
    query = data.get('query', '').strip()
    if not query: return jsonify({'error': 'query required'}), 400
    t0 = time.time()
    try:
        itinerary, log, steps = plan_travel(query)
        itinerary = _postprocess_itinerary(itinerary)
        return jsonify({'status': 'done', 'query': query, 'itinerary': itinerary,
            'elapsed': round(time.time()-t0, 1), 'steps': steps, 'tools_called': len(log), 'tool_log': log})
    except Exception as e:
        return jsonify({'status': 'error', 'error': str(e)}), 500

@app.route('/api/health')
def health():
    try:
        amap_weather("北京")
        amap_ok = True
    except:
        amap_ok = False
    return jsonify({'status': 'ok', 'agent': 'travelagent67-v3-react',
        'model': MODEL, 'amap_api': 'connected' if amap_ok else 'error',
        'function_calling': True, 'tools': list(TOOL_FNS.keys())})

@app.route('/')
def index():
    return jsonify({'service': 'travelagent67-v3 (Native Function Calling + Amap API)',
        'tools': list(TOOL_FNS.keys()), 'function_calling': True})

if __name__ == '__main__':
    print("travelagent67 v3 (Native Function Calling)")
    print(f"Tools: {list(TOOL_FNS.keys())}")
    try:
        w = amap_weather("北京")
        print(f"Amap: OK ({w['city']} {w['forecasts'][0]['day']})")
    except Exception as e:
        print(f"Amap: {e}")
    print("http://0.0.0.0:8081")
    app.run(host='0.0.0.0', port=8081, debug=False, threaded=True)
