from typing import List

from fastapi import FastAPI, UploadFile, File, Depends
from sqlalchemy.ext.asyncio import AsyncSession
from starlette.middleware.cors import CORSMiddleware
from starlette.responses import StreamingResponse

from database import get_db
from schemas.agent import QueryRequest
from schemas.library import LibraryFileRequest, LibraryFileResponse
from service.agent_service import query_agent
from service.library_service import add_file, get_file, list_files
from storage import upload_rustfs
from ocr import ocr_from_url

app = FastAPI(title="AI Study Service")

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


@app.post("/upload")
def upload(file: UploadFile = File(...)):
    url = upload_rustfs(file.file, file.filename)
    text = ocr_from_url(url)

    return {
        "text": text,
    }


@app.post("/generate")
def generate(query: QueryRequest):
    """流式对话"""
    return StreamingResponse(
        query_agent(query),
        media_type="text/event-stream"
    )


@app.post("/files", response_model=bool)
async def add_library_file(data: LibraryFileRequest, db: AsyncSession = Depends(get_db)):
    await add_file(db, data)
    return True


@app.get("/files", response_model=List[LibraryFileResponse])
async def list_library_file(db: AsyncSession = Depends(get_db)):
    return await list_files(db)


@app.get("/files/{file_id}", response_model=LibraryFileResponse)
async def get_library_file(file_id: int, db: AsyncSession = Depends(get_db)):
    return await get_file(db, file_id)