server.py

import re
import uuid
import hashlib
import asyncio
import shutil
from datetime import datetime
from typing import Optional
from dataclasses import dataclass, field
from urllib.parse import quote

from fastapi import FastAPI, HTTPException, BackgroundTasks
from fastapi.responses import FileResponse, JSONResponse, Response, StreamingResponse
from fastapi.middleware.cors import CORSMiddleware
from pydantic import BaseModel
from pathlib import Path
from contextlib import asynccontextmanager
import json
import os

from bili_download import download_video
from awasr import aw_asw
from llm_inference import llm_inference
from format_note import format_note


# 获取当前文件所在目录（Backend/）
BASE_DIR = Path(__file__).parent
CONFIG_PATH = BASE_DIR / "config.json"
STATIC_DIR = BASE_DIR / "static"
PROMPT_DIR = BASE_DIR / "prompt"
TRASH_DIR = BASE_DIR / ".trash"

# 全局配置变量
config = None
aw_asr_api_key = None
llm_api_key = None
base_url = None
model_name = None
output_path = None


@dataclass
class TaskInfo:
    """任务信息类，包含任务队列和状态"""
    task_id: str
    urls: list[str]
    queue: asyncio.Queue = field(default_factory=asyncio.Queue)
    status: str = "pending"  # pending, running, completed, failed
    results: list[dict] = field(default_factory=list)
    total: int = 0
    success_count: int = 0
    failed_count: int = 0


# 全局任务字典
tasks_db: dict[str, TaskInfo] = {}

# note_id -> 文件路径 的缓存映射
note_path_map: dict = {}


class TaskRequest(BaseModel):
    urls: list[str]


@asynccontextmanager
async def lifespan(app: FastAPI):
    """启动事件：检测 ffmpeg，创建必要目录，加载配置"""
    global config, aw_asr_api_key, llm_api_key, base_url, model_name, output_path

    # 检测 ffmpeg
    if not shutil.which("ffmpeg"):
        print("警告：未检测到 ffmpeg，视频转音频功能可能无法正常工作")

    # 创建必要目录
    for folder in ["asr_result", "notebook_output", "wav_cache"]:
        folder_path = BASE_DIR / folder
        if not folder_path.exists():
            os.makedirs(folder_path)
            print(f"已创建目录: {folder_path}")

    # 创建 .trash 目录
    if not TRASH_DIR.exists():
        os.makedirs(TRASH_DIR)

    # 加载 config.json
    try:
        with open(CONFIG_PATH, "r") as f:
            config = json.load(f)
            aw_asr_api_key = config[0]["aw_asr"]["api_key"]
            llm_api_key = config[1]['llm']['api_key']
            base_url = config[1]['llm']['base_url']
            model_name = config[1]['llm']['model_name']
            output_path = config[2]['output_path']["obsidian_path"]
        print("配置加载成功")
    except Exception as e:
        print(f'程序运行异常，请检查 config.json 文件: {e}')
        raise

    yield


# 创建 FastAPI 应用
app = FastAPI(lifespan=lifespan)

# 添加 CORS 中间件（允许 localhost）
app.add_middleware(
    CORSMiddleware,
    allow_origins=["http://localhost", "http://127.0.0.1"],
    allow_credentials=True,
    allow_methods=["*"],
    allow_headers=["*"],
)


# ============================================================
# 辅助函数
# ============================================================

def validate_url(url: str) -> bool:
    """校验 URL 是否为 B站链接或 BV 号"""
    url = url.strip()
    # 纯 BV 号: BV + 10位字母数字
    if re.match(r'^BV[a-zA-Z0-9]{10}$', url):
        return True
    # B站链接（含 BV 号）
    if re.match(r'^https?://(www\.)?bilibili\.com/video/BV[a-zA-Z0-9]{10}', url):
        return True
    return False


def normalize_url(url: str) -> str:
    """将 BV 号标准化为完整 URL"""
    url = url.strip()
    if url[:2] == "BV":
        return "https://www.bilibili.com/video/" + url
    return url


def get_note_id(filename: str) -> str:
    """根据文件名生成 note_id（md5 前12位）"""
    return hashlib.md5(filename.encode()).hexdigest()[:12]


def read_prompt_file() -> str:
    """直接读取 prompt/ 目录下唯一的 .txt 文件（绕过 select_prompt_file）"""
    prompt_files = [f for f in os.listdir(PROMPT_DIR) if f.endswith('.txt')]
    if not prompt_files:
        raise FileNotFoundError("prompt/ 目录下未找到 .txt 提示词文件")
    with open(PROMPT_DIR / prompt_files[0], 'r', encoding='utf-8') as f:
        return f.read()


def get_output_dir() -> Path:
    """获取有效的输出目录"""
    obsidian = Path(output_path) if output_path else None
    if obsidian and obsidian.exists():
        return obsidian
    return BASE_DIR / "notebook_output"


