"""
VoiceOfML Search — FastAPI 后端 (最终版)

功能:
  - 启动时加载 search_data.json 到内存
  - 提供搜索 API（全局 + 单仓库）+ 拼写纠错
  - 多维度过滤 + 分页
  - 文件夹内容浏览 API（进入式）
  - 目录树 API（过滤面板）
  - 流式代理下载（解决中文文件名乱码）
  - 在线阅读 txt 服务
  - 静态文件服务
"""

import asyncio
import gzip
import json
import posixpath
import random
import re
import time
from datetime import datetime
from pathlib import Path
from typing import Optional
from urllib.parse import unquote, quote, urlparse

import aiohttp
import jieba
from fastapi import FastAPI, Query
from fastapi.responses import FileResponse, HTMLResponse, JSONResponse, PlainTextResponse, StreamingResponse
from fastapi.staticfiles import StaticFiles
from fastapi.middleware.cors import CORSMiddleware
from fastapi.middleware.gzip import GZipMiddleware
from pydantic import BaseModel
from contextlib import asynccontextmanager
from stream_zip import stream_zip, ZIP_32

# ═══════════════════════════════════════════════════════════
# 启动时加载数据
# ═══════════════════════════════════════════════════════════

DATA_PATH = Path("data/search_data.json.gz")
FOLDER_TREE_PATH = Path("data/folder_tree.json.gz")
FOLDER_BROWSER_PATH = Path("data/folder_browser.json.gz")
records: list[dict] = []
folder_tree_data: dict[str, list[dict]] = {}
folder_browser_data: dict[str, dict[str, dict]] = {}

word_index = {}
folder_index = {}
extension_counts = {}
repo_counts = {}
repo_list = []
extension_list = []
did_you_mean_vocab = {}
vocab_by_len = {}
repo_records_map = {}
sorted_by_name = []
sorted_by_size = []
repo_sorted_by_name = {}
repo_sorted_by_size = {}

RECORD_KEY_MAP = {
    "r": "Repo",
    "f": "File",
    "e": "Extension",
    "d": "Folder",
    "s": "Size",
    "t": "HasTxt",
}

TREE_KEY_MAP = {
    "n": "name",
    "c": "count",
    "df": "hasDirectFiles",
    "ch": "children",
}

BROWSER_ENTRY_KEY_MAP = {
    "d": "folders",
    "f": "files",
}

BROWSER_FOLDER_ITEM_KEY_MAP = {
    "n": "name",
    "c": "count",
}

BROWSER_FILE_ITEM_KEY_MAP = {
    "n": "name",
    "e": "ext",
    "t": "hasTxt",
    "s": "size",
}


def build_record_relative_path(rec: dict) -> str:
    file_name = rec.get("File", "")
    extension = rec.get("Extension", "")
    full_name = f"{file_name}.{extension}" if extension else file_name
    folders = rec.get("Folder", []) or []
    return posixpath.join(*folders, full_name) if folders else full_name


def build_record_link(rec: dict) -> str:
    repo = rec.get("Repo", "")
    rel_path = build_record_relative_path(rec)
    return f"https://huggingface.co/datasets/{repo}/resolve/main/{quote(rel_path, safe='/')}"


def build_record_path_url(rec: dict) -> str:
    repo = rec.get("Repo", "")
    rel_path = build_record_relative_path(rec)
    return f"https://huggingface.co/datasets/{repo}/blob/main/{quote(rel_path, safe='/')}"


def build_browser_file_name(file_item: dict) -> str:
    name = str(file_item.get("name", "") or "file")
    ext = str(file_item.get("ext", "") or "")
    if not ext:
        return name
    if name.lower().endswith(f".{ext.lower()}"):
        return name
    return f"{name}.{ext}"


def enrich_browser_entry(repo: str, path: str, entry: dict) -> dict:
    decoded = dict(entry or {})
    decoded["current_path"] = decoded.get("current_path", path)
    folders = []
    for folder in decoded.get("folders", []) or []:
        item = dict(folder)
        if "path" not in item:
            name = item.get("name", "")
            item["path"] = f"{path}/{name}" if path else name
        folders.append(item)
    files = []
    for file_item in decoded.get("files", []) or []:
        item = dict(file_item)
        full_name = build_browser_file_name(item)
        rel_path = f"{path}/{full_name}" if path else full_name
        item["link"] = item.get("link") or f"https://huggingface.co/datasets/{repo}/resolve/main/{quote(rel_path, safe='/')}"
        item["path"] = item.get("path") or f"https://huggingface.co/datasets/{repo}/blob/main/{quote(rel_path, safe='/')}"
        files.append(item)
    decoded["folders"] = folders
    decoded["files"] = files
    return decoded


