#!/usr/bin/env python3

import os
import re
import glob
import json
import base64
import zipfile
import random
import requests
import openai
from PIL import Image
from urllib.parse import quote

import streamlit as st
import streamlit.components.v1 as components

# If you do model inference via huggingface_hub:
from huggingface_hub import InferenceClient

# ----------------------------
# Configurable BASE_URL
# ----------------------------
BASE_URL = "https://huggingface.co/spaces/awacke1/MermaidMarkdownDiagramEditor"

# Example placeholders for prompt prefixes
PromptPrefix = "AI-Search: "
PromptPrefix2 = "AI-Refine: "
PromptPrefix3 = "AI-JS: "

# Example roleplaying glossary
roleplaying_glossary = {
    "Core Rulebooks": {
        "Dungeons and Dragons": ["Player's Handbook", "Dungeon Master's Guide", "Monster Manual"],
        "GURPS": ["Basic Set Characters", "Basic Set Campaigns"]
    },
    "Campaigns & Adventures": {
        "Pathfinder": ["Rise of the Runelords", "Curse of the Crimson Throne"]
    }
}

# Example transhuman glossary
transhuman_glossary = {
    "Neural Interfaces": ["Cortex Jack", "Mind-Machine Fusion"],
    "Cybernetics": ["Robotic Limbs", "Augmented Eyes"],
}

# Simple function stubs
def process_text(text):
    st.write(f"process_text called with: {text}")

def search_arxiv(text):
    st.write(f"search_arxiv called with: {text}")

def SpeechSynthesis(text):
    st.write(f"SpeechSynthesis called with: {text}")

def process_image(image_file, prompt):
    return f"[process_image placeholder] Processing {image_file} with prompt: {prompt}"

def process_video(video_file, seconds_per_frame):
    st.write(f"[process_video placeholder] Video: {video_file}, seconds/frame: {seconds_per_frame}")

# Stub if you have a Hugging Face endpoint
API_URL = "https://huggingface-inference-endpoint-placeholder"
API_KEY = "hf_XXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXX"

@st.cache_resource
def InferenceLLM(prompt):
    return f"[InferenceLLM placeholder response to prompt: {prompt}]"

# ------------------------------------------
# Glossary & File Utility
# ------------------------------------------
@st.cache_resource
def display_glossary_entity(k):
    """
    Creates multiple link emojis for a single entity.
    """
    search_urls = {
        "🚀🌌ArXiv": lambda x: f"/?q={quote(x)}",
        "🃏Analyst": lambda x: f"/?q={quote(x)}-{quote(PromptPrefix)}",
        "📚PyCoder": lambda x: f"/?q={quote(x)}-{quote(PromptPrefix2)}",
        "🔬JSCoder": lambda x: f"/?q={quote(x)}-{quote(PromptPrefix3)}",
        "📖": lambda x: f"https://en.wikipedia.org/wiki/{quote(x)}",
        "🔍": lambda x: f"https://www.google.com/search?q={quote(x)}",
        "🔎": lambda x: f"https://www.bing.com/search?q={quote(x)}",
        "🎥": lambda x: f"https://www.youtube.com/results?search_query={quote(x)}",
        "🐦": lambda x: f"https://twitter.com/search?q={quote(x)}",
    }
    links_md = ' '.join([f"[{emoji}]({url(k)})" for emoji, url in search_urls.items()])
    st.markdown(f"**{k}** <small>{links_md}</small>", unsafe_allow_html=True)

def display_content_or_image(query):
    """
    If a query matches something in transhuman_glossary or a local image, show it.
    """
    for category, term_list in transhuman_glossary.items():
        for term in term_list:
            if query.lower() in term.lower():
                st.subheader(f"Found in {category}:")
                st.write(term)
                return True
    image_path = f"images/{query}.png"
    if os.path.exists(image_path):
        st.image(image_path, caption=f"Image for {query}")
        return True
    st.warning("No matching content or image found.")
    return False

def clear_query_params():
    """
    For clearing URL params, you'd typically use a new link or st.experimental_set_query_params().
    Here, we just warn the user. 
    """
    st.warning("Define a redirect or link without query params if you want to truly clear them.")

# -----------------------
# File Handling
# -----------------------
def load_file(file_path):
    try:
        with open(file_path, "r", encoding='utf-8') as f:
            return f.read()
    except:
        return ""

