app.py

import re
import json
import time
import html as _html
import streamlit as st
import streamlit.components.v1 as components
from pathlib import Path
from typing import TypedDict, List
from dotenv import load_dotenv
from datetime import datetime
from pydantic import BaseModel, Field

from langchain_deepseek import ChatDeepSeek
from langchain_core.messages import HumanMessage, SystemMessage
from langgraph.graph import StateGraph, END
from langgraph.checkpoint.memory import InMemorySaver

from lang_config import LANG_MAP, CELL_STYLES, CELL_STYLES_DEFAULT, _EXT_TO_STRIP

load_dotenv(dotenv_path=Path(__file__).parent / ".env")

HISTORY_DIR     = Path("code_history")
HISTORY_DIR.mkdir(exist_ok=True)
MAX_CODE_CHARS  = 8000
STREAM_THROTTLE = 0.05

NODE_LABELS = {
    "file_loader":  "📂 文件导入与识别",
    "preprocessor": "🔍 预处理与结构划分",
    "analyzer":     "🧩 单元格拆分分析",
    "summarizer":   "📋 整体总结生成",
}

MODEL_OPTIONS = {
    "DeepSeek Chat (通用)":  "deepseek-chat",
    "DeepSeek Coder (代码)": "deepseek-coder",
}

_HLJS_CSS  = "https://cdnjs.cloudflare.com/ajax/libs/highlight.js/11.9.0/styles/atom-one-dark.min.css"
_HLJS_JS   = "https://cdnjs.cloudflare.com/ajax/libs/highlight.js/11.9.0/highlight.min.js"
_HLJS_LNJS = "https://cdnjs.cloudflare.com/ajax/libs/highlightjs-line-numbers.js/2.8.0/highlightjs-line-numbers.min.js"

DIFF_LABELS = {1: "简单", 2: "中等", 3: "复杂"}


# ════════════════════════════════════════════════════════════════════════════
# 工具函数
# ════════════════════════════════════════════════════════════════════════════
def load_css() -> str:
    p = Path(__file__).parent / "styles.css"
    return f"<style>{p.read_text(encoding='utf-8')}</style>" if p.exists() else ""

def load_css_content() -> str:
    p = Path(__file__).parent / "styles.css"
    return p.read_text(encoding="utf-8") if p.exists() else ""

def normalize_code(code: str) -> str:
    text = str(code)
    for _ in range(4):
        prev = text
        text = _html.unescape(text)
        if text == prev:
            break
    return text

def difficulty_dots(level: int) -> str:
    level = max(1, min(3, int(level)))
    color_cls = {1: "nb-dot-on-1", 2: "nb-dot-on-2", 3: "nb-dot-on-3"}
    dots = "".join(
        f'<span class="nb-dot {color_cls[level] if i <= level else "nb-dot-off"}"></span>'
        for i in range(1, 4)
    )
    return f'<span class="nb-difficulty" title="难度：{DIFF_LABELS[level]}">{dots}</span>'

def skeleton_html(n: int = 3) -> str:
    cards = ""
    for _ in range(n):
        cards += """
        <div class="skeleton-card">
            <div class="skeleton-block" style="width:35%;height:14px;margin-bottom:14px;"></div>
            <div class="skeleton-block" style="width:100%;height:12px;"></div>
            <div class="skeleton-block" style="width:88%;height:12px;"></div>
            <div class="skeleton-block" style="width:75%;height:12px;"></div>
            <div class="skeleton-block" style="width:92%;height:12px;margin-bottom:0;"></div>
        </div>"""
    return cards


# ════════════════════════════════════════════════════════════════════════════
# Pydantic 模型
# ════════════════════════════════════════════════════════════════════════════
class CodeCell(BaseModel):
    type: str = Field(
        description="单元格类型：import / class / function / constant / config / ui / logic / decorator / error"
    )
    title: str = Field(description="简明标题，10字以内")
    code:  str = Field(description="原始代码，严格保持缩进，禁止对 < > \" & 做 HTML 转义")
    explanation: str = Field(
        description="40-150字，先一句话概括整体作用，再逐属性说明，格式：名称：作用，用分号分隔"
    )
    difficulty: int = Field(
        default=1,
        description="难度：1=简单（基础语法/配置），2=中等（业务逻辑），3=复杂（算法/设计模式/高级特性）"
    )

class AnalysisResult(BaseModel):
    cells: List[CodeCell] = Field(description="按逻辑顺序排列的代码讲解单元列表")


