from pydantic import BaseModel
from fastapi import FastAPI, UploadFile, File
from app.pdf_handler import process_pdf, process_image_pdf, process_image_pdf_dummy
from app.rag_pipeline import generate_answer, generate_answer_am
from app.test import query_test
from fastapi import FastAPI, File, UploadFile
from PIL import Image
import pytesseract
import io

app = FastAPI()


class ChatRequest(BaseModel):
    query: str
    top_k: int = 3


@app.post("/chat/")
async def chat(request: ChatRequest):
    return generate_answer(request.query, request.top_k)


@app.post("/am/chat/")
async def chat(request: ChatRequest):
    return generate_answer_am(request.query, request.top_k)



@app.post("/upload/")
async def upload_pdf(file: UploadFile = File(...)):
    return process_pdf(file)


@app.post("/upload/imgpdf/")
async def upload_pdf_image(file: UploadFile = File(...)):
    return process_image_pdf(file)


@app.post("/upload/imgpdf/dummy/")
async def upload_pdf_image_dummy(f):
    return process_image_pdf_dummy()


@app.post("/test/")
async def chat(query: str, top_k: int = 3):
    return query_test(query)


@app.post("/ocr/")
async def extract_text(file: UploadFile = File(...)):
    image_data = await file.read()
    image = Image.open(io.BytesIO(image_data))
    text = pytesseract.image_to_string(image)
    return {"extracted_text": text}