def decode_record(rec: dict) -> dict:
    if "Repo" in rec:
        return rec
    if isinstance(rec, list):
        return rec
    decoded = {}
    for short_key, full_key in RECORD_KEY_MAP.items():
        if short_key in rec:
            decoded[full_key] = rec[short_key]
    decoded.setdefault("Repo", "")
    decoded.setdefault("File", "")
    decoded.setdefault("Extension", "")
    decoded.setdefault("Folder", [])
    decoded.setdefault("Size", "")
    decoded.setdefault("HasTxt", False)
    return decoded


def decode_search_payload(data) -> list[dict]:
    if isinstance(data, list):
        if not data:
            return []
        if isinstance(data[0], list):
            return []
        return [decode_record(rec) for rec in data]
    if not isinstance(data, dict):
        return []
    repos = data.get("rp", []) or []
    folders = data.get("fd", []) or []
    records = []
    for item in data.get("rc", []) or []:
        if not isinstance(item, list) or len(item) < 6:
            continue
        repo = repos[item[0]] if isinstance(item[0], int) and 0 <= item[0] < len(repos) else ""
        folder = folders[item[3]] if isinstance(item[3], int) and 0 <= item[3] < len(folders) else []
        records.append({
            "Repo": repo,
            "File": item[1],
            "Extension": item[2],
            "Folder": folder,
            "Size": item[4],
            "HasTxt": bool(item[5]),
        })
    return records


def decode_tree_node(node: dict, parent_path: str = "", is_root: bool = False) -> dict:
    if "name" in node:
        decoded = dict(node)
        path = decoded.get("path")
        if path is None:
            name = decoded.get("name", "")
            decoded["path"] = "" if is_root else (f"{parent_path}/{name}" if parent_path else name)
        decoded["children"] = [
            decode_tree_node(child, decoded["path"], False)
            for child in decoded.get("children", [])
        ]
        decoded["hasChildren"] = len(decoded["children"]) > 0
        decoded["showSelfToggle"] = bool(decoded["path"] and decoded.get("hasDirectFiles") and decoded["hasChildren"])
        if is_root:
            decoded["isRoot"] = True
        return decoded
    name = node.get("n", "")
    path = "" if is_root else (f"{parent_path}/{name}" if parent_path else name)
    children = [decode_tree_node(child, path, False) for child in node.get("ch", [])]
    decoded = {
        "name": name,
        "path": path,
        "count": node.get("c", 0),
        "hasDirectFiles": node.get("df", False),
        "hasChildren": len(children) > 0,
        "showSelfToggle": bool(path and node.get("df", False) and len(children) > 0),
        "children": children,
    }
    if is_root:
        decoded["isRoot"] = True
    return decoded


def decode_folder_tree_data(data: dict) -> dict:
    decoded = {}
    for repo, nodes in (data or {}).items():
        decoded[repo] = [decode_tree_node(node, "", True) for node in nodes]
    return decoded


def decode_browser_entry(entry: dict, current_path: str = "") -> dict:
    if "folders" in entry:
        decoded = dict(entry)
        decoded.setdefault("current_path", current_path)
        for folder in decoded.get("folders", []):
            if "path" not in folder:
                name = folder.get("name", "")
                folder["path"] = f"{current_path}/{name}" if current_path else name
        return decoded
    return {
        "folders": [
            {
                "name": item.get("n", ""),
                "path": f"{current_path}/{item.get('n', '')}" if current_path else item.get("n", ""),
                "count": item.get("c", 0),
            }
            for item in entry.get("d", [])
        ],
        "files": [
            {
                "name": item.get("n", ""),
                "ext": item.get("e", ""),
                "hasTxt": item.get("t", False),
                "size": item.get("s", ""),
            }
            for item in entry.get("f", [])
        ],
        "current_path": current_path,
    }


def decode_folder_browser_data(data: dict) -> dict:
    decoded = {}
    for repo, repo_browser in (data or {}).items():
        decoded[repo] = {}
        for path, entry in (repo_browser or {}).items():
            decoded[repo][path] = decode_browser_entry(entry, path)
    return decoded

def tokenize(text: str) -> list[str]:
    tokens = []
    text_lower = text.lower()

    # 英文/数字 token
    alpha_tokens = re.findall(r"[a-z0-9]+", text_lower)
    tokens.extend(alpha_tokens)

    # 中文用 jieba 分词
    chinese_text = re.sub(r"[a-z0-9\s]+", " ", text_lower)
    chinese_text = re.sub(r"[^\u4e00-\u9fff\u3400-\u4dbf\s]+", " ", chinese_text)
    chinese_tokens = jieba.lcut(chinese_text)
    tokens.extend([t.strip() for t in chinese_tokens if t.strip()])

    # 兜底：单字切割（保留 jieba 未能覆盖的字符）
    chinese_tokens_str = " ".join(chinese_tokens)
    for ch in text_lower:
        if "\u4e00" <= ch <= "\u9fff" or "\u3400" <= ch <= "\u4dbf":
            if ch not in chinese_tokens_str:
                tokens.append(ch)

    return list(set(tokens))


