main.py

import io
import os
import uuid
from contextlib import asynccontextmanager
from pathlib import Path
from typing import Optional

from fastapi import FastAPI, UploadFile, File, HTTPException, status, Depends, BackgroundTasks, Form
from fastapi.responses import FileResponse
from fastapi.security import HTTPBearer, HTTPAuthorizationCredentials
from pydantic import BaseModel
from sqlalchemy.orm import Session

from _markitdown import MarkItDown
from base import DocumentConverterResult
from model_manager import ModelConfigurator
from repository.db import get_db, Job

# 安全验证
security = HTTPBearer()

# 从环境变量获取API密钥
API_KEY = os.getenv("MARKIT_API_KEY", "secret-key")
OUTPUT_DIR = Path("output_files")
OUTPUT_DIR.mkdir(exist_ok=True)
MINER_RUNNING_DEVICE = os.getenv("MINER_RUNNING_DEVICE", "cpu")
port = int(os.getenv("PORT", 20926))


# 依赖项：API Key 验证
async def verify_api_key(
        credentials: HTTPAuthorizationCredentials = Depends(security)
):
    if credentials.scheme != "Bearer" or credentials.credentials != API_KEY:
        raise HTTPException(
            status_code=status.HTTP_401_UNAUTHORIZED,
            detail="Invalid API Key",
        )
    return credentials


@asynccontextmanager
async def lifespan(app: FastAPI):
    """服务启动和关闭时的生命周期管理"""
    try:
        # 初始化模型
        configurator = ModelConfigurator(
            device=os.getenv("MINERU_DEVICE", MINER_RUNNING_DEVICE),
            use_modelscope=True
        )
        configurator.setup_environment()
        print("模型初始化完成")
    except Exception as e:
        print(f"模型初始化失败: {str(e)}")
        raise

    yield  # 应用运行期间

    # 清理逻辑（可选）
    print("服务关闭，清理资源...")

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


# from slowapi import Limiter, _rate_limit_exceeded_handler
# from slowapi.errors import RateLimitExceeded
# from slowapi.util import get_remote_address
# limiter = Limiter(key_func=get_remote_address)
# app.state.limiter = limiter
# app.add_exception_handler(RateLimitExceeded, _rate_limit_exceeded_handler)
# @limiter.limit("100/minute")


# 数据模型
class JobStatusResponse(BaseModel):
    job_id: str
    status: str
    filename: str
    params: dict
    error: Optional[str]


class JobResultResponse(BaseModel):
    job_id: str
    download_url: str
    format: str


def process_file(db: Session, job_id: str, file_content: bytes, filename: str, pdf_mode: str = "simple"):
    """处理各种文件的后台任务"""
    try:
        # 更新任务状态为 processing
        job = db.query(Job).filter(Job.id == job_id).first()
        if not job:
            raise ValueError(f"Job {job_id} not found")

        job.status = "processing"
        db.commit()

        # 创建处理器
        markitdown = MarkItDown(pdf_mode=pdf_mode)

        # 根据输入类型处理
        if filename.endswith('.md'):
            result = DocumentConverterResult(text_content=file_content.decode('utf-8'))
        else:
            # 将字节内容转为文件流
            file_stream = io.BytesIO(file_content)
            result = markitdown.convert_stream(file_stream)

        # 保存结果到文件
        output_file = OUTPUT_DIR / f"{job_id}.md"
        with open(output_file, "w", encoding="utf-8") as f:
            f.write(result.text_content)

        # 更新任务状态为 completed
        job.status = "completed"
        job.result_file = str(output_file)
        db.commit()

    except Exception as e:
        # 更新任务状态为 failed
        job.status = "failed"
        job.error = f"{type(e).__name__}: {str(e)}"
        db.commit()


@app.post("/api/jobs", status_code=status.HTTP_202_ACCEPTED)
async def upload_file(
        background_tasks: BackgroundTasks,
        file: UploadFile = File(...),
        pdf_mode: str = Form("simple"),
        db: Session = Depends(get_db)
):
    """上传文件并启动转换任务"""
    # 生成任务ID
    job_id = str(uuid.uuid4())

    try:
        # 读取文件内容
        content = await file.read()

        # 创建任务记录
        job = Job(
            id=job_id,
            filename=file.filename,
            params={"pdf_mode": pdf_mode},
            status="pending"
        )
        db.add(job)
        db.commit()

        # 启动后台任务
        background_tasks.add_task(
            process_file,
            db=db,
            job_id=job_id,
            file_content=content,
            filename=file.filename,
            pdf_mode=pdf_mode
        )

        return {"job_id": job_id}

    except Exception as e:
        raise HTTPException(
            status_code=status.HTTP_500_INTERNAL_SERVER_ERROR,
            detail=f"File upload failed: {str(e)}"
        )


@app.get("/api/jobs/{job_id}", response_model=JobStatusResponse)
async def get_job_status(
        job_id: str,
        db: Session = Depends(get_db)
):
    """查询任务状态"""
    job = db.query(Job).filter(Job.id == job_id).first()
    if not job:
        raise HTTPException(
            status_code=status.HTTP_404_NOT_FOUND,
            detail="Job not found"
        )

    return JobStatusResponse(
        job_id=job.id,
        status=job.status,
        filename=job.filename,
        params=job.params,
        error=job.error
    )


@app.get("/api/jobs/{job_id}/result")
async def download_result(
        job_id: str,
        db: Session = Depends(get_db)
):
    """下载任务结果文件"""
    job = db.query(Job).filter(Job.id == job_id).first()
    if not job:
        raise HTTPException(
            status_code=status.HTTP_404_NOT_FOUND,
            detail="Job not found"
        )

    if job.status != "completed":
        raise HTTPException(
            status_code=status.HTTP_425_TOO_EARLY,
            detail="Job not completed"
        )

    result_file = job.result_file
    if not result_file or not os.path.exists(result_file):
        raise HTTPException(
            status_code=status.HTTP_404_NOT_FOUND,
            detail="Result file not found"
        )

    # 返回文件内容
    return FileResponse(
        result_file,
        filename=f"{job.filename}.md",
        media_type="text/markdown"
    )


if __name__ == "__main__":
    import uvicorn

    uvicorn.run(app, host="0.0.0.0", port=port)