@st.cache_resource
def create_zip_of_files(files):
    zip_name = "Arxiv-Paper-Search-QA-RAG-Streamlit-Gradio-AP.zip"
    with zipfile.ZipFile(zip_name, 'w') as zipf:
        for file in files:
            zipf.write(file)
    return zip_name

@st.cache_resource
def get_zip_download_link(zip_file):
    with open(zip_file, 'rb') as f:
        data = f.read()
    b64 = base64.b64encode(data).decode()
    return f'<a href="data:application/zip;base64,{b64}" download="{zip_file}">Download All</a>'

def get_table_download_link(file_path):
    """
    Creates a download link for a single file from your snippet.
    """
    try:
        with open(file_path, 'r', encoding='utf-8') as file:
            data = file.read()
        b64 = base64.b64encode(data.encode()).decode()
        file_name = os.path.basename(file_path)
        ext = os.path.splitext(file_name)[1]
        mime_map = {
            '.txt': 'text/plain',
            '.py': 'text/plain',
            '.xlsx': 'text/plain',
            '.csv': 'text/plain',
            '.htm': 'text/html',
            '.md': 'text/markdown',
            '.wav': 'audio/wav'
        }
        mime_type = mime_map.get(ext, 'application/octet-stream')
        return f'<a href="data:{mime_type};base64,{b64}" target="_blank" download="{file_name}">{file_name}</a>'
    except:
        return ''

def get_file_size(file_path):
    return os.path.getsize(file_path)

def FileSidebar():
    """
    Renders .md files, providing open/view/delete/run logic in the sidebar.
    """
    all_files = glob.glob("*.md")
    # Exclude short-named or special files if needed
    all_files = [f for f in all_files if len(os.path.splitext(f)[0]) >= 5]
    all_files.sort(key=lambda x: (os.path.splitext(x)[1], x), reverse=True)

    Files1, Files2 = st.sidebar.columns(2)
    with Files1:
        if st.button("🗑 Delete All"):
            for file in all_files:
                os.remove(file)
            st.rerun()
    with Files2:
        if st.button("⬇️ Download"):
            zip_file = create_zip_of_files(all_files)
            st.sidebar.markdown(get_zip_download_link(zip_file), unsafe_allow_html=True)

    file_contents = ''
    file_name = ''
    next_action = ''

    for file in all_files:
        col1, col2, col3, col4, col5 = st.sidebar.columns([1, 6, 1, 1, 1])
        with col1:
            if st.button("🌐", key="md_" + file):
                file_contents = load_file(file)
                file_name = file
                next_action = 'md'
                st.session_state['next_action'] = next_action
        with col2:
            st.markdown(get_table_download_link(file), unsafe_allow_html=True)
        with col3:
            if st.button("📂", key="open_" + file):
                file_contents = load_file(file)
                file_name = file
                next_action = 'open'
                st.session_state['lastfilename'] = file
                st.session_state['filename'] = file
                st.session_state['filetext'] = file_contents
                st.session_state['next_action'] = next_action
        with col4:
            if st.button("▶️", key="read_" + file):
                file_contents = load_file(file)
                file_name = file
                next_action = 'search'
                st.session_state['next_action'] = next_action
        with col5:
            if st.button("🗑", key="delete_" + file):
                os.remove(file)
                st.rerun()

    # If we loaded a file
    if file_contents:
        if next_action == 'open':
            open1, open2 = st.columns([0.8, 0.2])
            with open1:
                file_name_input = st.text_input('File Name:', file_name, key='file_name_input')
                file_content_area = st.text_area('File Contents:', file_contents, height=300, key='file_content_area')

                if st.button('💾 Save File'):
                    with open(file_name_input, 'w', encoding='utf-8') as f:
                        f.write(file_content_area)
                    st.markdown(f'Saved {file_name_input} successfully.')

        elif next_action == 'search':
            file_content_area = st.text_area("File Contents:", file_contents, height=500)
            user_prompt = PromptPrefix2 + file_contents
            st.markdown(user_prompt)
            if st.button('🔍Re-Code'):
                search_arxiv(file_contents)

        elif next_action == 'md':
            st.markdown(file_contents)
            SpeechSynthesis(file_contents)
            if st.button("🔍Run"):
                st.write("Running GPT logic placeholder...")