def matches_exact_query(text: str, query: str) -> bool:
    if "*" in query or "?" in query:
        pattern = re.escape(query).replace(r"\*", ".*").replace(r"\?", ".")
        return re.search(pattern, text, re.IGNORECASE) is not None
    return query.lower() in text.lower()

def edit_distance(s1: str, s2: str, max_dist: int = 2) -> int:
    if abs(len(s1) - len(s2)) > max_dist:
        return 999
    prev = list(range(len(s2) + 1))
    for i, c1 in enumerate(s1):
        curr = [i + 1]
        for j, c2 in enumerate(s2):
            cost = 0 if c1 == c2 else 1
            curr.append(min(prev[j + 1] + 1, curr[j] + 1, prev[j] + cost))
        if min(curr) > max_dist:
            return 999
        prev = curr
    return prev[-1]

def build_indexes():
    global word_index, folder_index, extension_counts, repo_counts, repo_list, extension_list, did_you_mean_vocab, repo_records_map, sorted_by_name, sorted_by_size, repo_sorted_by_name, repo_sorted_by_size
    word_index = {}
    folder_index = {}
    extension_counts = {}
    repo_counts = {}
    did_you_mean_vocab = {}
    repo_records_map = {}
    repo_sorted_by_name = {}
    repo_sorted_by_size = {}

    for idx, rec in enumerate(records):
        repo = rec.get("Repo", "")
        repo_counts[repo] = repo_counts.get(repo, 0) + 1
        ext = (rec.get("Extension") or "").lower()
        if ext:
            extension_counts[ext] = extension_counts.get(ext, 0) + 1

        file_name = rec.get("File", "")
        folders = rec.get("Folder", [])
        text_parts = [file_name] + folders
        full_text = " ".join(text_parts)
        tokens = tokenize(full_text)
        for tok in tokens:
            word_index.setdefault(tok, set()).add(idx)
            did_you_mean_vocab[tok] = did_you_mean_vocab.get(tok, 0) + 1

        if repo not in folder_index:
            folder_index[repo] = {}
        for depth in range(len(folders) + 1):
            folder_path = "/".join(folders[:depth]) if depth > 0 else ""
            folder_index[repo].setdefault(folder_path, []).append(idx)

        repo_records_map.setdefault(repo, []).append(idx)
        folders_lower = [f.lower() for f in (folders or [])]
        rec["_file_lower"] = file_name.lower()
        rec["_folders_lower"] = folders_lower
        rec["_folder_text"] = " ".join(folders_lower)
        rec["_repo_lower"] = repo.lower()

    repo_list = sorted(
        [{"name": r, "count": c} for r, c in repo_counts.items()],
        key=lambda x: x["name"],
    )
    global vocab_by_len
    vocab_by_len = {}
    for word, freq in did_you_mean_vocab.items():
        wlen = len(word)
        vocab_by_len.setdefault(wlen, {})[word] = freq
    extension_list = sorted(extension_counts.keys())
    all_indices = list(range(len(records)))
    sorted_by_name = sorted(all_indices, key=lambda i: (records[i].get("File") or "").lower())
    sorted_by_size = sorted(all_indices, key=lambda i: (
        -(records[i].get("Size") if isinstance(records[i].get("Size"), (int, float)) else 0),
        (records[i].get("File") or "").lower(),
    ))
    for repo, repo_indices in repo_records_map.items():
        repo_sorted_by_name[repo] = sorted(repo_indices, key=lambda i: (records[i].get("File") or "").lower())
        repo_sorted_by_size[repo] = sorted(repo_indices, key=lambda i: (
            -(records[i].get("Size") if isinstance(records[i].get("Size"), (int, float)) else 0),
            (records[i].get("File") or "").lower(),
        ))

def load_data():
    global records, folder_tree_data, folder_browser_data
    start = time.time()
    if DATA_PATH.exists():
        raw = gzip.decompress(DATA_PATH.read_bytes())
        records = decode_search_payload(json.loads(raw.decode("utf-8")))
    else:
        print("⚠ 数据文件不存在")
        records = []
        return

    if FOLDER_TREE_PATH.exists():
        folder_tree_data = decode_folder_tree_data(json.loads(gzip.decompress(FOLDER_TREE_PATH.read_bytes()).decode("utf-8")))
    else:
        folder_tree_data = {}

    if FOLDER_BROWSER_PATH.exists():
        folder_browser_data = decode_folder_browser_data(json.loads(gzip.decompress(FOLDER_BROWSER_PATH.read_bytes()).decode("utf-8")))
    else:
        folder_browser_data = {}

    print(f"📖 已加载 {len(records)} 条记录 ({time.time() - start:.2f}s)")
    build_indexes()