# ════════════════════════════════════════════════════════════════════════════
# 注释剥离
# ════════════════════════════════════════════════════════════════════════════
def strip_comments(content: str, ext: str) -> str:
    strip_type = _EXT_TO_STRIP.get(ext.lower(), "none")
    if strip_type == "hash":
        content = re.sub(r"(?m)^[ \t]*#.*$", "", content)
        content = re.sub(r"(?m)(?<=\S)[ \t]+#[^\"'\n]*$", "", content)
    elif strip_type in ("c_style", "c_style_vue"):
        content = re.sub(r"/\*.*?\*/", "", content, flags=re.DOTALL)
        content = re.sub(r"//[^\n]*", "", content)
        if strip_type == "c_style_vue":
            content = re.sub(r"<!--.*?-->", "", content, flags=re.DOTALL)
    elif strip_type == "html":
        content = re.sub(r"<!--.*?-->", "", content, flags=re.DOTALL)
    elif strip_type == "double_dash":
        content = re.sub(r"--[^\n]*", "", content)
    elif strip_type == "lua":
        content = re.sub(r"--\[\[.*?\]\]", "", content, flags=re.DOTALL)
        content = re.sub(r"--[^\n]*", "", content)
    elif strip_type == "percent":
        content = re.sub(r"%[^\n]*", "", content)
    elif strip_type == "semicolon":
        content = re.sub(r";[^\n]*", "", content)
    elif strip_type == "vb":
        content = re.sub(r"(?m)^[ \t]*'.*$", "", content)
        content = re.sub(r"(?m)(?<=\S)[ \t]+'[^\n]*$", "", content)
    elif strip_type == "ml_block":
        content = re.sub(r"\(\*.*?\*\)", "", content, flags=re.DOTALL)
    elif strip_type == "batch":
        content = re.sub(r"(?im)^[ \t]*(?:rem\b|::)[^\n]*$", "", content)
    return re.sub(r"\n{3,}", "\n\n", content).strip()


# ════════════════════════════════════════════════════════════════════════════
# 文件加载
# ════════════════════════════════════════════════════════════════════════════
def load_code_file(uploaded_file) -> dict:
    filename  = uploaded_file.name
    ext       = Path(filename).suffix.lower()
    lang_info = LANG_MAP.get(ext, ("未知语言", "# 注释", "text"))
    language, comment_style, highlight = lang_info
    try:
        raw = uploaded_file.read()
        try:    content = raw.decode("utf-8")
        except: content = raw.decode("gbk", errors="replace")
    except Exception as e:
        return {"error": str(e)}
    lines    = content.splitlines()
    stripped = strip_comments(content, ext)
    return {
        "filename": filename, "extension": ext,
        "language": language, "comment_style": comment_style,
        "highlight": highlight, "content": content,
        "content_to_analyze": stripped,
        "line_count": len(lines), "char_count": len(content),
    }


# ════════════════════════════════════════════════════════════════════════════
# 预处理
# ════════════════════════════════════════════════════════════════════════════
def preprocess_and_divide(file_info: dict) -> dict:
    content  = file_info["content_to_analyze"]
    keywords = []
    patterns = [
        (r"^(?:import|from)\s+([\w\.]+)",                              "导入"),
        (r"^([A-Z][A-Z0-9_]{2,})\s*=",                                "常量"),
        (r"^def\s+(\w+)",                                              "函数"),
        (r"^class\s+(\w+)",                                            "类"),
        (r"(?:function\s+(\w+)|const\s+(\w+)\s*=\s*(?:async\s*)?\()", "函数"),
    ]
    seen = set()
    for pattern, ktype in patterns:
        for m in re.finditer(pattern, content, re.MULTILINE):
            name = next((g for g in m.groups() if g), None)
            if name and (ktype, name) not in seen:
                seen.add((ktype, name))
                keywords.append(f"[{ktype}] {name}")
    return {"keywords_summary": keywords[:30]}


# ════════════════════════════════════════════════════════════════════════════
# 智能分块
# ════════════════════════════════════════════════════════════════════════════
def smart_split_code(content: str, max_chars: int = MAX_CODE_CHARS) -> List[str]:
    if len(content) <= max_chars:
        return [content]
    boundary_re = re.compile(
        r"^(?:async\s+)?(?:def |class |function |const \w+ = (?:async )?\()",
        re.MULTILINE,
    )
    boundaries = [0] + [m.start() for m in boundary_re.finditer(content)]
    buf = ""; chunks: List[str] = []
    for i, start in enumerate(boundaries):
        end     = boundaries[i + 1] if i + 1 < len(boundaries) else len(content)
        segment = content[start:end]
        if len(buf) + len(segment) > max_chars and buf:
            chunks.append(buf); buf = segment
        else:
            buf += segment
    if buf: chunks.append(buf)
    result: List[str] = []
    for chunk in chunks:
        if len(chunk) > max_chars:
            result.extend(chunk[i:i+max_chars] for i in range(0, len(chunk), max_chars))
        else:
            result.append(chunk)
    return result


# ════════════════════════════════════════════════════════════════════════════
# 合并孤立配对标签
# ════════════════════════════════════════════════════════════════════════════
_OPEN_TAG_ONLY  = re.compile(r"^\s*<[a-zA-Z][^/>\s]*[^>]*>\s*$")
_CLOSE_TAG_ONLY = re.compile(r"^\s*</[a-zA-Z][^>]*>\s*$")