def scan_notes():
    """扫描输出目录，构建 note_path_map"""
    global note_path_map
    note_path_map.clear()
    out_dir = get_output_dir()
    if not out_dir.exists():
        return
    pattern = "【智能笔记】*.md"
    for filepath in out_dir.glob(pattern):
        note_id = get_note_id(filepath.name)
        note_path_map[note_id] = filepath


def format_sse_event(event_type: str, data: dict) -> str:
    """格式化 SSE 事件"""
    return f"event: {event_type}\ndata: {json.dumps(data, ensure_ascii=False)}\n\n"


# ============================================================
# 后台处理流程
# ============================================================

async def process_single_url(url: str, queue: asyncio.Queue) -> dict:
    """处理单个 URL: download -> asr -> llm -> format -> save"""
    result_info = {"url": url, "status": "pending", "error": None, "filename": None}

    try:
        url = normalize_url(url)

        # 推送下载阶段事件
        await queue.put({
            "type": "stage",
            "stage": "downloading",
            "url": url,
            "message": "正在下载音频..."
        })

        # 1. 下载视频音频
        download_file_name = await asyncio.to_thread(
            download_video, url, str(BASE_DIR / "wav_cache"), 32
        )

        # 推送 ASR 阶段事件
        await queue.put({
            "type": "stage",
            "stage": "asr",
            "message": "正在语音转文字..."
        })

        # 2. ASR 语音识别
        await asyncio.to_thread(
            aw_asw, download_file_name, aw_asr_api_key
        )

        # 推送 LLM 阶段事件
        await queue.put({
            "type": "stage",
            "stage": "llm",
            "message": "正在生成笔记..."
        })

        # 3. 读取 ASR 结果
        asr_file = BASE_DIR / "asr_result" / f"{download_file_name.split('.')[0]}.txt"
        text = await asyncio.to_thread(
            lambda: asr_file.read_text(encoding="utf-8")
        )

        # 4. LLM 推理 - 直接读取 prompt 文件，绕过 select_prompt_file
        # llm_inference 内部调用 select_prompt_file()，它会在只有一个文件时自动选择
        # 但它使用了 input() 来选择，所以我们需要用直接方式
        # 这里我们直接内联 LLM 调用逻辑
        from openai import OpenAI
        system_prompt = read_prompt_file()

        def _call_llm():
            client = OpenAI(api_key=llm_api_key, base_url=base_url)
            response = client.chat.completions.create(
                model=model_name or "deepseek-chat",
                messages=[
                    {"role": "system", "content": system_prompt},
                    {"role": "user", "content": text},
                ],
                stream=False,
            )
            return response.choices[0].message.content

        llm_result = await asyncio.to_thread(_call_llm)

        # 推送格式化阶段事件
        await queue.put({
            "type": "stage",
            "stage": "formatting",
            "message": "正在格式化..."
        })

        # 5. 格式化笔记
        formatted = await asyncio.to_thread(format_note, llm_result)

        # 6. 保存文件
        base_name = download_file_name.split('.')[0]
        note_filename = f"【智能笔记】{base_name}.md"
        content = formatted + f"\n\n---\n视频链接：[{download_file_name}]({url})"

        out_dir = get_output_dir()
        save_path = out_dir / note_filename

        def _save_file():
            try:
                save_path.write_text(content, encoding="utf-8")
            except Exception:
                fallback_dir = BASE_DIR / "notebook_output"
                fallback_dir.mkdir(exist_ok=True)
                (fallback_dir / note_filename).write_text(content, encoding="utf-8")

        await asyncio.to_thread(_save_file)

        result_info["status"] = "done"
        result_info["filename"] = note_filename

        # 推送完成事件
        await queue.put({
            "type": "complete",
            "url": url,
            "title": download_file_name,
            "status": "success"
        })

    except Exception as e:
        result_info["status"] = "error"
        result_info["error"] = str(e)

        # 推送错误事件
        await queue.put({
            "type": "error",
            "url": url,
            "error": f"处理失败: {str(e)}"
        })

    return result_info


async def process_task(task_id: str, urls: list[str]):
    """后台任务：串行处理每个 URL，推送进度事件到队列"""
    if task_id not in tasks_db:
        return

    task = tasks_db[task_id]
    task.status = "running"

    for url in urls:
        url = url.strip()
        if not url:
            continue
        result = await process_single_url(url, task.queue)
        task.results.append(result)

        if result["status"] == "done":
            task.success_count += 1
        else:
            task.failed_count += 1

    # 判断整体状态
    has_error = any(r["status"] == "error" for r in task.results)
    has_done = any(r["status"] == "done" for r in task.results)
    if has_done and has_error:
        task.status = "partial"
    elif has_done:
        task.status = "done"
    else:
        task.status = "error"

    # 推送完成事件
    await task.queue.put({
        "type": "done",
        "total": task.total,
        "success": task.success_count,
        "failed": task.failed_count
    })

    # 重新扫描笔记
    scan_notes()