# ═══════════════════════════════════════════════════════════
# 搜索核心
# ═══════════════════════════════════════════════════════════

def score_record(rec_idx, query_tokens, search_folders=True):
    rec = records[rec_idx]
    score = 0
    file_name = rec.get("_file_lower", "")
    repo = rec.get("_repo_lower", "")
    folder_text = rec.get("_folder_text", "")
    for tok in query_tokens:
        if tok in file_name:
            score += 3
        if search_folders and tok in folder_text:
            score += 2
        if tok in repo:
            score += 1
    return score

def apply_filters(indices, repos=None, extensions=None, folders=None, min_size=None, max_size=None, folder_match_mode="prefix"):
    extensions_lower = [e.lower() for e in extensions] if extensions else None
    result = []
    for idx in indices:
        rec = records[idx]
        if repos and rec.get("Repo") not in repos:
            continue
        if extensions_lower:
            ext = (rec.get("Extension") or "").lower()
            if ext not in extensions_lower:
                continue
        if folders:
            rec_folders = rec.get("Folder", []) or []
            matched = False
            for f in folders:
                f_clean = f.strip("/")
                if folder_match_mode == "exact":
                    if "/".join(rec_folders) == f_clean:
                        matched = True
                        break
                elif f_clean == "" and len(rec_folders) == 0:
                    matched = True
                    break
                else:
                    for depth in range(len(rec_folders) + 1):
                        prefix = "/".join(rec_folders[:depth]) if depth > 0 else ""
                        if prefix == f_clean:
                            matched = True
                            break
                    if matched:
                        break
            if not matched:
                continue
        size = rec.get("Size")
        if size and isinstance(size, (int, float)) and size > 0:
            if min_size is not None and size < min_size:
                continue
            if max_size is not None and size > max_size:
                continue
        result.append(idx)
    return result

def trim_results(raw):
    keys = ("Repo", "File", "Extension", "Link", "Path", "Folder", "Size", "HasTxt")
    trimmed = []
    for r in raw:
        item = {k: r.get(k, "" if k != "Folder" else []) for k in keys}
        item["Link"] = r.get("Link") or build_record_link(r)
        item["Path"] = r.get("Path") or build_record_path_url(r)
        trimmed.append(item)
    return trimmed


def sort_key_name(idx: int):
    return (records[idx].get("File") or "").lower()


def sort_key_size(idx: int):
    return (
        -(records[idx].get("Size") if isinstance(records[idx].get("Size"), (int, float)) else 0),
        (records[idx].get("File") or "").lower(),
    )


def paginate_indices(indices: list[int], page: int, page_size: int) -> tuple[int, list[dict]]:
    total = len(indices)
    start = (page - 1) * page_size
    return total, [records[i] for i in indices[start:start + page_size]]


def build_response(results: list[dict], total: int, page: int, page_size: int, did_you_mean=None) -> dict:
    return {
        "results": trim_results(results),
        "total": total,
        "page": page,
        "page_size": page_size,
        "did_you_mean": did_you_mean,
    }

def apply_mixed_folder_filters(indices, self_folders=None, subtree_folders=None):
    self_set = {str(path).strip("/") for path in (self_folders or []) if str(path).strip("/")}
    subtree_set = {str(path).strip("/") for path in (subtree_folders or []) if str(path).strip("/")}
    if not self_set and not subtree_set:
        return indices

    result = []
    for idx in indices:
        rec = records[idx]
        rec_folders = rec.get("Folder", []) or []
        rec_path = "/".join(rec_folders)
        if rec_path in self_set:
            result.append(idx)
            continue
        matched = False
        for depth in range(1, len(rec_folders) + 1):
            prefix = "/".join(rec_folders[:depth])
            if prefix in subtree_set:
                matched = True
                break
        if matched:
            result.append(idx)
    return result

