import os
import time
import json
import gradio as gr
from neo4j import GraphDatabase 
from langchain_community.graphs import Neo4jGraph
from langchain.chains import GraphCypherQAChain
from langchain_openai import ChatOpenAI
from langchain_core.prompts import PromptTemplate
from langchain_community.utilities import WikipediaAPIWrapper
from langchain_community.tools import WikipediaQueryRun
from pyvis.network import Network
import html
import tempfile

BIOCHAR_QUERY_TEMPLATES = """
Biochar remediation query patterns:
1. Pollutant <- treats - Biochar
   MATCH (b:Biochar)-[r:TREATS]->(p:Pollutant)
   WHERE toLower(p.name) CONTAINS toLower($keyword) OR $keyword IN coalesce(p.aliases, [])
   RETURN b.name, type(r), p.name

2. Biochar -> prepared_by -> PreparationMethod
   MATCH (b:Biochar)-[:PREPARED_BY]->(m:PreparationMethod)
   WHERE toLower(b.name) CONTAINS toLower($keyword) OR $keyword IN coalesce(b.aliases, [])
   RETURN b.name, m.name

3. Biochar -> has_property -> Property
   MATCH (b:Biochar)-[:HAS_PROPERTY]->(prop:Property)
   WHERE toLower(b.name) CONTAINS toLower($keyword) OR $keyword IN coalesce(b.aliases, [])
   RETURN b.name, prop.name

4. Biochar / Pollutant -> removes_via -> Mechanism
   MATCH (n)-[:REMOVES_VIA]->(m:Mechanism)
   WHERE (n:Biochar OR n:Pollutant)
   RETURN n.name, labels(n), m.name

5. Biochar -> derived_from -> Feedstock
   MATCH (b:Biochar)-[:DERIVED_FROM]->(f:Feedstock)
   RETURN b.name, f.name

6. Biochar / ApplicationScenario -> applied_in -> EnvironmentMedium
   MATCH (n)-[:APPLIED_IN]->(env:EnvironmentMedium)
   WHERE n:Biochar OR n:ApplicationScenario
   RETURN n.name, env.name

7. Biochar / PreparationMethod / ApplicationScenario -> performs_under -> Condition
   MATCH (n)-[:PERFORMS_UNDER]->(c:Condition)
   RETURN n.name, labels(n), c.name

8. Biochar / Property -> characterized_by -> CharacterizationMethod
   MATCH (n)-[:CHARACTERIZED_BY]->(cm:CharacterizationMethod)
   WHERE n:Biochar OR n:Property
   RETURN n.name, cm.name
"""

BIOCHAR_CYPHER_TEMPLATE = """You are an expert Neo4j Cypher generator for a biochar environmental remediation knowledge graph.

Task:
- Generate a read-only Cypher query that answers the user's question.
- Use only the relationship types and node labels present in the provided schema.
- Prefer precise, compact queries over broad graph expansion.

Domain priorities:
- Most important paths:
  Biochar-[:TREATS]->Pollutant
  Biochar-[:PREPARED_BY]->PreparationMethod
  Biochar-[:HAS_PROPERTY]->Property
  Biochar-[:REMOVES_VIA]->Mechanism
  Biochar-[:DERIVED_FROM]->Feedstock
  Biochar-[:APPLIED_IN]->EnvironmentMedium
  Biochar-[:PERFORMS_UNDER]->Condition
  Biochar|Property-[:CHARACTERIZED_BY]->CharacterizationMethod
- `Chunk` nodes are evidence containers. Use them only when the question asks for evidence, source text, or supporting context.
- Prefer matching on `name`; when helpful, also consider `aliases`.
- When the user asks "which", "what", "list", or "show", return distinct names and a small number of relevant supporting fields.
- When the user asks about mechanisms, prioritize `REMOVES_VIA`.
- When the user asks about synthesis or modification, prioritize `PREPARED_BY` and `DERIVED_FROM`.
- When the user asks about material characteristics, prioritize `HAS_PROPERTY`.
- When the user asks about environmental application, prioritize `APPLIED_IN` and `PERFORMS_UNDER`.

Query rules:
- Return at most 15 rows unless the user explicitly asks for more.
- Use `DISTINCT` whenever duplicates are likely.
- Do not write, delete, merge, call procedures, or use APOC.
- Do not invent labels or relationships.
- If multiple hop paths are needed, keep them biologically and chemically meaningful.
- If the question names a specific pollutant or biochar, filter with `toLower(n.name) CONTAINS toLower("...")`
  and optionally `ANY(alias IN coalesce(n.aliases, []) WHERE toLower(alias) CONTAINS toLower("..."))`.

Helpful query patterns:
{query_templates}

Schema:
{schema}

Question:
{question}

Return only the Cypher query text.
"""