def merge_paired_cells(cells: list) -> list:
    if not cells: return cells
    result: list = []; i = 0
    while i < len(cells):
        cell = cells[i]; code = cell.get("code", "").strip()
        if _OPEN_TAG_ONLY.match(code) and i + 1 < len(cells):
            nxt = cells[i + 1]
            result.append({
                **nxt,
                "code": cell["code"].rstrip() + "\n" + nxt.get("code", ""),
            }); i += 2
        elif _CLOSE_TAG_ONLY.match(code) and result:
            result[-1]["code"] = result[-1]["code"].rstrip() + "\n" + cell["code"]; i += 1
        else:
            result.append(cell); i += 1
    return result


# ════════════════════════════════════════════════════════════════════════════
# LangGraph 状态
# ════════════════════════════════════════════════════════════════════════════
class AnalysisState(TypedDict):
    file_info: dict; preprocessed: dict; cells: List[dict]; summary: str

def llm_invoke_with_retry(llm, messages, max_retries: int = 2):
    last_exc = None
    for attempt in range(max_retries + 1):
        try:    return llm.invoke(messages)
        except Exception as e:
            last_exc = e
            if attempt < max_retries: time.sleep(2 ** attempt)
    raise last_exc


# ════════════════════════════════════════════════════════════════════════════
# LangGraph 分析图
# ════════════════════════════════════════════════════════════════════════════
def build_analysis_graph(llm):
    structured_llm = llm.with_structured_output(AnalysisResult)

    def analyzer(state: AnalysisState):
        fi   = state["file_info"]
        pre  = state["preprocessed"]
        lang = fi["language"]
        kw   = "\n".join(pre["keywords_summary"][:20]) or "（无）"

        system_msg = (
            f"Role: {lang}资深开发导师。\n"
            "Task: 将代码按逻辑分块，返回结构化的讲解单元列表。\n\n"
            "Rules:\n"
            "1. import 合并为一块；全局配置归 constant/config；函数/类/核心逻辑单独成块。\n"
            "2. title 简明概括，10字以内。\n"
            "3. code 保持原始缩进，一字不改。\n"
            "4. explanation 分两层，总字数40-150字：\n"
            "   ① 一句话概括整体作用（15字以内，句末加句号）。\n"
            "   ② 逐属性说明，格式：名称：作用，分号分隔。\n"
            "5. difficulty 按代码复杂度评分：1=简单，2=中等，3=复杂。\n"
            "6. HTML/XML 配对标签必须整体放入同一单元格。\n"
            "7. 代码注释已预先剥离，无需在 explanation 复述。\n"
            "8. code 严禁将 < > \" & 转义为 HTML 实体，直接输出原始字符。"
        )

        chunks = smart_split_code(fi["content_to_analyze"])
        all_cells: List[dict] = []

        for idx, chunk in enumerate(chunks, 1):
            chunk_label = f"（第 {idx}/{len(chunks)} 块）" if len(chunks) > 1 else ""
            user_msg = (
                f"文件：{fi['filename']}  语言：{lang}{chunk_label}\n"
                f"关键元素：\n{kw}\n\n原始代码：\n{chunk}"
            )
            try:
                result: AnalysisResult = llm_invoke_with_retry(
                    structured_llm,
                    [SystemMessage(content=system_msg), HumanMessage(content=user_msg)],
                )
                all_cells.extend(cell.model_dump() for cell in result.cells)
            except Exception as e:
                all_cells.append({
                    "type": "error", "title": f"解析失败（块{idx}）",
                    "code": chunk[:500],
                    "explanation": f"该块解析出错：{e}",
                    "difficulty": 1,
                })

        all_cells = merge_paired_cells(all_cells)
        return {"cells": all_cells}

    def summarizer(state: AnalysisState):
        fi   = state["file_info"]
        lang = fi["language"]
        code = fi["content_to_analyze"][:2000]
        system_msg = (
            f"Role: {lang}资深架构师。\n"
            "Task: 用3-5句话总结该文件。\n\n"
            "Rules:\n"
            "1. 涵盖：①文件用途 ②主要技术/库 ③核心逻辑 ④代码风格评价。\n"
            "2. 通俗易懂，120字以内。\n"
            "3. 禁止输出标题、序号或 Markdown，直接输出纯文本。"
        )
        user_msg = (
            f"文件：{fi['filename']}（{lang}，{fi['line_count']} 行）\n\n"
            f"代码预览：\n{code}"
        )
        result = llm_invoke_with_retry(
            llm, [SystemMessage(content=system_msg), HumanMessage(content=user_msg)]
        )
        return {"summary": result.content.strip()}

    g = StateGraph(AnalysisState)
    g.add_node("analyzer",   analyzer)
    g.add_node("summarizer", summarizer)
    g.set_entry_point("analyzer")
    g.add_edge("analyzer",   "summarizer")
    g.add_edge("summarizer", END)
    return g.compile(checkpointer=InMemorySaver())