def search(q="", repos=None, extensions=None, folders=None, min_size=None, max_size=None, page=1, page_size=100, search_folders=True, sort="relevance", exact=False, folder_match_mode="prefix", folder_selfs=None, folder_subtrees=None):
    q = q.strip()

    # ── 精准搜索 ──
    if exact and q:
        matched_indices = []
        for idx, rec in enumerate(records):
            file_name = rec.get("_file_lower", "")
            folder_path = "/".join(rec.get("Folder", [])).lower() if rec.get("Folder") else ""
            repo_name = rec.get("_repo_lower", "")
            if matches_exact_query(file_name, q) or matches_exact_query(folder_path, q) or matches_exact_query(repo_name, q):
                matched_indices.append(idx)
        if folder_match_mode == "mixed":
            filtered = apply_mixed_folder_filters(matched_indices, folder_selfs, folder_subtrees)
            filtered = apply_filters(filtered, repos, extensions, None, min_size, max_size)
        else:
            filtered = apply_filters(matched_indices, repos, extensions, folders, min_size, max_size, folder_match_mode)
        if sort == "name":
            filtered.sort(key=sort_key_name)
        elif sort == "size":
            filtered.sort(key=sort_key_size)
        total, results = paginate_indices(filtered, page, page_size)
        return build_response(results, total, page, page_size)

    # ── 空查询 ──
    if not q:
        has_filters = repos or extensions or folders or min_size is not None or max_size is not None
        if not has_filters:
            if repos and len(repos) == 1:
                repo = repos[0]
                if sort == "name":
                    indices = repo_sorted_by_name.get(repo, [])
                elif sort == "size":
                    indices = repo_sorted_by_size.get(repo, [])
                else:
                    indices = repo_records_map.get(repo, [])
            else:
                if sort == "name":
                    indices = sorted_by_name
                elif sort == "size":
                    indices = sorted_by_size
                else:
                    indices = list(range(len(records)))
            total, results = paginate_indices(indices, page, page_size)
        else:
            if repos:
                indices = []
                for r in repos:
                    indices.extend(repo_records_map.get(r, []))
            else:
                indices = list(range(len(records)))
            if folder_match_mode == "mixed":
                filtered = apply_mixed_folder_filters(indices, folder_selfs, folder_subtrees)
                filtered = apply_filters(filtered, repos, extensions, None, min_size, max_size)
            else:
                filtered = apply_filters(indices, repos, extensions, folders, min_size, max_size, folder_match_mode)
            if sort == "name":
                filtered.sort(key=sort_key_name)
            elif sort == "size":
                filtered.sort(key=sort_key_size)
            total, results = paginate_indices(filtered, page, page_size)
        return build_response(results, total, page, page_size)

    # ── 普通搜索 ──
    query_tokens = tokenize(q)

    # 精确匹配：按最小集合优先取交集
    exact_matches = None
    exact_tokens = [(tok, word_index[tok]) for tok in query_tokens if tok in word_index]
    if exact_tokens:
        exact_tokens.sort(key=lambda x: len(x[1]))
        exact_matches = exact_tokens[0][1].copy()
        for _, idx_set in exact_tokens[1:]:
            exact_matches.intersection_update(idx_set)

    # 模糊匹配，同时缓存结果给 did_you_mean 复用
    fuzzy_cache = {}
    fuzzy_matches = set()
    for i, tok in enumerate(query_tokens):
        if tok in word_index:
            continue
        candidates = set()
        best_word = None
        best_dist = 999
        tok_len = len(tok)
        for delta in (-2, -1, 0, 1, 2):
            for vocab_word in vocab_by_len.get(tok_len + delta, {}):
                dist = edit_distance(tok, vocab_word, max_dist=2)
                if dist <= 2:
                    candidates.update(word_index.get(vocab_word, set()))
                    if dist < best_dist:
                        best_dist = dist
                        best_word = vocab_word
                    if len(candidates) >= 200:
                        break
            if len(candidates) >= 200:
                break
        if candidates:
            fuzzy_cache[tok] = (best_word, best_dist)
            if i == 0 and not exact_matches:
                fuzzy_matches = candidates
            else:
                fuzzy_matches.update(candidates)

    if exact_matches:
        all_matched = list(exact_matches)
        if fuzzy_matches:
            all_matched.extend([i for i in fuzzy_matches if i not in exact_matches])
    elif fuzzy_matches:
        all_matched = list(fuzzy_matches)
    else:
        all_matched = []

    # 关闭文件夹搜索时剔除仅命中文件夹的记录
    if not search_folders and query_tokens and all_matched:
        filtered_by_file = []
        for idx in all_matched:
            rec = records[idx]
            file_name = rec.get("_file_lower", "")
            repo = rec.get("_repo_lower", "")
            if all(tok in file_name or tok in repo for tok in query_tokens):
                filtered_by_file.append(idx)
        all_matched = filtered_by_file

    if folder_match_mode == "mixed":
        filtered = apply_mixed_folder_filters(all_matched, folder_selfs, folder_subtrees)
        filtered = apply_filters(filtered, repos, extensions, None, min_size, max_size)
    else:
        filtered = apply_filters(all_matched, repos, extensions, folders, min_size, max_size, folder_match_mode)

    # 打分 + 排序
    scored = [(idx, score_record(idx, query_tokens, search_folders)) for idx in filtered]

    if sort == "relevance":
        scored.sort(key=lambda x: -x[1])
    elif sort == "name":
        scored.sort(key=lambda x: sort_key_name(x[0]))
    elif sort == "size":
        scored.sort(key=lambda x: sort_key_size(x[0]))

    total = len(scored)
    start = (page - 1) * page_size
    paged = scored[start:start + page_size]
    results = [records[idx] for idx, _ in paged]

    # did_you_mean：复用模糊匹配阶段缓存的编辑距离结果
    did_you_mean = None
    if total < 10 and query_tokens:
        suggestions = []
        for tok in query_tokens:
            if tok in word_index:
                continue
            if tok in fuzzy_cache:
                best_word, best_dist = fuzzy_cache[tok]
                if best_word and best_dist <= 2:
                    suggestions.append(best_word)
            else:
                best_word = None
                best_dist = 999
                for vocab_word in did_you_mean_vocab:
                    if abs(len(vocab_word) - len(tok)) > 2:
                        continue
                    dist = edit_distance(tok, vocab_word, max_dist=2)
                    if dist < best_dist:
                        best_dist = dist
                        best_word = vocab_word
                if best_word and best_dist <= 2:
                    suggestions.append(best_word)
        if suggestions:
            did_you_mean = " ".join(suggestions)

    return build_response(results, total, page, page_size, did_you_mean)