def extract_visualization_keywords(message, api_key, base_url):
    llm_kwargs = {"model": "gpt-4o-mini", "temperature": 0, "openai_api_key": api_key}
    if base_url and base_url.strip():
        llm_kwargs["base_url"] = base_url.strip()

    llm = ChatOpenAI(**llm_kwargs)
    prompt = f"""你要为知识图谱可视化提取检索关键词。

用户问题可能是中文、英文或中英混合。请尽量返回图谱中最可能存在的实体关键词候选，优先返回英文学术名称；如果原问题里有中文术语，请同时保留中文候选。

候选可以来自这些类别：
- 污染物
- 生物炭材料
- 原料
- 制备方法
- 性质
- 机理

要求：
1. 返回 1 到 3 个候选短语。
2. 优先返回文献和知识图谱中最可能使用的标准英文名称。
3. 如果问题是中文，尽量给出对应英文术语。
4. 只返回 JSON，不要解释。

输出格式：
{{"keywords": ["candidate 1", "candidate 2"]}}

问题：{message}
"""
    content = llm.invoke(prompt).content.strip()
    try:
        payload = json.loads(content)
        keywords = payload.get("keywords") or []
        keywords = [str(k).strip() for k in keywords if str(k).strip()]
    except Exception:
        keywords = [content]

    keywords.append(message.strip())
    normalized = []
    seen = set()
    for keyword in keywords:
        if keyword and keyword not in seen:
            seen.add(keyword)
            normalized.append(keyword)
    return normalized[:4]


def rewrite_question_for_graph_search(message, api_key, base_url):
    llm_kwargs = {"model": "gpt-4o-mini", "temperature": 0, "openai_api_key": api_key}
    if base_url and base_url.strip():
        llm_kwargs["base_url"] = base_url.strip()

    llm = ChatOpenAI(**llm_kwargs)
    prompt = f"""你要把用户问题改写成更适合英文知识图谱检索的问句。

要求：
1. 保留原问题语义，不要扩展没提到的事实。
2. 如果原问题是中文，请优先改写为简洁、标准的英文科研问句。
3. 把中文术语改成更可能出现在知识图谱节点里的英文表达。
4. 保留核心实体，例如污染物、生物炭材料、原料、制备方法、性质、机理。
5. 只返回改写后的单句问句，不要解释。

示例：
- 哪些生物炭可以处理镉、铅或砷等污染物？
  -> Which biochars can treat pollutants such as cadmium, lead, or arsenic?
- 稻壳生物炭常见的制备方法和关键条件有哪些？
  -> What are the common preparation methods and key conditions for rice husk biochar?

原问题：{message}
"""
    return llm.invoke(prompt).content.strip()


def extract_wikipedia_keyword(message, api_key, base_url):
    llm_kwargs = {"model": "gpt-4o-mini", "temperature": 0, "openai_api_key": api_key}
    if base_url and base_url.strip():
        llm_kwargs["base_url"] = base_url.strip()

    llm = ChatOpenAI(**llm_kwargs)
    prompt = f"""请为 Wikipedia 检索提取一个最具体的实体关键词。

要求：
1. 如果原问题是中文，尽量输出对应的标准英文术语。
2. 只输出一个词或一个短语。
3. 不要输出宽泛类别词，如 Pollutant、Property、Material。
4. 不要解释。

问题：{message}
"""
    return llm.invoke(prompt).content.strip()

# ==========================================
# 1. Read Database Credentials
# ==========================================
NEO4J_URI = os.environ.get("NEO4J_URI", "")
NEO4J_USERNAME = os.environ.get("NEO4J_USERNAME", "")
NEO4J_PASSWORD = os.environ.get("NEO4J_PASSWORD", "")
DEFAULT_BASE_URL = ""
NEO4J_DATABASE_NAME = "f8c5f809" 