@st.cache_resource
def get_cached_graph(_llm):
    """单例缓存分析图，避免每次重建。"""
    return build_analysis_graph(_llm)


# ════════════════════════════════════════════════════════════════════════════
# 历史记录管理
# ════════════════════════════════════════════════════════════════════════════
def save_history(file_info: dict, cells: list, summary: str,
                 elapsed: float = 0.0) -> Path:
    clean_cells = [{**c, "code": normalize_code(c.get("code", ""))} for c in cells]
    ts   = datetime.now().strftime("%Y%m%d_%H%M%S")
    slug = re.sub(r'[^\w\u4e00-\u9fff]', "_", Path(file_info["filename"]).stem)[:20]
    out  = HISTORY_DIR / f"{ts}_{slug}.json"
    out.write_text(json.dumps({
        "filename":   file_info["filename"],
        "language":   file_info["language"],
        "highlight":  file_info["highlight"],
        "timestamp":  datetime.now().isoformat(),
        "line_count": file_info["line_count"],
        "char_count": file_info["char_count"],
        "cell_count": len(clean_cells),
        "elapsed_s":  round(elapsed, 1),
        "summary":    summary,
        "cells":      clean_cells,
    }, ensure_ascii=False, indent=2), encoding="utf-8")
    return out

def list_histories() -> list:
    return sorted(HISTORY_DIR.glob("*.json"), reverse=True)

def load_history(p: Path) -> dict:
    return json.loads(p.read_text(encoding="utf-8"))

def delete_history(p: Path):
    if p.exists(): p.unlink()


# ════════════════════════════════════════════════════════════════════════════
# Markdown 导出
# ════════════════════════════════════════════════════════════════════════════
def export_markdown(filename: str, highlight: str, cells: list, summary: str) -> str:
    lines = [f"# 📓 {filename} — 代码讲解\n"]
    if summary:
        lines.append(f"## 📋 整体总结\n\n{summary}\n")
    for i, cell in enumerate(cells, 1):
        diff  = cell.get("difficulty", 1)
        stars = "⭐" * diff
        lines.append(
            f"## [{i}] {cell.get('title','代码块')} "
            f"`{cell.get('type','')}` {stars}\n"
        )
        lines.append(f"```{highlight}\n{normalize_code(cell.get('code',''))}\n```\n")
        if cell.get("explanation"):
            lines.append(f"> 💡 {cell['explanation']}\n")
    return "\n".join(lines)


# ════════════════════════════════════════════════════════════════════════════
# 追问 LLM
# ════════════════════════════════════════════════════════════════════════════
def followup_llm(llm, cell: dict, question: str, highlight: str) -> str:
    code = normalize_code(cell.get("code", ""))
    system_msg = (
        "你是一位专业的代码导师。根据提供的代码块和用户问题，给出清晰简洁的回答。"
        "使用中文回答，150字以内，可使用代码片段辅助说明。"
    )
    user_msg = (
        f"代码类型：{cell.get('type','')}，标题：{cell.get('title','')}\n\n"
        f"```{highlight}\n{code[:1000]}\n```\n\n"
        f"问题：{question}"
    )
    result = llm_invoke_with_retry(
        llm, [SystemMessage(content=system_msg), HumanMessage(content=user_msg)]
    )
    return result.content.strip()


