import json
import os
import re
from pathlib import Path

import streamlit as st
from google import genai
from google.genai import types

from app import agent, config


def format_output(response: list[dict]) -> tuple[str, str]:
    try:
        answer = ""
        citations = {}
        for statement in response:
            text = statement["text"].strip()
            answer = (
                answer + f"\n{text}"
                if text.startswith("*") or text.startswith("-")
                else answer + f" {text}"
            )
            citation_ids = []
            for source in statement.get("sources", []):
                source_str = f"[{source['title']}]({source['url']})"
                if not (citation_id := citations.get(source_str)):
                    citation_id = len(citations) + 1
                    citations[source_str] = citation_id
                citation_ids.append(citation_id)
            if citation_ids:
                answer += " ".join(f"[^{i}]" for i in sorted(citation_ids))
    except KeyError as err:
        print(err)
        return str(response), ""

    footnotes = "\n".join(f"[^{id}]: {citation}" for citation, id in citations.items())
    return answer, footnotes


def main():
    gemini_client = genai.Client(api_key=config.settings.google_api_key)

    st.title("Elna")
    with st.form("search", border=False):
        query = st.text_input("Your medical question")
        submit = st.form_submit_button("Ask")
        response = st.empty()

        if submit:
            with st.spinner("Thinking...", show_time=True):
                output = agent.respond(gemini_client, query)

            answer, footnotes = format_output(output)
            response.markdown(f"{answer}\n\n{footnotes}")


if __name__ == "__main__":
    main()