# ==========================================
# 2. Database Initialization (LangChain RAG)
# ==========================================
try:
    graph = Neo4jGraph(
        url=NEO4J_URI,
        username=NEO4J_USERNAME,
        password=NEO4J_PASSWORD,
        database=NEO4J_DATABASE_NAME 
    )
    print("✅ LangChain Graph RAG connected successfully!")
except Exception as e:
    print(f"❌ LangChain Graph RAG connection failed: {e}")
    graph = None 

# Native driver for visualization
neo4j_driver = GraphDatabase.driver(NEO4J_URI, auth=(NEO4J_USERNAME, NEO4J_PASSWORD))

# ==========================================
# 3. Dynamic Graph Chain Initialization
# ==========================================
def get_graph_chain(api_key: str, base_url: str):
    # 动态组装大模型参数，支持自定义 Base URL
    llm_kwargs = {
        "model": "gpt-4o-mini",
        "temperature": 0,
        "openai_api_key": api_key
    }
    if base_url and base_url.strip():
        llm_kwargs["base_url"] = base_url.strip()
        
    llm = ChatOpenAI(**llm_kwargs)
    cypher_prompt = PromptTemplate(
        template=BIOCHAR_CYPHER_TEMPLATE,
        input_variables=["schema", "question"],
        partial_variables={"query_templates": BIOCHAR_QUERY_TEMPLATES}
    )

    qa_template = """你是一名严谨的生物炭环境修复、环境工程与材料学专家。
    请严格基于图数据库返回的 [Context] 作答。
    
    [关键规则]:
    1. 尽最大努力使用给定 Context。即使上下文稍微凌乱，也要尽量提取其中与生物炭、污染物、原料、制备方法、性质、机理、环境介质相关的信息。
    2. 回答要围绕生物炭环境修复展开，尽量区分处理对象、制备方式、材料性质和修复机理。
    3. 除非上下文明确把某设备当作表征方法，否则不要把纯实验辅助设备写进答案。
    4. 只有当 Context 完全为空时，才允许输出完全等于 "Not found"。只要有任何数据，就必须组织出有用答案。
    5. 最终请用专业、自然、清晰的中文回答。

    Context: {context}
    Question: {question}

    中文专业回答:"""
    qa_prompt = PromptTemplate(template=qa_template, input_variables=["context", "question"])

    kg_rag_chain = GraphCypherQAChain.from_llm(
        llm=llm,
        graph=graph,
        verbose=True,
        qa_prompt=qa_prompt,
        cypher_prompt=cypher_prompt,
        top_k=15,
        allow_dangerous_requests=True
    )
    return kg_rag_chain

