import os
import shutil
import tempfile
from tempfile import NamedTemporaryFile
from pathlib import Path

from fastapi import FastAPI, UploadFile
from fastapi.middleware import Middleware
from fastapi.middleware.cors import CORSMiddleware
from fastapi.responses import StreamingResponse
from .rag import ChatPDF

middleware = [
    Middleware(
        CORSMiddleware,
        allow_origins=["*"],
        allow_methods=['*'],
        allow_headers=['*']
    )
]

app = FastAPI(middleware=middleware)

files_dir = os.path.expanduser("~/wtp_be_files/")
session_assistant = ChatPDF()
# session_messages = []

# async def stream_generator()

@app.get("/query")
async def process_input(text: str):
    if text and len(text.strip()) > 0:
        text = text.strip()
        print("PRINTING STREAM")
        agent_text_stream = session_assistant.ask(text)
        print(agent_text_stream)
        for text in agent_text_stream:
            print(text)
            # yield text
        # session_messages.append((text, True))
        # session_messages.append((agent_text, False))
        return StreamingResponse(agent_text_stream, media_type='text/event-stream')


@app.post("/upload")
def upload(files: list[UploadFile]):
    session_assistant.clear()
    # session_messages = []
    
    try:
        os.makedirs(files_dir)
        for file in files:
            try:
                path = f"{files_dir}/{file.filename}"
                file.file.seek(0)
                with open(path, 'wb') as destination:
                    shutil.copyfileobj(file.file, destination)
            finally:
                file.file.close()
    finally:
        session_assistant.ingest(files_dir)
        shutil.rmtree(files_dir)

    return "Files inserted!"


@app.get("/")
def ping():
    return "Pong!"