# ════════════════════════════════════════════════════════════════════════════
# Notebook 渲染
# ════════════════════════════════════════════════════════════════════════════
def render_notebook(cells: list, highlight: str, summary: str = ""):
    css_content = load_css_content()

    summary_html = (
        f'<div class="nb-summary"><b>📋 整体总结</b><br><br>'
        f'{_html.escape(summary)}</div>'
    ) if summary else ""

    cells_html:   List[str] = []
    code_data_js: List[str] = []
    set_code_js:  List[str] = []
    total_height = (200 if summary else 20) + 64

    for i, cell in enumerate(cells, 1):
        ctype             = cell.get("type", "logic")
        icon, accent, hbg = CELL_STYLES.get(ctype, CELL_STYLES_DEFAULT)
        title             = cell.get("title", "代码块")
        code              = normalize_code(cell.get("code", ""))
        explanation       = cell.get("explanation", "")
        diff              = max(1, min(3, int(cell.get("difficulty", 1))))

        line_count   = max(len(code.splitlines()), 1)
        cell_height  = 52 + line_count * 24 + (52 if explanation else 0)
        total_height += cell_height + 26

        title_safe   = _html.escape(str(title))
        explain_html = (
            f'<div class="nb-explain">'
            f'<span class="nb-explain-icon">💡</span>'
            f'<span class="nb-explain-text">{_html.escape(explanation)}</span>'
            f'</div>'
        ) if explanation else ""

        search_text = _html.escape(f"{title} {ctype} {explanation} {code[:300]}")

        code_data_js.append(f"codeData[{i}] = {json.dumps(code)};")
        set_code_js.append(
            f"var el{i}=document.getElementById('nbcode_{i}');"
            f"if(el{i})el{i}.textContent=codeData[{i}];"
        )

        cells_html.append(f"""
        <div class="nb-cell-outer"
             style="animation-delay:{i*0.035:.2f}s;border-left:3px solid {accent};"
             data-search="{search_text}">
            <div class="nb-header" style="background:{hbg};">
                <span class="nb-in-label">In&nbsp;[{i}]</span>
                <span class="nb-icon">{icon}</span>
                <span class="nb-title" style="color:{accent};">{title_safe}</span>
                <span class="nb-type-badge"
                      style="color:{accent};border:1px solid {accent}33;">{ctype}</span>
                {difficulty_dots(diff)}
                <div class="nb-toolbar">
                    <button class="nb-toolbar-btn" id="copybtn_{i}"
                            onclick="copyCell({i})" title="复制代码">⎘ 复制</button>
                    <button class="nb-toolbar-btn"
                            onclick="toggleCell(this)" title="折叠">⌄</button>
                </div>
            </div>
            <div class="nb-cell-body">
                <pre class="nb-code-pre"><code class="language-{highlight}"
                     id="nbcode_{i}"></code></pre>
                {explain_html}
            </div>
        </div>""")

    full_html = f"""<!DOCTYPE html>
<html><head><meta charset="utf-8">
<link rel="stylesheet" href="{_HLJS_CSS}">
<style>{css_content}</style>
</head>
<body>
{summary_html}

<div class="nb-search-wrap">
    <input id="nb-search" class="nb-search-input"
           type="text" placeholder="🔍 搜索代码块标题、类型或内容...">
    <span id="nb-search-stats" class="nb-search-stats"></span>
    <button id="nb-search-clear" class="nb-search-clear"
            onclick="clearSearch()">✕</button>
</div>
<div id="nb-no-results" class="nb-no-results">未找到匹配的代码块 🔍</div>

{''.join(cells_html)}

<script src="{_HLJS_JS}"></script>
<script src="{_HLJS_LNJS}"></script>
<script>
var codeData = {{}};
{chr(10).join(code_data_js)}
{chr(10).join(set_code_js)}

document.querySelectorAll('pre.nb-code-pre code').forEach(function(el) {{
    hljs.highlightElement(el);
    hljs.lineNumbersBlock(el);
}});

// ── 搜索 ──────────────────────────────────────────────
var allCells    = document.querySelectorAll('.nb-cell-outer');
var searchInput = document.getElementById('nb-search');
var searchStats = document.getElementById('nb-search-stats');
var searchClear = document.getElementById('nb-search-clear');
var noResults   = document.getElementById('nb-no-results');

if (searchInput) {{
    searchInput.addEventListener('input', function() {{
        var q = this.value.toLowerCase().trim();
        var count = 0;
        allCells.forEach(function(cell) {{
            var match = !q || (cell.dataset.search||'').toLowerCase().includes(q);
            cell.style.display = match ? '' : 'none';
            if (match) count++;
        }});
        searchStats.textContent = q ? count+' / '+allCells.length+' 个结果' : '';
        searchClear.style.display = q ? 'inline' : 'none';
        noResults.style.display   = (q && count===0) ? 'block' : 'none';
    }});
}}

function clearSearch() {{
    searchInput.value = '';
    searchInput.dispatchEvent(new Event('input'));
    searchInput.focus();
}}

// ── 折叠/展开 ─────────────────────────────────────────
function toggleCell(btn) {{
    var body = btn.closest('.nb-cell-outer').querySelector('.nb-cell-body');
    if (!body) return;
    var collapsed = body.classList.contains('nb-collapsed');
    body.classList.toggle('nb-collapsed', !collapsed);
    btn.textContent = collapsed ? '⌄' : '⌃';
    btn.title       = collapsed ? '折叠' : '展开';
    // 通知父页面调整 iframe 高度
    window.parent.postMessage({{
        type: 'nb-resize',
        height: document.body.scrollHeight + 40
    }}, '*');
}}

// ── 复制代码 ──────────────────────────────────────────
function copyCell(idx) {{
    var text = codeData[idx] || '';
    var btn  = document.getElementById('copybtn_'+idx);
    if (!navigator.clipboard) {{
        var ta = document.createElement('textarea');
        ta.value = text;
        ta.style.position = 'fixed';
        ta.style.opacity  = '0';
        document.body.appendChild(ta);
        ta.select();
        document.execCommand('copy');
        document.body.removeChild(ta);
        flashCopied(btn); return;
    }}
    navigator.clipboard.writeText(text).then(function() {{ flashCopied(btn); }});
}}

function flashCopied(btn) {{
    if (!btn) return;
    var prev = btn.textContent;
    btn.textContent = '✓ 已复制';
    btn.classList.add('copied');
    setTimeout(function() {{
        btn.textContent = prev;
        btn.classList.remove('copied');
    }}, 1800);
}}

// ── iframe 高度自适应（postMessage 上报）────────────────
function reportHeight() {{
    window.parent.postMessage({{
        type: 'nb-resize',
        height: document.body.scrollHeight + 40
    }}, '*');
}}
window.addEventListener('load', reportHeight);
new ResizeObserver(reportHeight).observe(document.body);
</script>
</body></html>"""

    components.html(full_html, height=total_height + 60, scrolling=False)


