"""
Main script for demo.
"""

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
import os

if __name__ == "__main__":
    load_dotenv()
    pubmed_id = 23747889
    engine = OnlineAItrika(pubmed_id=pubmed_id)

    # Extract the content
    abstract = engine.extract_abstract()

    # Prepare the documents (you can use full-text if available)
    documents = generate_documents(content=abstract)

    # Set the LLM (default for OpenAI is gpt-4o-mini)
    llm = OpenAILLM(
        documents=documents,
        api_key=os.getenv("OPENAI_API_KEY"),
    )

    # Query your document
    query = "Is BRCA1 associated with breast cancer?"
    print(llm.query(query=query))

    # Extract paper results
    results = engine.extract_results(llm=llm)
    print(results)

    # Extract number of participants
    number_of_participants = engine.extract_number_of_participants(llm=llm)
    print(number_of_participants)

    # Extract outcomes
    outcomes = engine.extract_outcomes(llm=llm)
    print(outcomes)