# ═══════════════════════════════════════════════════════════
# 目录树（过滤面板用）
# ═══════════════════════════════════════════════════════════

def build_filter_folder_tree(repo: str) -> list[dict]:
    if repo in folder_tree_data:
        return folder_tree_data[repo]
    if repo not in folder_index:
        return []

    repo_name_short = repo.split("/")[-1]
    tree_root = {
        "name": repo_name_short,
        "path": "",
        "children": [],
        "count": 0,
        "hasDirectFiles": False,
        "hasChildren": False,
        "showSelfToggle": False,
        "isRoot": True,
    }
    node_map = {"": tree_root}

    for path, idx_list in folder_index[repo].items():
        if path == "":
            tree_root["count"] = len(idx_list)
            continue
        parts = path.split("/")
        for depth in range(1, len(parts) + 1):
            key = "/".join(parts[:depth])
            if key not in node_map:
                parent_key = "/".join(parts[:depth - 1]) if depth > 1 else ""
                parent = node_map[parent_key]
                node = {
                    "name": parts[depth - 1],
                    "path": key,
                    "children": [],
                    "count": 0,
                    "hasDirectFiles": False,
                    "hasChildren": False,
                    "showSelfToggle": False,
                }
                node_map[key] = node
                parent["children"].append(node)
        node_map[path]["count"] = len(idx_list)

    direct_file_dirs = set()
    for idx in repo_records_map.get(repo, []):
        rec = records[idx]
        folders = rec.get("Folder", []) or []
        direct_file_dirs.add("/".join(folders))

    for path_key, node in node_map.items():
        node["hasDirectFiles"] = path_key in direct_file_dirs
        node["children"].sort(key=lambda child: child["name"])
        node["hasChildren"] = len(node["children"]) > 0
        node["showSelfToggle"] = bool(path_key and node["hasDirectFiles"] and node["hasChildren"])

    return [tree_root]

# ═══════════════════════════════════════════════════════════
# 文件夹内容浏览（进入式）
# ═══════════════════════════════════════════════════════════

def get_folder_contents(repo: str, path: str) -> dict:
    if repo in folder_browser_data:
        repo_browser = folder_browser_data.get(repo, {}) or {}
        if path in repo_browser:
            return enrich_browser_entry(repo, path, repo_browser[path])
        if path == "" and "" in repo_browser:
            return enrich_browser_entry(repo, path, repo_browser[""])

    path_parts = [p for p in path.strip("/").split("/") if p]
    path_depth = len(path_parts)
    repo_indices = repo_records_map.get(repo, [])
    repo_records = [records[i] for i in repo_indices]
    def is_under(rec_folders, prefix):
        return rec_folders[:len(prefix)] == prefix
    matching = [rec for rec in repo_records if is_under(rec.get("Folder", []) or [], path_parts)]

    subfolder_names = set()
    subfolder_counts = {}
    for rec in matching:
        folders = rec.get("Folder", []) or []
        if len(folders) > path_depth:
            name = folders[path_depth]
            subfolder_names.add(name)
            subfolder_counts[name] = subfolder_counts.get(name, 0) + 1
    result_folders = []
    for fname in sorted(subfolder_names):
        full_path = "/".join(path_parts + [fname]) if path_parts else fname
        result_folders.append({"name": fname, "path": full_path, "count": subfolder_counts.get(fname, 0)})

    direct_files = []
    for rec in matching:
        folders = rec.get("Folder", []) or []
        if len(folders) == path_depth:
            direct_files.append({
                "name": rec.get("File", ""),
                "ext": rec.get("Extension", ""),
                "link": rec.get("Link") or build_record_link(rec),
                "path": rec.get("Path") or build_record_path_url(rec),
                "hasTxt": rec.get("HasTxt", False),
                "size": rec.get("Size", ""),
            })
    direct_files.sort(key=lambda f: f["name"])
    return {"folders": result_folders, "files": direct_files, "current_path": path}

