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, edit_file
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("/library/upload")
def upload(file: UploadFile = File(...)):
    url = upload_rustfs(file.file, file.filename)
    return ocr_from_url(url)


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


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


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


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


@app.put("/library/files/{id}", response_model=bool)
async def edit_library_file(id: int, data: LibraryFileRequest, db: AsyncSession = Depends(get_db)):
    return await edit_file(db, id, data)