# ==========================================
# 4. 🕸️ Core Visualization Function: Generate Pyvis HTML
# ==========================================
def generate_vis_subgraph_html(message, api_key, base_url):
    if not neo4j_driver:
        return "⚠️ 图数据库未连接"

    if not api_key or not api_key.startswith("sk-"):
        return "<h3>⚠️ 未填写 API Key，无法生成可视化结果</h3>"

    # Extract keyword candidates using LLM
    try:
        keywords = extract_visualization_keywords(message, api_key, base_url)
    except Exception as e:
        keywords = [message]
        print(f"Keyword extraction failed: {e}")

    if not keywords:
        keywords = [message]

    print(f"🔍 Extracted graph keywords for visualization: {keywords}")

    vis_cypher = """
    MATCH (core_entity)
    WHERE any(keyword IN $keywords WHERE
        toLower(coalesce(core_entity.name, "")) CONTAINS toLower(keyword)
        OR toLower(coalesce(core_entity.id, "")) CONTAINS toLower(keyword)
        OR toLower(coalesce(core_entity.canonical_name, "")) CONTAINS toLower(keyword)
        OR any(alias IN coalesce(core_entity.aliases, []) WHERE toLower(alias) CONTAINS toLower(keyword))
    )
    MATCH (core_entity)-[r]-(neighbor)
    RETURN core_entity, r, neighbor 
    LIMIT 60
    """
    
    color_map = {
        'Biochar': '#7C5C3B',
        'Pollutant': '#E76F51',
        'Feedstock': '#8AB17D',
        'PreparationMethod': '#E9C46A',
        'Property': '#4EA8DE',
        'Mechanism': '#9D4EDD',
        'EnvironmentMedium': '#2A9D8F',
        'Condition': '#F4A261',
        'CharacterizationMethod': '#7B8CDE',
        'ApplicationScenario': '#B56576',
        'Chunk': '#E5E7E9'
    }

    nodes = {}
    edges = []

    try:
        with neo4j_driver.session(database=NEO4J_DATABASE_NAME) as session:
            result = session.run(vis_cypher, keywords=keywords)
            
            for record in result:
                core_node = record['core_entity']
                rel = record['r']
                neighbor_node = record['neighbor']
                
                core_id = core_node.element_id
                if core_id not in nodes:
                    core_label = list(core_node.labels)[0] 
                    core_name = core_node.get('name', core_node.get('id', 'Unknown'))
                    nodes[core_id] = {
                        'id': core_id,
                        'label': core_name,
                        'title': f"Label: {core_label}\n{json.dumps(dict(core_node), indent=2, ensure_ascii=False)}",
                        'color': color_map.get(core_label, '#D2E5FF') 
                    }

                neighbor_id = neighbor_node.element_id
                if neighbor_id not in nodes:
                    neighbor_label = list(neighbor_node.labels)[0]
                    if neighbor_label == 'Chunk':
                         raw_name = neighbor_node.get('text', '')[:12] + '...' 
                    else:
                        raw_name = neighbor_node.get('name', neighbor_node.get('id', 'Unknown'))
                    
                    nodes[neighbor_id] = {
                        'id': neighbor_id,
                        'label': raw_name,
                        'title': f"Label: {neighbor_label}\n{json.dumps(dict(neighbor_node), indent=2, ensure_ascii=False)}",
                        'color': color_map.get(neighbor_label, '#D2E5FF')
                    }

                rel_type = rel.type
                edges.append((rel.start_node.element_id, rel.end_node.element_id, rel_type))

    except Exception as e:
        print(f"❌ 可视化查询失败: {str(e)}")
        return f"<h3>❌ 可视化关系查询失败：{str(e)}</h3>"

    if not nodes:
        return f"<h3>⚠️ 检索完成，但没有找到可视化关系网络。</h3><p>💡 当前尝试匹配的关键词为：<b>{', '.join(keywords)}</b>。这可能意味着这些关键词未命中图谱中的节点名称或别名，或相关节点尚未建立一阶关系。</p>"

    net = Network(height='600px', width='100%', bgcolor='#ffffff', font_color='#333333', notebook=False)
    for node_id, node_data in nodes.items():
        net.add_node(node_data['id'], label=node_data['label'], title=node_data['title'], color=node_data['color'])
    for edge in edges:
        net.add_edge(edge[0], edge[1], label=edge[2], width=1, color='#DDDDDD')
    
    net.toggle_physics(True)
    net.set_options("""
    var options = {
      "interaction": {
        "hover": true,
        "hoverConnectedEdges": true
      }
    }
    """)
    
    path = tempfile.mktemp(suffix='.html')
    net.save_graph(path)
    with open(path, 'r', encoding='utf-8') as f:
        html_content = f.read()
        
    escaped_html = html.escape(html_content)
    return f'<iframe style="width: 100%; height: 600px; border: none;" srcdoc="{escaped_html}"></iframe>'

