from fastapi import FastAPI, HTTPException
from fastapi.responses import StreamingResponse, HTMLResponse
import json
from es_gpt import ESGPT
import asyncio

# Create an instance of the ESGPT class
es = ESGPT(index_name="papers")

# Create a FastAPI app
app = FastAPI()

# Define the search route


@app.get("/", response_class=HTMLResponse)
async def read_index():
    with open("index.html", "r") as file:
        html = file.read()
    return html


@app.get("/search")
async def search(q: str):
    # Perform a search for the query
    results = es.search(q)
    print(results)

    # Stream the search results to the client
    async def stream_response():
        for hit in results:
            # sleep(0.1)
            await asyncio.sleep(0.1)
            yield "data: " + json.dumps(hit) + "\n\n"
        yield "[DONE]"

    return StreamingResponse(stream_response(), media_type="text/event-stream")

# Define the summary route


@app.get("/summary")
async def summary(q: str):
    # Perform a search for the query
    results = es.search(q)

    # Generate summaries of the search results
    resp = es.summarize(q, results)

    if resp.status_code != 200:
        raise HTTPException(resp.status_code, resp.text)

    return StreamingResponse(resp.iter_content(1),
                             media_type="text/event-stream")