"""
Streamlit app for demo.
"""

import streamlit as st
from aitrika.engine.online_aitrika import OnlineAItrika
from aitrika.utils.text_parser import generate_documents
from aitrika.llm.openai import OpenAILLM
from dotenv import load_dotenv
import os
import time


load_dotenv()

st.set_page_config(
    page_title="AItrika",
    page_icon="images/logo.png",
)
st.title("AItrika 🧪")


def response_generator(query: str):
    response = llm.query(query=query)
    print(response)
    for word in response.split():
        yield word + " "
        time.sleep(0.05)


def reset_conversation():
    st.session_state.conversation = None
    st.session_state.chat_history = None


def format_list_to_markdown(items):
    markdown_text = ""
    if items[0].strip().startswith("**") and items[0].strip().endswith("**"):
        markdown_text += f"## {items[0].strip('* ')}\n\n"
        items = items[1:]
    for item in items:
        parts = item.split("\n\n")
        if len(parts) >= 2:
            for subitem in parts[1:]:
                markdown_text += f"- {subitem.strip()}\n"
        else:
            markdown_text += f"- {item.strip()}\n"
            markdown_text = markdown_text.replace('"', "")
        markdown_text += "\n"
    return markdown_text


pubmed_id = st.text_input("Enter the PubMed ID", placeholder="23747889")

if pubmed_id:
    engine = OnlineAItrika(pubmed_id=pubmed_id)
    documents = generate_documents(content=engine.extract_abstract())
    llm = OpenAILLM(documents=documents, api_key=os.getenv("GROQ_API_KEY"))
    with st.expander("Select the information that you want to extract: "):
        option = st.selectbox(
            " ",
            (
                "PubMed ID",
                "Title",
                "Abstract",
                "Authors",
                "Full text",
                "Journal",
                "Genes",
                "Diseases",
                "Species",
                "Chemicals",
                "Mutations",
                "Associations between genes and diseases",
                "Results",
                "Bibliography",
                "Methods",
                "Acknowledgements",
                "Introduction",
                "Paper results",
                "Number of participants",
                "Characteristics of participants",
            ),
        )
        if option == "PubMed ID":
            st.markdown("## PubMed ID")
            st.write(engine.extract_pubmed_id())
        elif option == "Title":
            st.markdown("## Title")
            st.write(engine.extract_title())
        elif option == "Abstract":
            st.markdown("## Abstract")
            st.write(engine.extract_abstract())
        elif option == "Authors":
            st.markdown("## Authors")
            st.write(engine.extract_authors())
        elif option == "Full text":
            st.markdown("## Full text")
            st.write(engine.extract_full_text())
        elif option == "Journal":
            st.markdown("## Journal")
            st.write(engine.extract_journal())
        elif option == "Genes":
            st.markdown("## Genes")
            st.dataframe(engine.extract_genes(dataframe=True))
            st.json(engine.genes())
        elif option == "Diseases":
            st.markdown("## Diseases")
            st.dataframe(engine.extract_diseases(dataframe=True))
            st.json(engine.diseases())
        elif option == "Associations between genes and diseases":
            st.markdown("## Associations between genes and diseases")
            st.dataframe(engine.extract_associations(dataframe=True))
            st.json(engine.extract_associations())
        elif option == "Species":
            st.markdown("## Species")
            st.dataframe(engine.extract_species(dataframe=True))
            st.json(engine.extract_species())
        elif option == "Mutations":
            st.markdown("## Mutations")
            st.dataframe(engine.extract_mutations(dataframe=True))
            st.json(engine.extract_mutations())
        elif option == "Chemicals":
            st.markdown("## Chemicals")
            st.dataframe(engine.extract_chemicals(dataframe=True))
            st.json(engine.extract_chemicals())
        elif option == "Results":
            st.markdown("## Results")
            st.write(engine.extract_results(llm=llm))
        elif option == "Bibliography":
            st.markdown("## Bibliography")
            st.write(engine.extract_bibliography(llm=llm))
        elif option == "Methods":
            st.markdown("## Methods")
            methods = engine.extract_methods(llm=llm).split("---")
            formatted_methods = format_list_to_markdown(methods)
            st.markdown(formatted_methods)
        elif option == "Acknowledgements":
            st.markdown("## Acknowledgements")
            st.write(engine.extract_acknowledgements(llm=llm))
        elif option == "Introduction":
            st.markdown("## Introduction")
            st.write(engine.extract_introduction(llm=llm))
        elif option == "Number of participants":
            st.markdown("## Number of participants")
            st.write(engine.extract_number_of_participants(llm=llm))
        elif option == "Paper results":
            st.markdown("## Paper results")
            paper_results = engine.extract_paper_results(llm=llm).split("---")
            formatted_paper_results = format_list_to_markdown(paper_results)
            st.markdown(formatted_paper_results)
        elif option == "Characteristics of participants":
            st.markdown("## Characteristics of participants")
            characteristics_of_participants = (
                engine.extract_characteristics_of_participants(llm=llm).split("---")
            )
            formatted_characteristics_of_participants = format_list_to_markdown(
                characteristics_of_participants
            )
            st.markdown(formatted_characteristics_of_participants)