import streamlit as st
import pandas as pd
import folium
from streamlit_folium import st_folium
import os
import requests
import io
from openai import OpenAI
import json


# --- 0. 配置与初始化 ---
st.set_page_config(
    page_title="李白生平GIS与RAG整合",
    page_icon="🐉",
    layout="wide"
)

# 初始化OpenAI客户端
client = OpenAI(
    api_key=os.getenv("DEEPSEEK_API_KEY", "sk-72997944466a4af2bcd52a068895f8cf"),
    base_url="https://api.deepseek.com"
)

# ----------------------------------------------------
# 全局变量定义
# ----------------------------------------------------
GITHUB_XLSX_URL = "https://raw.githubusercontent.com/seblee424/5571_data/main/李白人生重要节点与代表作地理位置.xlsx"
location_col = '地点（古称/今称）'
summary_col = '诗作/事件摘要'


# --- 1. RAG 补充函数：抓取CBDB李白人物资料 ---
@st.cache_data(ttl=3600)
def get_cbdb_data(name="李白"):
    url = f"https://cbdb.fas.harvard.edu/cbdbapi/person.php?name={name}&o=json"
    headers = {"User-Agent": "Mozilla/5.0 (Streamlit App)"}
    try:
        response = requests.get(url, headers=headers, timeout=10)
        if response.status_code == 200:
            return response.json()
        else:
            return None
    except Exception:
        return None


# --- 2. 地理坐标匹配表 ---
LOCATION_COORDS = {
    "碎叶城": {"lat": 42.8447, "lon": 75.1648, "match_keys": ["碎叶城"]},
    "峨眉山": {"lat": 29.5807, "lon": 103.3592, "match_keys": ["峨眉山"]},
    "蜀中": {"lat": 31.7828, "lon": 104.7570, "match_keys": ["蜀中", "江油"]},
    "荆门/南津关": {"lat": 30.5667, "lon": 111.4500, "match_keys": ["荆门", "南津关"]},
    "岳阳楼": {"lat": 29.3879, "lon": 113.1092, "match_keys": ["岳阳楼", "岳阳"]},
    "安陆": {"lat": 31.3653, "lon": 113.7077, "match_keys": ["安陆"]},
    "黄鹤楼": {"lat": 30.5484, "lon": 114.3168, "match_keys": ["黄鹤楼", "武汉"]},
    "金陵（凤凰台）": {"lat": 32.0415, "lon": 118.7781, "match_keys": ["金陵", "凤凰台", "南京"]},
    "庐山": {"lat": 29.5910, "lon": 115.9922, "match_keys": ["庐山", "九江"]},
    "天姥山": {"lat": 29.5000, "lon": 120.8900, "match_keys": ["天姥山"]},
    "金陵/长干里": {"lat": 32.0298, "lon": 118.7900, "match_keys": ["长干里"]},
    "长安": {"lat": 34.2652, "lon": 108.9500, "match_keys": ["长安", "西安"]},
    "长安/宫廷": {"lat": 34.2652, "lon": 108.9500, "match_keys": ["宫廷"]},
    "长安/洛阳": {"lat": 34.6859, "lon": 112.4600, "match_keys": ["洛阳"]},
    "桃花潭": {"lat": 30.4079, "lon": 118.4230, "match_keys": ["桃花潭", "泾县"]},
    "敬亭山": {"lat": 30.9822, "lon": 118.7844, "match_keys": ["敬亭山", "宣城"]},
    "天门山": {"lat": 31.4285, "lon": 118.3970, "match_keys": ["天门山", "芜湖"]},
    "扬州/旅店": {"lat": 32.3934, "lon": 119.4290, "match_keys": ["扬州"]},
    "夜郎": {"lat": 27.6888, "lon": 106.3773, "match_keys": ["夜郎", "桐梓"]},
    "白帝城": {"lat": 31.0450, "lon": 109.5780, "match_keys": ["白帝城", "奉节"]},
    "秋浦": {"lat": 30.6500, "lon": 117.4800, "match_keys": ["秋浦", "池州"]},
    "当涂": {"lat": 31.5453, "lon": 118.4870, "match_keys": ["当涂", "马鞍山"]},

    # 泛指类
    "蜀道": {"lat": 31.0000, "lon": 107.0000, "match_keys": ["蜀道"]},
    "月下独酌": {"lat": 34.2652, "lon": 108.9500, "match_keys": ["独酌", "月下"]},
    "静夜思": {"lat": 32.3934, "lon": 119.4290, "match_keys": ["静夜思"]},
    "长江沿线": {"lat": 30.5928, "lon": 114.3055, "match_keys": ["长江"]},
    "战城南": {"lat": 35.0000, "lon": 100.0000, "match_keys": ["边塞", "战争"]},
    "送友人": {"lat": 30.5928, "lon": 114.3055, "match_keys": ["送友人"]},
    "将进酒": {"lat": 34.2652, "lon": 108.9500, "match_keys": ["将进酒", "豪饮"]},
    "行路难": {"lat": 34.2652, "lon": 108.9500, "match_keys": ["行路难"]},
}