# ---------------------------
# Scoring / Glossaries
# ---------------------------
score_dir = "scores"
os.makedirs(score_dir, exist_ok=True)

def generate_key(label, header, idx):
    return f"{header}_{label}_{idx}_key"

def update_score(key, increment=1):
    """
    Track a 'score' for each glossary item or term, saved in JSON per key.
    """
    score_file = os.path.join(score_dir, f"{key}.json")
    if os.path.exists(score_file):
        with open(score_file, "r") as file:
            score_data = json.load(file)
    else:
        score_data = {"clicks": 0, "score": 0}
    score_data["clicks"] += increment
    score_data["score"] += increment
    with open(score_file, "w") as file:
        json.dump(score_data, file)
    return score_data["score"]

def load_score(key):
    file_path = os.path.join(score_dir, f"{key}.json")
    if os.path.exists(file_path):
        with open(file_path, "r") as file:
            score_data = json.load(file)
        return score_data["score"]
    return 0

def display_buttons_with_scores(num_columns_text):
    """
    Show glossary items as clickable buttons that increment a 'score'.
    """
    game_emojis = {
        "Dungeons and Dragons": "🐉",
        "Call of Cthulhu": "🐙",
        "GURPS": "🎲",
        "Pathfinder": "🗺️",
        "Kindred of the East": "🌅",
        "Changeling": "🍃",
    }
    topic_emojis = {
        "Core Rulebooks": "📚",
        "Maps & Settings": "🗺️",
        "Game Mechanics & Tools": "⚙️",
        "Monsters & Adversaries": "👹",
        "Campaigns & Adventures": "📜",
        "Creatives & Assets": "🎨",
        "Game Master Resources": "🛠️",
        "Lore & Background": "📖",
        "Character Development": "🧍",
        "Homebrew Content": "🔧",
        "General Topics": "🌍",
    }

    for category, games in roleplaying_glossary.items():
        category_emoji = topic_emojis.get(category, "🔍")
        st.markdown(f"## {category_emoji} {category}")
        for game, terms in games.items():
            game_emoji = game_emojis.get(game, "🎮")
            for term in terms:
                key = f"{category}_{game}_{term}".replace(' ', '_').lower()
                score_val = load_score(key)
                if st.button(f"{game_emoji} {category} {game} {term} {score_val}", key=key):
                    newscore = update_score(key.replace('?', ''))
                    st.markdown(f"Scored **{category} - {game} - {term}** -> {newscore}")

# -------------------------------
# Image & Video
# -------------------------------
def display_images_and_wikipedia_summaries(num_columns=4):
    image_files = [f for f in os.listdir('.') if f.endswith('.png')]
    if not image_files:
        st.write("No PNG images found in the current directory.")
        return

    image_files_sorted = sorted(image_files, key=lambda x: len(x.split('.')[0]))
    cols = st.columns(num_columns)
    col_index = 0

    for image_file in image_files_sorted:
        with cols[col_index % num_columns]:
            try:
                image = Image.open(image_file)
                st.image(image, use_column_width=True)
                k = image_file.split('.')[0]
                display_glossary_entity(k)
                image_text_input = st.text_input(f"Prompt for {image_file}", key=f"image_prompt_{image_file}")
                if image_text_input:
                    response = process_image(image_file, image_text_input)
                    st.markdown(response)
            except:
                st.write(f"Could not open {image_file}")
        col_index += 1

def display_videos_and_links(num_columns=4):
    video_files = [f for f in os.listdir('.') if f.endswith(('.mp4', '.webm'))]
    if not video_files:
        st.write("No MP4 or WEBM videos found in the current directory.")
        return

    video_files_sorted = sorted(video_files, key=lambda x: len(x.split('.')[0]))
    cols = st.columns(num_columns)
    col_index = 0

    for video_file in video_files_sorted:
        with cols[col_index % num_columns]:
            k = video_file.split('.')[0]
            st.video(video_file, format='video/mp4', start_time=0)
            display_glossary_entity(k)
            video_text_input = st.text_input(f"Video Prompt for {video_file}", key=f"video_prompt_{video_file}")
            if video_text_input:
                try:
                    seconds_per_frame = 10
                    process_video(video_file, seconds_per_frame)
                except ValueError:
                    st.error("Invalid input for seconds per frame!")
        col_index += 1