# ═══════════════════════════════════════════════════════════
# FastAPI 应用
# ═══════════════════════════════════════════════════════════

@asynccontextmanager
async def lifespan(app: FastAPI):
    load_data()
    yield

app = FastAPI(title="VoiceOfML Search", version="2.0", lifespan=lifespan)

app.add_middleware(GZipMiddleware, minimum_size=500)

app.add_middleware(
    CORSMiddleware,
    allow_origins=["*"],
    allow_methods=["*"],
    allow_headers=["*"],
)

# ── 搜索 API ──────────────────────────────────────────

class SearchRequest(BaseModel):
    q: str = ""
    repos: Optional[list[str]] = None
    extensions: Optional[list[str]] = None
    folders: Optional[list[str]] = None
    folder_match_mode: str = "prefix"
    folder_selfs: Optional[list[str]] = None
    folder_subtrees: Optional[list[str]] = None
    min_size: Optional[int] = None
    max_size: Optional[int] = None
    page: int = 1
    page_size: int = 100
    search_folders: bool = True
    sort: str = "relevance"
    exact: bool = False

@app.post("/api/search")
def api_search(body: SearchRequest):
    return JSONResponse(search(
        body.q,
        [r.strip() for r in body.repos] if body.repos else None,
        [e.strip().lower() for e in body.extensions] if body.extensions else None,
        [f.strip() for f in body.folders] if body.folders else None,
        body.min_size, body.max_size,
        body.page, body.page_size,
        body.search_folders, body.sort, body.exact,
        body.folder_match_mode,
        [f.strip() for f in body.folder_selfs] if body.folder_selfs else None,
        [f.strip() for f in body.folder_subtrees] if body.folder_subtrees else None,
    ))

@app.post("/api/search/{repo_name}")
def api_search_repo(repo_name: str, body: SearchRequest):
    full_repo = f"VoiceOfML/{repo_name}"
    if full_repo not in repo_counts:
        return JSONResponse({"error": f"仓库 {full_repo} 不存在", "results": [], "total": 0}, status_code=404)
    return JSONResponse(search(
        body.q, [full_repo],
        [e.strip().lower() for e in body.extensions] if body.extensions else None,
        [f.strip() for f in body.folders] if body.folders else None,
        body.min_size, body.max_size,
        body.page, body.page_size,
        body.search_folders, body.sort, body.exact,
        body.folder_match_mode,
        [f.strip() for f in body.folder_selfs] if body.folder_selfs else None,
        [f.strip() for f in body.folder_subtrees] if body.folder_subtrees else None,
    ))

@app.get("/api/repos")
def api_repos():
    return JSONResponse(repo_list)

@app.get("/api/extensions")
def api_extensions(repo: Optional[str] = Query(default=None)):
    if repo:
        full_repo = f"VoiceOfML/{repo}"
        repo_indices = repo_records_map.get(full_repo, [])
        pool = [records[i] for i in repo_indices]
    else:
        pool = records

    ext_count = {}
    for rec in pool:
        ext = (rec.get("Extension") or "").lower()
        if ext:
            ext_count[ext] = ext_count.get(ext, 0) + 1

    # 保持原有排序（按字母）
    sorted_exts = sorted(ext_count.keys())
    return JSONResponse([{"name": e, "count": ext_count[e]} for e in sorted_exts])
    
@app.get("/api/folders/{repo_name}")
def api_folders(repo_name: str):
    full_repo = f"VoiceOfML/{repo_name}"
    if full_repo not in repo_counts:
        return JSONResponse({"error": f"仓库 {full_repo} 不存在"}, status_code=404)
    return JSONResponse(build_filter_folder_tree(full_repo))

@app.get("/api/folders/{repo_name}/contents")
def api_folder_contents(repo_name: str, path: str = Query(default="")):
    full_repo = f"VoiceOfML/{repo_name}"
    if full_repo not in repo_counts:
        return JSONResponse({"error": f"仓库 {full_repo} 不存在"}, status_code=404)
    return JSONResponse(get_folder_contents(full_repo, path))

@app.get("/api/random")
def api_random(repo: Optional[str] = Query(default=None)):
    if repo:
        full_repo = f"VoiceOfML/{repo}"
        repo_indices = repo_records_map.get(full_repo, [])
        pool = [records[i] for i in repo_indices]
    else:
        pool = records
    if not pool:
        return JSONResponse({"error": "无可用记录"}, status_code=404)
    rec = dict(random.choice(pool))
    rec["Link"] = rec.get("Link") or build_record_link(rec)
    rec["Path"] = rec.get("Path") or build_record_path_url(rec)
    return JSONResponse(rec)

class ZipRequest(BaseModel):
    files: list[dict] = []