# --- 3. GitHub 在线加载 XLSX ---
@st.cache_data
def load_and_prepare_data_github(url):

    try:
        st.info("正在从 GitHub 加载 XLSX 数据...")
        res = requests.get(url, timeout=10)

        if res.status_code != 200:
            st.error(f"GitHub 下载失败，状态码: {res.status_code}")
            return pd.DataFrame()

        df = pd.read_excel(io.BytesIO(res.content), sheet_name=0)
        st.success("✅ 成功从 GitHub 加载数据！")

    except Exception as e:
        st.error(f"❌ 加载 GitHub XLSX 失败: {e}")
        return pd.DataFrame()

    df.columns = df.columns.str.strip()

    if location_col not in df.columns or summary_col not in df.columns:
        st.error(f"XLSX 缺少关键列: {location_col}, {summary_col}")
        return pd.DataFrame()

    # 经纬度匹配
    coords_list = []
    df["coords_key"] = ""

    for idx, row in df.iterrows():
        loc = str(row[location_col]).strip()
        match = None
        match_key = "未知"

        for key, data in LOCATION_COORDS.items():
            if loc == key:
                match = data
                match_key = key
                break
            if any(k in loc for k in data["match_keys"]):
                match = data
                match_key = key
                break

        if match:
            coords_list.append((match["lat"], match["lon"]))
            df.loc[idx, "coords_key"] = match_key
        else:
            coords_list.append((34.0478, 108.4357))

    df["Latitude"] = [x[0] for x in coords_list]
    df["Longitude"] = [x[1] for x in coords_list]

    return df


# --- 使用 GitHub 数据加载 ---
data_df = load_and_prepare_data_github(GITHUB_XLSX_URL)


# --- 4. Chatbot (RAG) ---
def run_chatbot(cbdb_data, prompt):
    cbdb_text = json.dumps(cbdb_data, ensure_ascii=False)[:5000] if cbdb_data else "无CBDB资料。"

    system_prompt_rag = (
        "你是一个李白生平研究的聊天机器人，擅长介绍李白的生平、作品和相关地点。"
        "回答中如提及地点，请使用与地图匹配的地名，如安陆、桃花潭、黄鹤楼、长安。"
        "引用 CBDB 信息时请在结尾标注“资料源自CBDB”。"
        f"\n\nCBDB数据如下：{cbdb_text}"
    )

    try:
        msgs = [{"role": "system", "content": system_prompt_rag}]
        msgs.extend(st.session_state.chat_history[-5:])

        resp = client.chat.completions.create(
            model="deepseek-chat",
            messages=msgs,
            stream=False
        )

        answer = resp.choices[0].message.content.strip()

        # 自动匹配回答的地名，对地图高亮
        highlight = None
        for key in data_df["coords_key"].unique():
            if key != "未知" and key in answer:
                highlight = key
                break

        st.session_state.highlight_location_key = highlight
        return answer

    except Exception as e:
        return f"Chatbot 发生错误: {e}"


# --- 5. GIS 地图 ---
def create_li_bai_map(df, highlight_key):

    if df.empty:
        return folium.Map(location=[34.0, 108.0], zoom_start=4)

    center_lat = df["Latitude"].mean()
    center_lon = df["Longitude"].mean()

    m = folium.Map(
        location=[center_lat, center_lon],
        zoom_start=4.5,
        tiles="cartodbdarkmatter"
    )

    # 轨迹线
    points = df[["Latitude", "Longitude"]].values.tolist()
    if len(points) > 1:
        folium.PolyLine(points, color="#00AEEF", weight=3, opacity=0.5).add_to(m)

    # 节点
    for _, row in df.iterrows():
        is_high = (row["coords_key"] == highlight_key)

        popup_html = f"""
        <b>序号:</b> {row['序号']}<br>
        <b>阶段:</b> {row['阶段（大致年份）']}<br>
        <b>地点:</b> {row['地点（古称/今称）']}<br>
        <b>事件/诗作:</b> {row['诗作/事件摘要']}<br>
        <b>核心情感:</b> {row['核心情感/主题']}<br>
        <b>节点类型:</b> {row['节点类型']}
        """

        if is_high:
            color, icon = "orange", "fire"
        elif "人生事件" in row["节点类型"]:
            color, icon = "blue", "user"
        else:
            color, icon = "green", "flag"

        folium.Marker(
            [row["Latitude"], row["Longitude"]],
            popup=folium.Popup(popup_html, max_width=300),
            icon=folium.Icon(color=color, icon=icon, prefix="fa", icon_color="white"),
        ).add_to(m)

    return m


# --- 6. 主界面 ---
cbdb_data = get_cbdb_data("李白")

st.header("🐉 李白生平 GIS + RAG 交互系统")

if "chat_history" not in st.session_state:
    st.session_state.chat_history = []

if "highlight_location_key" not in st.session_state:
    st.session_state.highlight_location_key = None


if data_df.empty:
    st.error("❌ 数据加载失败，请检查 GitHub 文件")
else:
    col1, col2 = st.columns([1, 1.5])

    # 左侧 Chatbot
    with col1:
        st.subheader("💬 李白 Chatbot")

        if cbdb_data:
            st.info("✅ 已加载 CBDB 资料")
        else:
            st.warning("⚠️ CBDB 资料加载失败")

        for msg in st.session_state.chat_history:
            with st.chat_message(msg["role"]):
                st.markdown(msg["content"])

        if prompt := st.chat_input("请输入你的问题…"):
            with st.chat_message("user"):
                st.markdown(prompt)

            st.session_state.chat_history.append({"role": "user", "content": prompt})

            with st.chat_message("assistant"):
                with st.spinner("AI 正在思考…"):
                    answer = run_chatbot(cbdb_data, prompt)
                    st.write(answer)

                if st.session_state.highlight_location_key:
                    st.success(f"地图已高亮：{st.session_state.highlight_location_key}")

            st.session_state.chat_history.append({"role": "assistant", "content": answer})
            st.rerun()

    # 右侧 GIS
    with col2:
        st.subheader("🗺️ 李白一生足迹地图")
        st.info("提示：左侧 Chatbot 会触发地图节点高亮")

        current_map = create_li_bai_map(data_df, st.session_state.highlight_location_key)
        st_folium(current_map, width=800, height=700)