# --------------------------------
# MERMAID DIAGRAM
# --------------------------------
def generate_mermaid_html(mermaid_code: str) -> str:
    """
    Returns HTML that centers the Mermaid diagram, loading from a CDN.
    """
    return f"""
    <html>
    <head>
        <script src="https://cdn.jsdelivr.net/npm/mermaid/dist/mermaid.min.js"></script>
        <style>
            .centered-mermaid {{
                display: flex;
                justify-content: center;
                margin: 20px auto;
            }}
            .mermaid {{
                max-width: 800px;
            }}
        </style>
    </head>
    <body>
        <div class="mermaid centered-mermaid">
            {mermaid_code}
        </div>
        <script>
            mermaid.initialize({{ startOnLoad: true }});
        </script>
    </body>
    </html>
    """

def append_model_param(url: str, model_selected: bool) -> str:
    """
    If user checks 'Append ?model=1', we append &model=1 or ?model=1 if not present.
    """
    if not model_selected:
        return url
    delimiter = "&" if "?" in url else "?"
    return f"{url}{delimiter}model=1"

def inject_base_url(url: str) -> str:
    """
    If a link does not start with http, prepend your BASE_URL
    so it becomes an absolute link to huggingface.co/spaces/...
    """
    if url.startswith("http"):
        return url
    return f"{BASE_URL}{url}"

# Our default diagram, containing the "click" lines with /?q=...
DEFAULT_MERMAID = """
flowchart LR
    U((User 😎)) -- "Talk 🗣️" --> LLM[LLM Agent 🤖\\nExtract Info]
    click U "/?q=User%20😎" _self
    click LLM "/?q=LLM%20Agent%20Extract%20Info" _blank

    LLM -- "Query 🔍" --> HS[Hybrid Search 🔎\\nVector+NER+Lexical]
    click HS "/?q=Hybrid%20Search%20Vector+NER+Lexical" _blank

    HS -- "Reason 🤔" --> RE[Reasoning Engine 🛠️\\nNeuralNetwork+Medical]
    click RE "/?q=Reasoning%20Engine%20NeuralNetwork+Medical" _blank

    RE -- "Link 📡" --> KG((Knowledge Graph 📚\\nOntology+GAR+RAG))
    click KG "/?q=Knowledge%20Graph%20Ontology+GAR+RAG" _blank
"""