# ==========================================
# 5. Q&A Function (Graph Priority + Wiki Fallback)
# ==========================================
def answer_question(message, history, api_key, base_url):
    history.append({"role": "user", "content": message})
    
    if not api_key or not api_key.startswith("sk-"):
        history.append({"role": "assistant", "content": "⚠️ 请先在上方填写有效的 LLM API Key。"})
        yield history
        return
    
    if graph is None:
        history.append({"role": "assistant", "content": "⚠️ 图数据库连接失败，请检查环境变量配置。"})
        yield history
        return
    
    try:
        chain = get_graph_chain(api_key, base_url) 
        graph_query = rewrite_question_for_graph_search(message, api_key, base_url)
        
        history.append({"role": "assistant", "content": "🧠 正在优先检索生物炭修复知识图谱，请稍候..."})
        yield history
        
        response = chain.invoke({"query": graph_query})
        final_answer = response["result"]
        
        if "Not found" in final_answer:
            history[-1]["content"] = "🧠 图谱中暂未命中，正在补充检索 Wikipedia..."
            yield history
            keyword = extract_wikipedia_keyword(message, api_key, base_url)
            
            wiki_wrapper = WikipediaAPIWrapper(lang="en", top_k_results=1, doc_content_chars_max=800)
            wiki_tool = WikipediaQueryRun(api_wrapper=wiki_wrapper)
            wiki_result = wiki_tool.run(keyword) 
            
            if "No good Wikipedia Search Result" in wiki_result or not wiki_result.strip():
                 final_answer = f"图谱中没有找到相关信息，Wikipedia 中也没有检索到 [{keyword}] 的合适结果。"
            else:
                 final_answer = f"**图谱中未找到该实体。以下是 Wikipedia 中关于 [{keyword}] 的基础说明：**\n\n{wiki_result}"
        
        streamed_text = ""
        for char in final_answer:
            streamed_text += char
            history[-1]["content"] = streamed_text 
            yield history
            time.sleep(0.005) 
            
    except Exception as e:
        history[-1]["content"] = f"处理过程中出现系统错误：{str(e)}"
        yield history


def get_schema():
    if graph is None:
        return "⚠️ 图数据库未连接"
    live_schema = str(graph.schema or "").strip()
    if not live_schema:
        return "⚠️ 当前未读取到实时 schema。可能原因包括：图数据库中尚无数据、连接已建立但 schema introspection 暂未返回结果，或当前库尚未完成导入。"
    return live_schema

# ==========================================
# 6. UI Construction (Professional Sans-serif Theme)
# ==========================================
custom_theme = gr.themes.Soft(
    font=[gr.themes.GoogleFont("Inter"), "ui-sans-serif", "system-ui", "sans-serif"],
    text_size=gr.themes.sizes.text_md
)

