from fastapi import FastAPI, HTTPException
from fastapi.responses import JSONResponse
from pydantic import BaseModel
from aitrika.engine.online_aitrika import OnlineAItrika
from aitrika.llm.openai import OpenAILLM
from aitrika.utils.text_parser import generate_documents
from dotenv import load_dotenv
from fastapi.openapi.utils import get_openapi
import json
import os
import uvicorn

app = FastAPI(title="AItrika API", description="AItrika API", version="0.1.0")
load_dotenv()


class PubMedRequest(BaseModel):
    pubmed_id: int


class QueryRequest(BaseModel):
    pubmed_id: int
    query: str


@app.post(
    "/associations",
    summary="Get associations from a PubMed article",
    responses={
        200: {"description": "Associations retrieved successfully"},
        422: {"description": "Invalid PubMed ID provided"},
    },
)
def get_associations(request: PubMedRequest):
    """
    Extracts associations from a PubMed article based on its ID.

    Args:
        request (PubMedRequest): An object containing the PubMed ID.

    Returns:
        dict: A dictionary containing the extracted associations.

    Raises:
        HTTPException: If the PubMed ID is invalid or the article is not found.
    """
    try:
        engine = OnlineAItrika(pubmed_id=request.pubmed_id)
        associations = engine.extract_associations()
        return JSONResponse(content={"associations": associations}, status_code=200)
    except ValueError as e:
        raise HTTPException(status_code=422, detail=str(e))
    except Exception:
        raise HTTPException(status_code=500, detail="An unexpected error occurred")


@app.post(
    "/abstract",
    summary="Get abstract of a PubMed article",
    responses={
        200: {"description": "Abstract retrieved successfully"},
        422: {"description": "Invalid PubMed ID provided"},
    },
)
def get_abstract(request: PubMedRequest):
    """
    Retrieves the abstract of a PubMed article based on its ID.

    Args:
        request (PubMedRequest): An object containing the PubMed ID.

    Returns:
        dict: A dictionary containing the abstract of the article.

    Raises:
        HTTPException: If the PubMed ID is invalid or the article is not found.
    """
    try:
        engine = OnlineAItrika(pubmed_id=request.pubmed_id)
        abstract = engine.extract_abstract()
        return JSONResponse(content={"abstract": abstract}, status_code=200)
    except ValueError as e:
        raise HTTPException(status_code=422, detail=str(e))
    except Exception:
        raise HTTPException(status_code=500, detail="An unexpected error occurred")


@app.post(
    "/query",
    summary="Query a PubMed article",
    responses={
        200: {"description": "Query result retrieved successfully"},
        422: {"description": "Invalid PubMed ID or query provided"},
    },
)
def query_document(request: QueryRequest):
    """
    Queries a PubMed article with a user-provided question.

    Args:
        request (QueryRequest): An object containing the PubMed ID and the query.

    Returns:
        dict: A dictionary containing the result of the query.

    Raises:
        HTTPException: If the PubMed ID is invalid, the article is not found, or the query fails.
    """
    try:
        engine = OnlineAItrika(pubmed_id=request.pubmed_id)
        abstract = engine.extract_abstract()
        documents = generate_documents(content=abstract)
        llm = OpenAILLM(
            documents=documents,
            api_key=os.getenv("OPENAI_API_KEY"),
        )
        result = llm.query(query=request.query)
        return JSONResponse(content={"result": result}, status_code=200)
    except ValueError as e:
        raise HTTPException(status_code=422, detail=str(e))
    except Exception:
        raise HTTPException(status_code=500, detail="An unexpected error occurred")


@app.post(
    "/results",
    summary="Get results from a PubMed article",
    responses={
        200: {"description": "Results retrieved successfully"},
        422: {"description": "Invalid PubMed ID provided"},
    },
)
def get_results(request: PubMedRequest):
    """
    Extracts results from a PubMed article based on its ID.

    Args:
        request (PubMedRequest): An object containing the PubMed ID.

    Returns:
        dict: A dictionary containing the extracted results.

    Raises:
        HTTPException: If the PubMed ID is invalid or the article is not found.
    """
    try:
        engine = OnlineAItrika(pubmed_id=request.pubmed_id)
        llm = OpenAILLM(
            documents=generate_documents(content=engine.extract_abstract()),
            api_key=os.getenv("OPENAI_API_KEY"),
        )
        results = engine.extract_results(llm=llm)
        return JSONResponse(content={"results": results}, status_code=200)
    except ValueError as e:
        raise HTTPException(status_code=422, detail=str(e))
    except Exception:
        raise HTTPException(status_code=500, detail="An unexpected error occurred")


@app.post(
    "/participants",
    summary="Get number of participants from a PubMed article",
    responses={
        200: {"description": "Number of participants retrieved successfully"},
        422: {"description": "Invalid PubMed ID provided"},
    },
)
def get_number_of_participants(request: PubMedRequest):
    """
    Extracts the number of participants from a PubMed article based on its ID.

    Args:
        request (PubMedRequest): An object containing the PubMed ID.

    Returns:
        dict: A dictionary containing the number of participants.

    Raises:
        HTTPException: If the PubMed ID is invalid or the article is not found.
    """
    try:
        engine = OnlineAItrika(pubmed_id=request.pubmed_id)
        llm = OpenAILLM(
            documents=generate_documents(content=engine.extract_abstract()),
            api_key=os.getenv("OPENAI_API_KEY"),
        )
        number_of_participants = engine.extract_number_of_participants(llm=llm)
        return JSONResponse(
            content={"number_of_participants": number_of_participants}, status_code=200
        )
    except ValueError as e:
        raise HTTPException(status_code=422, detail=str(e))
    except Exception:
        raise HTTPException(status_code=500, detail="An unexpected error occurred")


@app.post(
    "/outcomes",
    summary="Get outcomes from a PubMed article",
    responses={
        200: {"description": "Outcomes retrieved successfully"},
        422: {"description": "Invalid PubMed ID provided"},
    },
)
def get_outcomes(request: PubMedRequest):
    """
    Extracts outcomes from a PubMed article based on its ID.

    Args:
        request (PubMedRequest): An object containing the PubMed ID.

    Returns:
        dict: A dictionary containing the extracted outcomes.

    Raises:
        HTTPException: If the PubMed ID is invalid or the article is not found.
    """
    try:
        engine = OnlineAItrika(pubmed_id=request.pubmed_id)
        llm = OpenAILLM(
            documents=generate_documents(content=engine.extract_abstract()),
            api_key=os.getenv("OPENAI_API_KEY"),
        )
        outcomes = engine.extract_outcomes(llm=llm)
        return JSONResponse(content={"outcomes": outcomes}, status_code=200)
    except ValueError as e:
        raise HTTPException(status_code=422, detail=str(e))
    except Exception:
        raise HTTPException(status_code=500, detail="An unexpected error occurred")


if __name__ == "__main__":
    with open("docs/api-reference/openapi.json", "w") as f:
        json.dump(
            get_openapi(
                title=app.title,
                version=app.version,
                openapi_version=app.openapi_version,
                description=app.description,
                routes=app.routes,
            ),
            f,
        )
    uvicorn.run(app, host="0.0.0.0", port=8000)