def main():
    st.set_page_config(page_title="Mermaid + Clickable Links with Base URL", layout="wide")

    # ---------------------------------------------
    # Query Param Parsing (non-experimental)
    # ---------------------------------------------
    query_params = st.query_params
    query_list = (query_params.get('q') or query_params.get('query') or [''])
    q_or_query = query_list[0] if query_list else ''
    if q_or_query.strip():
        # If there's a q= or query= param, do some processing
        search_payload = PromptPrefix + q_or_query
        st.markdown(search_payload)
        process_text(search_payload)

    # If an 'action' param is present
    if 'action' in query_params:
        action_list = query_params['action']
        if action_list:
            action = action_list[0]
            if action == 'show_message':
                st.success("Showing a message because 'action=show_message' was found in the URL.")
            elif action == 'clear':
                clear_query_params()

    # If a 'query' param is present, show content or image
    if 'query' in query_params:
        query_val = query_params['query'][0]
        display_content_or_image(query_val)

    # ---------------------------------------------
    # Let user pick if we want ?model=1
    # ---------------------------------------------
    st.sidebar.write("## Diagram Link Settings")
    model_selected = st.sidebar.checkbox("Append ?model=1 to each link?")

    # ---------------------------------------------
    # Rebuild the clickable lines in the Mermaid code
    # ---------------------------------------------
    base_diagram = DEFAULT_MERMAID
    lines = base_diagram.strip().split("\n")
    new_lines = []

    for line in lines:
        # We look for lines like: click SOMENODE "/?q=Something" _self
        if "click " in line and '"/?' in line:
            # Try to extract the URL part
            parts = re.split(r'click\s+\S+\s+"([^"]+)"\s+("_self")', line)
            if len(parts) == 4:
                # Example:
                # parts[0] -> 'click LLM '
                # parts[1] -> '/?q=LLM%20Agent%20Extract%20Info'
                # parts[2] -> ' _self'
                # parts[3] -> '' or trailing
                old_url = parts[1]
                # 1) Prepend base if needed
                new_url = inject_base_url(old_url)
                # 2) Possibly add &model=1
                new_url = append_model_param(new_url, model_selected)

                # Recombine
                new_line = f"{parts[0]}\"{new_url}\" {parts[2]}"
                new_lines.append(new_line)
            else:
                # If we can't parse it, keep it as is
                new_lines.append(line)
        else:
            new_lines.append(line)

    mermaid_code = "\n".join(new_lines)

    # ---------------------------------------------
    # Render the top-centered Mermaid diagram
    # ---------------------------------------------
    st.title("Mermaid Diagram with Base URL Injection")
    diagram_html = generate_mermaid_html(mermaid_code)
    components.html(diagram_html, height=400, scrolling=True)

    # ---------------------------------------------
    # Two-column interface: Markdown & Mermaid
    # ---------------------------------------------
    left_col, right_col = st.columns(2)

    # --- Left: Markdown Editor
    with left_col:
        st.subheader("Markdown Side 📝")
        if "markdown_text" not in st.session_state:
            st.session_state["markdown_text"] = "## Hello!\nType some *Markdown* here.\n"
        markdown_text = st.text_area(
            "Edit Markdown:",
            value=st.session_state["markdown_text"],
            height=300
        )
        st.session_state["markdown_text"] = markdown_text

        colA, colB = st.columns(2)
        with colA:
            if st.button("🔄 Refresh Markdown"):
                st.write("**Markdown** content refreshed! 🍿")
        with colB:
            if st.button("❌ Clear Markdown"):
                st.session_state["markdown_text"] = ""
                st.rerun()

        st.markdown("---")
        st.markdown("**Preview:**")
        st.markdown(markdown_text)

    # --- Right: Mermaid Editor
    with right_col:
        st.subheader("Mermaid Side 🧜‍♂️")

        # We store the final code in session state, so user can edit
        if "current_mermaid" not in st.session_state:
            st.session_state["current_mermaid"] = mermaid_code

        mermaid_input = st.text_area(
            "Edit Mermaid Code:",
            value=st.session_state["current_mermaid"],
            height=300
        )

        colC, colD = st.columns(2)
        with colC:
            if st.button("🎨 Refresh Diagram"):
                st.session_state["current_mermaid"] = mermaid_input
                st.write("**Mermaid** diagram refreshed! 🌈")
                st.rerun()
        with colD:
            if st.button("❌ Clear Mermaid"):
                st.session_state["current_mermaid"] = ""
                st.rerun()

        st.markdown("---")
        st.markdown("**Mermaid Source:**")
        st.code(mermaid_input, language="python", line_numbers=True)

    # ---------------------------------------------
    # Media Galleries
    # ---------------------------------------------
    st.markdown("---")
    st.header("Media Galleries")

    num_columns_images = st.slider("Choose Number of Image Columns", 1, 15, 5, key="num_columns_images")
    display_images_and_wikipedia_summaries(num_columns_images)

    num_columns_video = st.slider("Choose Number of Video Columns", 1, 15, 5, key="num_columns_video")
    display_videos_and_links(num_columns_video)

    # (Optionally) Extended text interface
    showExtendedTextInterface = False
    if showExtendedTextInterface:
        # For example:
        # display_glossary_grid(roleplaying_glossary)
        # num_columns_text = st.slider("Choose Number of Text Columns", 1, 15, 4)
        # display_buttons_with_scores(num_columns_text)
        pass

    # ---------------------------------------------
    # File Sidebar
    # ---------------------------------------------
    FileSidebar()

    # ---------------------------------------------
    # Random Title at the bottom
    # ---------------------------------------------
    titles = [
        "🧠🎭 Semantic Symphonies & Episodic Encores",
        "🌌🎼 AI Rhythms of Memory Lane",
        "🎭🎉 Cognitive Crescendos & Neural Harmonies",
        "🧠🎺 Mnemonic Melodies & Synaptic Grooves",
        "🎼🎸 Straight Outta Cognition",
        "🥁🎻 Jazzy Jambalaya of AI Memories",
        "🏰 Semantic Soul & Episodic Essence",
        "🥁🎻 The Music Of AI's Mind"
    ]
    st.markdown(f"**{random.choice(titles)}**")


if __name__ == "__main__":
    main()