# ============================================================
# API 端点
# ============================================================

@app.get("/")
async def index():
    """GET / 端点，返回 static/index.html"""
    index_html = STATIC_DIR / "index.html"
    if not index_html.exists():
        return {"error": "index.html not found"}
    return FileResponse(index_html)


@app.post("/api/tasks")
async def create_task(req: TaskRequest):
    """POST /api/tasks - 创建处理任务"""
    # 校验 urls 非空
    if not req.urls:
        raise HTTPException(status_code=400, detail="urls 不能为空")

    # 校验每个 URL 格式
    for url in req.urls:
        if not validate_url(url.strip()):
            raise HTTPException(
                status_code=400,
                detail=f"无效的 URL 格式: {url}"
            )

    task_id = uuid.uuid4().hex
    task_info = TaskInfo(
        task_id=task_id,
        urls=req.urls,
        total=len(req.urls)
    )
    tasks_db[task_id] = task_info

    # 启动后台处理
    asyncio.create_task(process_task(task_id, req.urls))

    return JSONResponse(
        status_code=202,
        content={"task_id": task_id, "total": len(req.urls)}
    )


@app.get("/api/tasks/{task_id}/stream")
async def task_stream(task_id: str):
    """GET /api/tasks/{task_id}/stream - SSE 流端点，返回任务进度"""
    if task_id not in tasks_db:
        raise HTTPException(status_code=404, detail="任务不存在")

    task_info = tasks_db[task_id]

    async def event_generator():
        """SSE 事件生成器"""
        try:
            while True:
                # 从队列中获取事件
                event = await task_info.queue.get()

                # 检查是否为结束事件
                if event.get("type") == "done":
                    yield format_sse_event("done", event)
                    break

                # 发送其他事件
                event_type = event.get("type")
                yield format_sse_event(event_type, event)

        except asyncio.CancelledError:
            # 客户端断开连接
            pass

    return StreamingResponse(
        event_generator(),
        media_type="text/event-stream",
        headers={
            "Cache-Control": "no-cache",
            "Connection": "keep-alive",
            "X-Accel-Buffering": "no"
        }
    )


@app.get("/api/tasks/{task_id}")
async def get_task_status(task_id: str):
    """GET /api/tasks/{task_id} - 查询任务状态"""
    if task_id not in tasks_db:
        raise HTTPException(status_code=404, detail="任务不存在")
    return tasks_db[task_id]


@app.get("/api/notes")
async def list_notes():
    """GET /api/notes - 返回笔记列表"""
    scan_notes()
    notes = []
    for note_id, filepath in note_path_map.items():
        stat = filepath.stat()
        notes.append({
            "id": note_id,
            "title": filepath.stem,
            "created_at": datetime.fromtimestamp(stat.st_mtime).isoformat(),
            "filename": filepath.name,
        })
    # 按修改时间倒序
    notes.sort(key=lambda x: x["created_at"], reverse=True)
    return {"notes": notes}


@app.get("/api/notes/{note_id}")
async def get_note(note_id: str):
    """GET /api/notes/{note_id} - 返回笔记内容"""
    scan_notes()
    if note_id not in note_path_map:
        raise HTTPException(status_code=404, detail="笔记不存在")

    filepath = note_path_map[note_id]
    stat = filepath.stat()
    content = filepath.read_text(encoding="utf-8")

    return {
        "id": note_id,
        "title": filepath.stem,
        "content": content,
        "created_at": datetime.fromtimestamp(stat.st_mtime).isoformat(),
    }


@app.get("/api/notes/{note_id}/download")
async def download_note(note_id: str):
    """GET /api/notes/{note_id}/download - 下载笔记文件"""
    scan_notes()
    if note_id not in note_path_map:
        raise HTTPException(status_code=404, detail="笔记不存在")

    filepath = note_path_map[note_id]
    encoded_name = quote(filepath.name)
    return FileResponse(
        path=str(filepath),
        filename=filepath.name,
        media_type="text/markdown",
        headers={
            "Content-Disposition": f"attachment; filename*=UTF-8''{encoded_name}"
        }
    )


@app.delete("/api/notes/{note_id}")
async def delete_note(note_id: str):
    """DELETE /api/notes/{note_id} - 删除笔记（移至 .trash/）"""
    scan_notes()
    if note_id not in note_path_map:
        raise HTTPException(status_code=404, detail="笔记不存在")

    filepath = note_path_map[note_id]
    dest = TRASH_DIR / filepath.name

    shutil.move(str(filepath), str(dest))

    # 从缓存中移除
    del note_path_map[note_id]

    return Response(status_code=204)


if __name__ == "__main__":
    import uvicorn
    uvicorn.run(app, host="127.0.0.1", port=8000)