import uvicorn
import helper
from fastapi import FastAPI
from fastapi.middleware.cors import CORSMiddleware
from pydantic import BaseModel
from typing import List

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

class RequestData(BaseModel):
    abstract: str
    words: int
    papers: int

class ResponseData(BaseModel):
    summary: str
    ids: List[str]

base_prompt = """
Write a related work section for a scientific paper based on the following abstract and references. Your output should:

1. Introduce the topic (1 sentence)
2. Summarize key related works (3-4 sentences)
3. Compare and contrast approaches (2-3 sentences)
4. Connect to the proposed work (1-2 sentences)

Use [1], [2], etc. to cite references. Do not copy text directly. Be concise and coherent.

Abstract: {abstract}

References:
{references}

Write the related work section:
"""
sentence_plan = """
1. Introduction sentence (1 sentence)
2. Overview of relevant studies (2-3 sentences, cite [1], [2])
3. Detailed discussion on key papers (4-5 sentences, cite [3], [4], [5], if present)
4. Summary of related work and connection to proposed approach (2 sentences, cite remaining papers, if present)
"""

@app.post("/generateLiteratureSurvey/", response_model=ResponseData)
async def generate_literature_survey(request_data: RequestData):
    summary, ids = summarize(request_data.abstract, request_data.words, request_data.papers)
    return {"summary": summary,
            "ids": ids
            }

@app.get("/")
async def root():
    return {"status": 1}

def summarize(query, n_words, n_papers) :
   keywords = helper.extract_keywords(query)
   papers = helper.search_papers(keywords, n_papers*2)
   ranked_papers = helper.re_rank_papers(query, papers, n_papers)
   literature_review, ids = helper.generate_related_work(query, ranked_papers, base_prompt, sentence_plan, n_words)
   return literature_review, ids

if __name__ == '__main__':
   print("Program running")
   uvicorn.run(app)