@app.post("/api/zip")
async def api_zip(req: ZipRequest):
    if not req.files:
        return JSONResponse({"error": "无文件"}, status_code=400)
    if len(req.files) > 200:
        return JSONResponse({"error": "最多 200 个文件"}, status_code=400)

    timeout = aiohttp.ClientTimeout(total=300, connect=15)

    async def download_one(session, f):
        link = f.get("link", "")
        name = f.get("name", "unknown")
        parsed = urlparse(link)
        if parsed.hostname not in ("huggingface.co", "hf.co", "cdn-lfs.hf.co", "cdn-lfs-us-1.hf.co"):
            return None
        try:
            async with session.get(
                link, allow_redirects=True,
                headers={"User-Agent": "VoiceOfML-Search/1.0"},
            ) as upstream:
                if upstream.status != 200:
                    return None
                return (name, await upstream.read())
        except Exception as e:
            print(f"ZIP download failed {name}: {e}")
            return None

    async with aiohttp.ClientSession(timeout=timeout) as session:
        tasks = [download_one(session, f) for f in req.files]
        results = await asyncio.gather(*tasks)

    valid = [r for r in results if r is not None]
    if not valid:
        return JSONResponse({"error": "所有文件下载失败"}, status_code=500)

    try:
        from stream_zip import stream_zip, ZIP_32
        def member_files():
            modified_at = datetime.now()
            for name, data in valid:
                yield (name.encode("utf-8"), modified_at, ZIP_32, data)
        return StreamingResponse(
            stream_zip(member_files()),
            media_type="application/zip",
            headers={"Content-Disposition": "attachment; filename=voml_batch.zip"},
        )
    except Exception as e:
        print(f"stream-zip failed, falling back to zipfile: {e}")
        import io, zipfile
        buf = io.BytesIO()
        with zipfile.ZipFile(buf, "w", zipfile.ZIP_DEFLATED) as zf:
            for name, data in valid:
                zf.writestr(name, data)
        buf.seek(0)
        def chunks():
            while True:
                c = buf.read(65536)
                if not c: break
                yield c
        return StreamingResponse(
            chunks(),
            media_type="application/zip",
            headers={"Content-Disposition": "attachment; filename=voml_batch.zip"},
        )

# ── 下载代理 ──────────────────────────────────────────
@app.get("/api/download")
async def api_download(file: str = Query(...), link: str = Query(...)):
    target_url = unquote(link)
    from urllib.parse import urlparse
    parsed = urlparse(target_url)
    if parsed.hostname not in ("huggingface.co", "hf.co", "cdn-lfs.hf.co", "cdn-lfs-us-1.hf.co"):
        return JSONResponse({"error": "只允许 HuggingFace 域名"}, status_code=403)
    encoded_filename = quote(file, safe="")
    content_disposition = f"attachment; filename*=UTF-8''{encoded_filename}"

    async def stream_response():
        timeout = aiohttp.ClientTimeout(total=120, connect=15)
        async with aiohttp.ClientSession(timeout=timeout) as session:
            try:
                async with session.get(
                    target_url,
                    allow_redirects=True,
                    headers={"User-Agent": "VoiceOfML-Search/1.0"},
                ) as upstream:
                    if upstream.status != 200:
                        return
                    async for chunk in upstream.content.iter_chunked(65536):
                        yield chunk
            except Exception as e:
                print(f"下载代理异常: {e}")
                return

    return StreamingResponse(
        stream_response(),
        media_type="application/octet-stream",
        headers={
            "Content-Disposition": content_disposition,
            "Cache-Control": "no-cache",
        },
    )

# ── 静态文件 ──────────────────────────────────────────
txt_dir = Path("txt")
if not txt_dir.exists():
    txt_dir.mkdir(parents=True, exist_ok=True)
app.mount("/txt", StaticFiles(directory="txt"), name="txt")
app.mount("/static", StaticFiles(directory="static", html=True), name="static")
app.mount("/icons", StaticFiles(directory="static/icons"), name="icons")

# ── PWA files at root ─────────────────────────────────
@app.get("/manifest.json")
def serve_manifest():
    p = Path("static/manifest.json")
    if p.exists():
        return JSONResponse(json.loads(p.read_text(encoding="utf-8")))
    return JSONResponse({"error": "manifest not found"}, status_code=404)

@app.get("/sw.js")
def serve_sw():
    p = Path("static/sw.js")
    if p.exists():
        return PlainTextResponse(p.read_text(encoding="utf-8"), media_type="application/javascript")
    return PlainTextResponse("", status_code=404)

@app.get("/{rest_of_path:path}")
async def serve_spa(rest_of_path: str):
    index_path = Path("static/index.html")
    if index_path.exists():
        return HTMLResponse(index_path.read_text(encoding="utf-8"))
    return HTMLResponse("<h1>VoiceOfML Search</h1>", status_code=200)