def render_stream_box(slot, label: str, text: str):
    safe = text.replace("&","&amp;").replace("<","&lt;").replace(">","&gt;")
    slot.markdown(
        f'<div class="stream-box">'
        f'<div class="stream-label">{label}<span class="cursor"></span></div>'
        f'{safe}</div>',
        unsafe_allow_html=True,
    )

def render_skeleton(slot, n: int = 3):
    css_content = load_css_content()
    slot.components_html = None
    components.html(
        f"<!DOCTYPE html><html><head><meta charset='utf-8'>"
        f"<style>{css_content}</style></head><body>"
        f"{skeleton_html(n)}</body></html>",
        height=n * 160 + 20,
    )

def _badge(css_cls, icon, label):
    return f'<div class="{css_cls}"><span>{icon}</span><span>{label}</span></div>'


# ════════════════════════════════════════════════════════════════════════════
# Streamlit 应用初始化
# ════════════════════════════════════════════════════════════════════════════
st.set_page_config(page_title="代码讲解笔记本", page_icon="📓", layout="wide")
st.markdown(load_css(), unsafe_allow_html=True)

for key, val in {
    "nb_cells":      [],
    "nb_summary":    "",
    "nb_highlight":  "python",
    "nb_filename":   "",
    "is_analyzing":  False,
    "selected_model": "DeepSeek Chat (通用)",
    "followup_idx":  None,
    "followup_ans":  {},
}.items():
    if key not in st.session_state:
        st.session_state[key] = val

@st.cache_resource
def init_llm(model_name: str):
    return ChatDeepSeek(model=model_name, temperature=0.2)

llm = init_llm(MODEL_OPTIONS[st.session_state.selected_model])


# ════════════════════════════════════════════════════════════════════════════
# 侧边栏
# ════════════════════════════════════════════════════════════════════════════
with st.sidebar:
    st.markdown(
        """
        <div style="text-align:center;padding:6px 0 4px 0">
            <div style="font-size:2em">📓</div>
            <div style="font-weight:700;font-size:1.05em;margin-top:2px">代码讲解笔记本</div>
            <div style="font-size:0.74em;color:#888;margin-top:2px">
                上传代码 × 自动识别 × Notebook 讲解
            </div>
        </div>
        """, unsafe_allow_html=True,
    )
    st.divider()

    # 模型选择
    st.markdown("#### 🤖 模型选择")
    selected = st.selectbox(
        "选择模型",
        options=list(MODEL_OPTIONS.keys()),
        index=list(MODEL_OPTIONS.keys()).index(st.session_state.selected_model),
        disabled=st.session_state.is_analyzing,
        label_visibility="collapsed",
    )
    if selected != st.session_state.selected_model:
        st.session_state.selected_model = selected
        st.rerun()
    st.caption(f"当前：`{MODEL_OPTIONS[selected]}`")

    st.divider()

    with st.expander("📂 支持的语言", expanded=False):
        lang_groups = {
            "Web":    [".html",".css",".js",".ts",".vue"],
            "系统级": [".c",".cpp",".cs",".rs",".go"],
            "JVM":    [".java",".kt",".scala"],
            "脚本":   [".py",".rb",".php",".lua",".jl"],
            "Shell":  [".sh",".ps1",".bat"],
            "移动端": [".swift",".dart"],
            "数据":   [".sql",".graphql"],
            "配置":   [".yaml",".toml",".xml",".tf"],
            "函数式": [".hs",".ex",".erl",".clj",".ml"],
            "其他":   [".r",".vb",".asm",".nim"],
        }
        for grp, exts in lang_groups.items():
            st.caption(f"**{grp}**: " + "  ".join(f"`{e}`" for e in exts))

    st.divider()

    ch, cb = st.columns([3, 1])
    with ch: st.markdown("#### 🗂️ 历史记录")
    with cb:
        if st.session_state.nb_cells and not st.session_state.is_analyzing:
            if st.button("🏠", help="返回主页", use_container_width=True):
                st.session_state.nb_cells    = []
                st.session_state.nb_summary  = ""
                st.session_state.nb_filename = ""
                st.session_state.followup_ans = {}
                st.rerun()

    histories = list_histories()
    if not histories:
        st.caption("暂无历史记录")
    else:
        if st.button("🗑️ 删除全部", use_container_width=True,
                     disabled=st.session_state.is_analyzing):
            for p in histories: delete_history(p)
            st.toast("✅ 已清空历史记录"); st.rerun()

        for p in histories[:20]:
            parts = p.stem.split("_", 2)
            label = parts[2].replace("_", " ") if len(parts) > 2 else p.stem
            ts    = f"{parts[0][4:6]}/{parts[0][6:8]}" if len(parts[0]) >= 8 else ""
            cr, cd = st.columns([5, 1], vertical_alignment="center")
            with cr:
                if st.button(f"📄 {label[:12]}  {ts}", key=f"v_{p}",
                             use_container_width=True,
                             disabled=st.session_state.is_analyzing):
                    h = load_history(p)
                    st.session_state.nb_cells     = h["cells"]
                    st.session_state.nb_summary   = h["summary"]
                    st.session_state.nb_highlight = h["highlight"]
                    st.session_state.nb_filename  = h["filename"]
                    st.session_state.followup_ans = {}
                    st.rerun()
            with cd:
                if st.button("🗑️", key=f"d_{p}", use_container_width=True,
                             help="删除", disabled=st.session_state.is_analyzing):
                    delete_history(p)
                    st.session_state.nb_cells = []
                    st.toast(f"已删除 {label[:12]}"); st.rerun()