custom_css = """
:root {
  --paper: #f6f2e8;
  --ink: #1e2a22;
  --muted: #5d665e;
  --sand: #e9dfc8;
  --clay: #7c5c3b;
  --leaf: #365f48;
  --water: #d9e9e8;
  --panel: rgba(255, 252, 246, 0.92);
}

.gradio-container {
  background:
    radial-gradient(circle at top right, rgba(54, 95, 72, 0.12), transparent 26%),
    radial-gradient(circle at left center, rgba(124, 92, 59, 0.10), transparent 24%),
    linear-gradient(180deg, #f7f4ee 0%, #efe5d4 100%);
}

.app-shell {
  max-width: 1440px;
  margin: 0 auto;
}

.hero-banner {
  background: linear-gradient(135deg, rgba(250, 248, 243, 0.96), rgba(232, 223, 203, 0.92));
  border: 1px solid rgba(124, 92, 59, 0.18);
  border-radius: 28px;
  padding: 28px 30px;
  box-shadow: 0 18px 40px rgba(67, 49, 28, 0.10);
}

.hero-kicker {
  display: inline-block;
  padding: 6px 12px;
  border-radius: 999px;
  background: rgba(54, 95, 72, 0.10);
  color: var(--leaf);
  font-size: 12px;
  letter-spacing: 0.08em;
  text-transform: uppercase;
  font-weight: 700;
}

.hero-title {
  margin: 14px 0 10px;
  font-size: 42px;
  line-height: 1.08;
  color: var(--ink);
  font-weight: 800;
}

.hero-subtitle {
  margin: 0;
  max-width: 860px;
  color: var(--muted);
  font-size: 16px;
  line-height: 1.75;
}

.metric-strip {
  display: grid;
  grid-template-columns: repeat(3, minmax(0, 1fr));
  gap: 12px;
  margin-top: 18px;
}

.metric-card {
  border-radius: 18px;
  padding: 16px 18px;
  background: rgba(255,255,255,0.55);
  border: 1px solid rgba(54, 95, 72, 0.12);
}

.metric-label {
  font-size: 12px;
  color: #6a7067;
  text-transform: uppercase;
  letter-spacing: 0.08em;
}

.metric-value {
  margin-top: 8px;
  font-size: 20px;
  font-weight: 700;
  color: #223228;
}

.workspace-row {
  gap: 20px;
}

.side-panel, .main-panel {
  background: var(--panel);
  border: 1px solid rgba(54, 95, 72, 0.14);
  border-radius: 26px;
  box-shadow: 0 16px 36px rgba(44, 34, 20, 0.08);
}

.side-panel {
  padding: 20px;
}

.main-panel {
  padding: 22px;
}

.panel-title {
  font-size: 14px;
  color: #617064;
  letter-spacing: 0.08em;
  text-transform: uppercase;
  margin-bottom: 8px;
}

.panel-headline {
  font-size: 26px;
  color: var(--ink);
  font-weight: 800;
  margin-bottom: 8px;
}

.panel-copy {
  color: var(--muted);
  line-height: 1.7;
  font-size: 14px;
  margin-bottom: 16px;
}

.gradio-container .gr-button-primary {
  background: linear-gradient(135deg, #355f49, #4e7b60);
  border: none !important;
}

.gradio-container .gr-button-secondary {
  border-color: rgba(54, 95, 72, 0.18) !important;
}

.viz-intro {
  padding: 14px 16px;
  border-radius: 16px;
  background: rgba(233, 223, 200, 0.55);
  color: #58452f;
  font-size: 13px;
  line-height: 1.7;
  margin-bottom: 10px;
}

.path-card {
  margin-top: 18px;
  padding: 18px;
  border-radius: 22px;
  background:
    radial-gradient(circle at top right, rgba(54, 95, 72, 0.10), transparent 38%),
    linear-gradient(160deg, rgba(241, 235, 221, 0.95), rgba(232, 220, 196, 0.92));
  border: 1px solid rgba(124, 92, 59, 0.14);
}

.path-title {
  font-size: 12px;
  color: #72685a;
  text-transform: uppercase;
  letter-spacing: 0.08em;
  margin-bottom: 8px;
}

.path-headline {
  font-size: 22px;
  font-weight: 800;
  color: #26342a;
  margin-bottom: 10px;
}

.path-copy {
  font-size: 14px;
  line-height: 1.75;
  color: #57584f;
  margin-bottom: 14px;
}

.path-steps {
  display: grid;
  gap: 10px;
}

.path-step {
  padding: 12px 14px;
  border-radius: 16px;
  background: rgba(255,255,255,0.52);
  border: 1px solid rgba(54, 95, 72, 0.10);
}

.path-step strong {
  display: block;
  color: #304036;
  margin-bottom: 4px;
  font-size: 14px;
}

.path-step span {
  color: #61655b;
  font-size: 13px;
  line-height: 1.65;
}

@media (max-width: 900px) {
  .hero-title {
    font-size: 34px;
  }

  .metric-strip {
    grid-template-columns: 1fr;
  }
}
"""

