import streamlit as st
from streamlit_markmap import markmap
from langchain.llms import OpenAI
from langchain.prompts import PromptTemplate
import os
from langchain_openai import ChatOpenAI


def remove_markdown_fences(text: str) -> str:
    # Entferne führende Code-Fence (entweder "```markdown" oder "```")
    if text.startswith("```markdown"):
        text = text[len("```markdown"):].strip()
    elif text.startswith("```"):
        text = text[len("```"):].strip()
    # Entferne abschließende Code-Fence "```", falls vorhanden
    if text.endswith("```"):
        text = text[:-3].strip()
    return text

# Streamlit-Konfiguration
st.set_page_config(page_title="Mindmap Generator", layout="wide")
st.title("Mindmap Ideen-Generator mit LangChain und OpenAI")

# Sidebar: Eingabe des OpenAI API Keys
st.sidebar.header("Einstellungen")
api_key = st.sidebar.text_input("OpenAI API Key", type="password")
os.environ['OPENAI_API_KEY'] = api_key

# Texteingabe für den Brainstorming-Prompt
user_prompt = st.text_input(
    "Thema"
)

# Button, um die Generierung zu starten
if st.button("Generiere Mindmap"):
    if not user_prompt.strip():
        st.error("Bitte gib ein Thema ein!")
    else:
        with st.spinner("Generiere Mindmap..."):
            # Prompt-Vorlage definieren, die den LLM anweist, eine Markmap im Markdown-Format zu erzeugen.
            template = """
            Du bist ein kreativer Ideengeber und sollst ein Brainstorming zu folgendem Thema erstellen:
            {topic}

            Generiere aus fünf Perspektiven mindestens jeweils zehn Ideen.
            
            Erstelle bitte eine Mindmap im Markdown-Format (Markmap), in der die Hauptideen und Unterideen hierarchisch strukturiert sind.
            
            Nutze folgendes Format:
            # Hauptidee
            ## Unteridee 1
            ## Unteridee 2
            ### Weitere Unteridee
            """
            prompt = PromptTemplate(template=template, input_variables=["topic"])
            full_prompt = prompt.format(topic=user_prompt)
            
            llm = ChatOpenAI(
                model="gpt-4o",
                temperature=0,
                max_tokens=None,
                timeout=None,
                max_retries=2
            )
            
            response = llm.invoke(full_prompt)
            # Anfrage an das LLM senden
            mindmap_markdown = response.content
            mindmap_markdown = remove_markdown_fences(mindmap_markdown)

            # Ausgabe der generierten Mindmap
            with st.expander("Generated Markwodn"):
                st.markdown("## Generierte Mindmap (Markdown):")
                st.code(mindmap_markdown, language="markdown")
            
            with st.expander("Prompt"):
                st.markdown("## Prompt")
                st.code(full_prompt, language="markdown")


            # Darstellung der Mindmap mit streamlit_markmap
            markmap(mindmap_markdown, height=800)