# ════════════════════════════════════════════════════════════════════════════
# 主界面
# ════════════════════════════════════════════════════════════════════════════
st.markdown("## 📓 代码讲解笔记本")
st.caption("上传代码文件 → AI 自动拆成讲解单元 → Notebook 格式逐块展示，每块附通俗解释与难度标注")

uploaded = st.file_uploader(
    "上传代码文件",
    type=[e.lstrip(".") for e in LANG_MAP.keys()],
    disabled=st.session_state.is_analyzing,
    label_visibility="collapsed",
)
run_btn = st.button(
    "🚀 开始讲解", type="primary",
    disabled=(uploaded is None or st.session_state.is_analyzing),
)
st.divider()

cols         = st.columns(4)
status_slots = {k: cols[i].empty() for i, k in enumerate(NODE_LABELS)}
report_box   = st.empty()


def reset_statuses(active=False):
    icon = "⏳" if active else "⬜"
    for k, s in status_slots.items():
        s.markdown(_badge("badge-idle", icon, NODE_LABELS[k]), unsafe_allow_html=True)

def mark_running(node, label=None):
    if node in status_slots:
        status_slots[node].markdown(
            _badge("badge-running", "⚙️", label or NODE_LABELS[node]),
            unsafe_allow_html=True,
        )

def mark_complete(node):
    if node in status_slots:
        status_slots[node].markdown(
            _badge("badge-done", "✅", NODE_LABELS[node]),
            unsafe_allow_html=True,
        )


reset_statuses()

# 展示已有结果
if st.session_state.nb_cells and not run_btn:
    # 导出工具栏
    ex_md, ex_info = st.columns([1, 6])
    with ex_md:
        md_str = export_markdown(
            st.session_state.nb_filename,
            st.session_state.nb_highlight,
            st.session_state.nb_cells,
            st.session_state.nb_summary,
        )
        st.download_button(
            "⬇️ 导出 Markdown",
            data=md_str.encode("utf-8"),
            file_name=f"{Path(st.session_state.nb_filename).stem}_讲解.md",
            mime="text/markdown",
            use_container_width=True,
        )
    with ex_info:
        cell_count = len(st.session_state.nb_cells)
        diff_counts = {1: 0, 2: 0, 3: 0}
        for c in st.session_state.nb_cells:
            diff_counts[max(1, min(3, int(c.get("difficulty", 1))))] += 1
        st.caption(
            f"共 **{cell_count}** 个单元格　"
            f"🟢 简单 {diff_counts[1]}　"
            f"🟡 中等 {diff_counts[2]}　"
            f"🔴 复杂 {diff_counts[3]}"
        )

    with report_box.container():
        render_notebook(
            st.session_state.nb_cells,
            st.session_state.nb_highlight,
            st.session_state.nb_summary,
        )

    # 追问区
    st.divider()
    st.markdown("#### 💬 追问代码块")
    st.caption("选择一个代码块，针对它向 AI 提问")

    cell_options = {
        f"[{i+1}] {c.get('title','代码块')} ({c.get('type','')})": i
        for i, c in enumerate(st.session_state.nb_cells)
    }
    col_sel, col_q, col_btn = st.columns([3, 5, 1])
    with col_sel:
        chosen_label = st.selectbox(
            "选择代码块", options=list(cell_options.keys()),
            label_visibility="collapsed",
        )
    chosen_idx = cell_options[chosen_label]
    with col_q:
        question = st.text_input(
            "问题", placeholder="例：这个函数的时间复杂度是多少？",
            label_visibility="collapsed",
        )
    with col_btn:
        ask_btn = st.button("提问", type="secondary", use_container_width=True)

    if ask_btn and question.strip():
        with st.spinner("AI 思考中..."):
            try:
                ans = followup_llm(
                    llm,
                    st.session_state.nb_cells[chosen_idx],
                    question.strip(),
                    st.session_state.nb_highlight,
                )
                st.session_state.followup_ans[chosen_idx] = {
                    "q": question.strip(), "a": ans
                }
            except Exception as e:
                st.error(f"追问失败：{e}")

    if chosen_idx in st.session_state.followup_ans:
        qa = st.session_state.followup_ans[chosen_idx]
        # Q：转义显示，A：交由 st.markdown 渲染保留代码块格式
        st.markdown(
            f'<div style="'
            f'border-left:3px solid #7aa2f7;'
            f'background:rgba(122,162,247,0.07);'
            f'border-radius:6px;'
            f'padding:10px 16px 8px 16px;'
            f'margin-top:6px;">'
            f'<span style="color:#7aa2f7;font-weight:600;">Q：</span>'
            f'{_html.escape(qa["q"])}'
            f'</div>',
            unsafe_allow_html=True,
        )
        st.markdown(
            f'<div style="'
            f'border-left:3px solid #9ece6a;'
            f'background:rgba(158,206,106,0.07);'
            f'border-radius:6px 6px 0 0;'
            f'padding:8px 16px 2px 16px;'
            f'margin-top:4px;">'
            f'<span style="color:#9ece6a;font-weight:600;">A：</span>'
            f'</div>',
            unsafe_allow_html=True,
        )
        st.markdown(qa["a"])