with gr.Blocks(theme=custom_theme, css=custom_css) as demo:
    with gr.Column(elem_classes=["app-shell"]):
        gr.HTML("""
        <div class="hero-banner">
          <div class="hero-kicker">Designed by Zhou's Group</div>
          <div class="hero-title">生物炭环境修复动态资源库</div>
          <p class="hero-subtitle">
            支持用户围绕污染物、生物炭、制备方法、基础性质与修复机理做连续探索。
          </p>
        </div>
        """)

        with gr.Row(elem_classes=["workspace-row"]):
            with gr.Column(scale=4, elem_classes=["side-panel"]):
                gr.HTML("""
                <div class="panel-title">控制台</div>
                <div class="panel-headline">连接、提示与图谱总览</div>
                <div class="panel-copy">
                  请先填写凭证，再从示例问题或自定义问题开始。
                </div>
                """)

                api_key_input = gr.Textbox(
                    label="🔑 API Key",
                    placeholder="sk-...",
                    type="password"
                )
                base_url_input = gr.Textbox(
                    label="🌐 Base URL",
                    placeholder="https://api.openai.com/v1",
                    value=DEFAULT_BASE_URL
                )

                gr.Markdown("""
                **提问模板**

                1. 哪些生物炭可以处理镉、铅或砷？
                2. 稻壳生物炭常见的制备方法和条件有哪些？
                3. 某类改性生物炭通常有哪些基础性质？
                4. 生物炭去除磷酸盐可能涉及哪些机理？
                5. 哪些表征方法常用于分析生物炭结构？
                """)

                gr.HTML("""
                <div class="path-card">
                  <div class="path-title">研究路径</div>
                  <div class="path-headline">建议的探索顺序</div>
                  <div class="path-copy">
                    如果你想连续追问并快速得到更稳定的图谱结果，可以按下面这条路径逐步展开。
                  </div>
                  <div class="path-steps">
                    <div class="path-step">
                      <strong>先问污染物</strong>
                      <span>先锁定镉、砷、磷酸盐、抗生素等目标污染物，找到相关生物炭材料。</span>
                    </div>
                    <div class="path-step">
                      <strong>再问材料与制备</strong>
                      <span>继续追问原料、热解方式、改性手段和关键条件，补全材料侧信息。</span>
                    </div>
                    <div class="path-step">
                      <strong>最后问性质与机理</strong>
                      <span>查看表面性质、作用机理和子图关系，确认知识连接是否完整。</span>
                    </div>
                  </div>
                </div>
                """)

            with gr.Column(scale=8, elem_classes=["main-panel"]):
                gr.HTML("""
                <div class="panel-title">研究工作区</div>
                <div class="panel-headline">问答与关系导航</div>
                <div class="panel-copy">
                  专注于连续提问、结果阅读和关系网络浏览。每次提交问题后，系统会先回答，再自动生成相关实体的一阶子图。
                </div>
                """)

                chatbot_ui = gr.Chatbot(
                    avatar_images=("user.png", "ai.png"),
                    height=430,
                    render=True
                )

                msg_input = gr.Textbox(
                    label="输入问题",
                    placeholder="例如：哪些生物炭可以处理镉、铅或砷等污染物？",
                    lines=2,
                    max_lines=4
                )

                with gr.Row():
                    submit_btn = gr.Button("开始分析", variant="primary")
                    clear_btn = gr.Button("清空当前会话", variant="secondary")

                with gr.Tabs():
                    with gr.Tab("关系子图"):
                        gr.HTML("""
                        <div class="viz-intro">
                          当前问题最相关的关系网络。支持拖拽、缩放和查看节点属性，适合快速确认某一污染物、生物炭或机理在图谱中的连接方式。
                        </div>
                        """)
                        html_vis_output = gr.HTML(
                            value="<h3>💬 请输入凭证并提交问题...</h3><p>系统会在这里绘制与问题关联最紧密的动态子图。</p>",
                            height="640px",
                            render=True
                        )

                    with gr.Tab("提问建议"):
                        gr.Markdown(
                            "\n".join(
                                [
                                    "1. 问污染物：哪些生物炭可以处理镉、砷、磷酸盐？",
                                    "2. 问材料：某种改性生物炭的制备方法、原料和性质是什么？",
                                    "3. 问机理：某类污染物去除可能涉及哪些机制？",
                                    "4. 问条件：某种处理效果在什么 pH、温度或投加量下表现更好？",
                                    "5. 问表征：哪些方法常用于分析生物炭结构与表面性质？",
                                ]
                            )
                        )

    # ==========================================
    # Wiring logic (现在包含了 base_url_input)
    # ==========================================
    answer_event = msg_input.submit(
        fn=answer_question,
        inputs=[msg_input, chatbot_ui, api_key_input, base_url_input],
        outputs=[chatbot_ui]
    )
    answer_event.then(
        fn=generate_vis_subgraph_html,
        inputs=[msg_input, api_key_input, base_url_input],
        outputs=[html_vis_output]
    )
    answer_event.then(lambda: "", outputs=[msg_input])

    btn_event = submit_btn.click(
        fn=answer_question,
        inputs=[msg_input, chatbot_ui, api_key_input, base_url_input],
        outputs=[chatbot_ui]
    )
    btn_event.then(
        fn=generate_vis_subgraph_html,
        inputs=[msg_input, api_key_input, base_url_input],
        outputs=[html_vis_output]
    )
    btn_event.then(lambda: "", outputs=[msg_input])

    clear_btn.click(lambda: [], outputs=chatbot_ui, queue=False)
    clear_btn.click(lambda: "<h3>💬 请输入一个生物炭修复相关问题...</h3>", outputs=html_vis_output)

if __name__ == "__main__":
    demo.launch(theme=custom_theme)