import asyncio
import json
from typing import Annotated, Optional
from fastapi import APIRouter, HTTPException, Depends, Query
from fastapi.responses import StreamingResponse
from sqlalchemy.ext.asyncio import AsyncSession
from sqlalchemy import select

from app.database import get_db
from app.models.transcript import Transcript
from app.state import task_manager
from app.services.extractor import run_extraction

router = APIRouter(prefix="/extraction", tags=["故事提取"])


@router.post("/start")
async def start_extraction(
    task_id: Optional[str] = Query(None),
    db: Annotated[AsyncSession, Depends(get_db)] = None,
):
    """启动提取任务，支持多用户并发"""
    # 如果传了 task_id，复用已有任务
    if task_id:
        state = task_manager.get_task(task_id)
        if state and state.is_running:
            return {"task_id": task_id, "message": "任务正在进行中"}
    else:
        state = None

    # 检查是否有待处理的文字稿
    result = await db.execute(select(Transcript).where(Transcript.status == "pending"))
    pending = result.scalars().all()

    if not pending:
        # 没有 pending 的，把所有非 pending 的文件重置为 pending（包括 completed/processing/error）
        result2 = await db.execute(
            select(Transcript).where(Transcript.status != "pending")
        )
        others = result2.scalars().all()
        if others:
            for t in others:
                t.status = "pending"
                t.error_message = ""
            await db.commit()
            pending = others
        else:
            raise HTTPException(400, "没有可提取的文字稿，请先上传文件")

    total = len(pending)

    # 创建新任务
    state = task_manager.create_task()
    state.reset()
    state.is_running = True

    asyncio.create_task(run_extraction(state))

    return {"task_id": state.task_id, "total_files": total}


@router.get("/status")
async def get_status(task_id: Optional[str] = Query(None)):
    """获取提取进度"""
    if task_id:
        state = task_manager.get_task(task_id)
        if state:
            return {**state.progress, "task_id": state.task_id}
    return {"status": "idle", "percent": 0}


@router.get("/status/stream")
async def stream_status(task_id: Optional[str] = Query(None)):
    """SSE 实时进度推送（支持多用户）"""
    if not task_id:
        raise HTTPException(400, "缺少 task_id 参数")

    state = task_manager.get_task(task_id)
    if not state:
        raise HTTPException(404, "任务不存在")

    async def event_generator():
        # 立即推送当前状态
        yield f"event: progress\ndata: {json.dumps({**state.progress, 'task_id': state.task_id}, ensure_ascii=False)}\n\n"
        if state.progress.get("status") in ("completed", "error", "stopped"):
            return
        while True:
            data = await state.get_update()
            yield f"event: progress\ndata: {json.dumps({**data, 'task_id': state.task_id}, ensure_ascii=False)}\n\n"
            if data.get("status") in ("completed", "error", "stopped"):
                break

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


@router.post("/stop")
async def stop_extraction(task_id: Optional[str] = Query(None)):
    """停止提取任务"""
    if task_id:
        state = task_manager.get_task(task_id)
    else:
        state = None

    if not state or not state.is_running:
        raise HTTPException(400, "没有正在进行的提取任务")

    state.cancel_event.set()
    return {"message": "正在停止...", "task_id": state.task_id}