# ════════════════════════════════════════════════════════════════════════════
# 分析流程
# ════════════════════════════════════════════════════════════════════════════
if run_btn and uploaded:
    st.session_state.is_analyzing = True
    st.session_state.nb_cells     = []
    st.session_state.nb_summary   = ""
    st.session_state.followup_ans = {}
    report_box.empty()
    reset_statuses(active=True)
    t_start = time.time()

    try:
        mark_running("file_loader")
        file_info = load_code_file(uploaded)
        if "error" in file_info:
            st.error(f"❌ 读取失败：{file_info['error']}"); st.stop()
        mark_complete("file_loader")

        mark_running("preprocessor")
        preprocessed = preprocess_and_divide(file_info)
        mark_complete("preprocessor")

        chunks     = smart_split_code(file_info["content_to_analyze"])
        chunk_info = f"（将分 {len(chunks)} 块分析）" if len(chunks) > 1 else ""
        with report_box.container():
            st.markdown(
                f'<div class="info-card">'
                f'📄 <b>{file_info["filename"]}</b>　'
                f'语言：<b>{file_info["language"]}</b>　'
                f'{file_info["line_count"]} 行 / {file_info["char_count"]} 字符'
                f'{chunk_info}</div>',
                unsafe_allow_html=True,
            )

        # 骨架屏
        skel_slot = st.empty()
        css_c = load_css_content()
        components.html(
            f"<!DOCTYPE html><html><head><meta charset='utf-8'>"
            f"<style>{css_c}</style></head><body>"
            f"{skeleton_html(min(len(chunks)*2, 5))}</body></html>",
            height=min(len(chunks)*2, 5) * 160 + 20,
        )

        mark_running("analyzer", "🧩 拆分代码单元...")
        app    = get_cached_graph(llm)
        config = {"configurable": {"thread_id": f"run_{int(time.time())}"}}
        init_s: AnalysisState = {
            "file_info":    file_info,
            "preprocessed": preprocessed,
            "cells":        [],
            "summary":      "",
        }

        current_node = ""; current_text = ""; last_render = 0.0

        for chunk in app.stream(init_s, config, stream_mode="messages", version="v2"):
            try:
                if chunk.get("type") != "messages": continue
                msg, metadata = chunk["data"]
                node = metadata.get("langgraph_node", "")
                if node != current_node:
                    if current_node == "analyzer":
                        mark_complete("analyzer")
                        mark_running("summarizer", "📋 生成总结...")
                        current_text = ""
                    current_node = node; last_render = 0.0
                if node == "analyzer": continue
                if hasattr(msg, "content") and msg.content:
                    current_text += msg.content
                    now = time.time()
                    if now - last_render >= STREAM_THROTTLE:
                        render_stream_box(report_box, "📋 正在生成总结...", current_text)
                        last_render = now
            except (KeyError, ValueError, TypeError): continue

        mark_complete("analyzer"); mark_complete("summarizer")
        skel_slot.empty()

        final   = app.get_state(config)
        cells   = final.values.get("cells",   [])
        summary = final.values.get("summary", "")
        elapsed = time.time() - t_start

        if not cells:
            st.warning("⚠️ 单元格解析失败，请重试")
        else:
            save_history(file_info, cells, summary, elapsed)
            st.session_state.nb_cells    = cells
            st.session_state.nb_summary  = summary
            st.session_state.nb_highlight= file_info["highlight"]
            st.session_state.nb_filename = file_info["filename"]

            report_box.empty()
            with report_box.container():
                render_notebook(cells, file_info["highlight"], summary)

            st.success(
                f"✅ 讲解完成！共 {len(cells)} 个单元格，"
                f"耗时 {elapsed:.1f}s"
            )
            st.balloons()

    except Exception as e:
        st.error(f"❌ 分析出错：{e}")
    finally:
        st.session_state.is_analyzing = False
        st.rerun()