#!/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

# For demonstration, we'll import from huggingface_hub
# (You can omit if you're not using HF or adapt your own client)
from huggingface_hub import InferenceClient

# ----------------------------
# Placeholder data structures
# ----------------------------

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

# Minimal example of a 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"]
    }
}

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

# Just to demonstrate how your "search_arxiv" or "SpeechSynthesis" etc. might be placeholders
def process_text(text):
    st.write(f"process_text called with: {text}")

def process_text2(text_input):
    return f"[process_text2 placeholder] Received: {text_input}"

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}")

def search_glossary(content):
    st.write(f"search_glossary called with: {content}")

# If you have HF Inference endpoint, set them here, else placeholders
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}]"


# --------------------------------------
# Display Entities & Glossary Functions
# --------------------------------------
@st.cache_resource
def display_glossary_entity(k):
    """
    Example of how you'd create multiple links for a glossary entity.
    This was in your original snippet. We'll keep it short.
    """
    search_urls = {
        "🚀🌌ArXiv": lambda k: f"/?q={quote(k)}",
        "🃏Analyst": lambda k: f"/?q={quote(k)}-{quote(PromptPrefix)}",
        "📚PyCoder": lambda k: f"/?q={quote(k)}-{quote(PromptPrefix2)}",
        "🔬JSCoder": lambda k: f"/?q={quote(k)}-{quote(PromptPrefix3)}",
        "📖": lambda k: f"https://en.wikipedia.org/wiki/{quote(k)}",
        "🔍": lambda k: f"https://www.google.com/search?q={quote(k)}",
        "🔎": lambda k: f"https://www.bing.com/search?q={quote(k)}",
        "🎥": lambda k: f"https://www.youtube.com/results?search_query={quote(k)}",
        "🐦": lambda k: f"https://twitter.com/search?q={quote(k)}",
    }
    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)


@st.cache_resource
def display_glossary_grid(roleplaying_glossary):
    """
    Displays a glossary in columns with multiple link emojis.
    """
    search_urls = {
        "🚀🌌ArXiv": lambda k: f"/?q={quote(k)}",
        "🃏Analyst": lambda k: f"/?q={quote(k)}-{quote(PromptPrefix)}",
        "📚PyCoder": lambda k: f"/?q={quote(k)}-{quote(PromptPrefix2)}",
        "🔬JSCoder": lambda k: f"/?q={quote(k)}-{quote(PromptPrefix3)}",
        "📖": lambda k: f"https://en.wikipedia.org/wiki/{quote(k)}",
        "🔍": lambda k: f"https://www.google.com/search?q={quote(k)}",
        "🔎": lambda k: f"https://www.bing.com/search?q={quote(k)}",
        "🎥": lambda k: f"https://www.youtube.com/results?search_query={quote(k)}",
        "🐦": lambda k: f"https://twitter.com/search?q={quote(k)}",
    }

    for category, details in roleplaying_glossary.items():
        st.write(f"### {category}")
        cols = st.columns(len(details))
        for idx, (game, terms) in enumerate(details.items()):
            with cols[idx]:
                st.markdown(f"#### {game}")
                for term in terms:
                    links_md = ' '.join([f"[{emoji}]({url(term)})" for emoji, url in search_urls.items()])
                    st.markdown(f"**{term}** <small>{links_md}</small>", unsafe_allow_html=True)


# --------------------
# File-Handling Logic
# --------------------
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()
    href = f'<a href="data:application/zip;base64,{b64}" download="{zip_file}">Download All</a>'
    return href

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')
        href = f'<a href="data:{mime_type};base64,{b64}" target="_blank" download="{file_name}">{file_name}</a>'
        return href
    except:
        return ''

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

def compare_and_delete_files(files):
    """
    Compare file sizes. If duplicates exist, keep only the latest.
    """
    if not files:
        st.warning("No files to compare.")
        return
    file_sizes = {}
    for file in files:
        size = os.path.getsize(file)
        file_sizes.setdefault(size, []).append(file)
    # Remove all but the latest file for each size
    for size, paths in file_sizes.items():
        if len(paths) > 1:
            latest_file = max(paths, key=os.path.getmtime)
            for file in paths:
                if file != latest_file:
                    os.remove(file)
                    st.success(f"Deleted {file} as a duplicate.")
    st.rerun()

def FileSidebar():
    """
    Renders the file sidebar with all the open/view/run/delete logic.
    """
    all_files = glob.glob("*.md")
    # Example logic filtering filenames
    all_files = [file for file in all_files if len(os.path.splitext(file)[0]) >= 5]
    all_files.sort(key=lambda x: (os.path.splitext(x)[1], x), reverse=True)

    # Buttons for "Delete All" and "Download"
    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 = ''

    # Each file row
    for file in all_files:
        col1, col2, col3, col4, col5 = st.sidebar.columns([1,6,1,1,1])
        with col1:
            # Show an emoji button to do "md"
            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)
                file_name = file
                st.rerun()
                next_action = 'delete'
                st.session_state['next_action'] = next_action

    # Duplicate detection
    file_sizes = [get_file_size(file) for file in all_files]
    previous_size = None
    st.sidebar.title("File Operations")
    for file, size in zip(all_files, file_sizes):
        duplicate_flag = "🚩" if size == previous_size else ""
        with st.sidebar.expander(f"File: {file} {duplicate_flag}"):
            st.text(f"Size: {size} bytes")
            if st.button("View", key=f"view_{file}"):
                try:
                    with open(file, "r", encoding='utf-8') as f:
                        file_content = f.read()
                    st.code(file_content, language="markdown")
                except UnicodeDecodeError:
                    st.error("Failed to decode the file with UTF-8.")
            if st.button("Delete", key=f"delete3_{file}"):
                os.remove(file)
                st.rerun()
        previous_size = size

    # If we have loaded something
    if len(file_contents) > 0:
        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', on_change=None)
                file_content_area = st.text_area('File Contents:', file_contents, height=300, key='file_content_area')

                # Minimal “Save” stubs
                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':
            # Example usage
            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...")

# ---------------------------
# Basic 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):
    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):
    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)
        return score_data["score"]
    return 0

def display_buttons_with_scores(num_columns_text):
    """
    Show buttons that track a 'score' from your glossary data.
    """
    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 = load_score(key)
                if st.button(f"{game_emoji} {category} {game} {term} {score}", key=key):
                    newscore = update_score(key.replace('?',''))
                    st.markdown(f"Scored **{category} - {game} - {term}** -> {newscore}")

# --------------------
# Image & Video Grids
# --------------------
def display_images_and_wikipedia_summaries(num_columns=4):
    """
    Display all .png images in the current directory in a grid, referencing the name as a 'keyword'.
    """
    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

    # Sort by length of filename, just as an example
    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)
                # Provide a text input for user interactions
                image_text_input = st.text_input(f"Prompt for {image_file}", key=f"image_prompt_{image_file}")
                if len(image_text_input) > 0:
                    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):
    """
    Displays all .mp4 or .webm videos found in the current directory in a grid.
    """
    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)
            # Provide a text input
            video_text_input = st.text_input(f"Video Prompt for {video_file}", key=f"video_prompt_{video_file}")
            if video_text_input:
                try:
                    # Hard-coded example
                    seconds_per_frame = 10
                    process_video(video_file, seconds_per_frame)
                except ValueError:
                    st.error("Invalid input for seconds per frame!")
        col_index += 1


# -------------------------------------
# Query Param Helpers from your snippet
# -------------------------------------
def get_all_query_params(key):
    return st.query_params().get(key, [])

def clear_query_params():
    st.query_params()

def display_content_or_image(query):
    """
    If a query matches something in transhuman_glossary or
    a local image, show it. Otherwise warn no match.
    """
    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


# ------------------------------------
# MERMAID DIAGRAM with Clickable Links
# ------------------------------------
def generate_mermaid_html(mermaid_code: str) -> str:
    """
    Returns HTML embedding a Mermaid diagram. We embed the code
    in <div class="mermaid"> and center it with CSS.
    """
    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 {{
                /* Let the diagram scale or otherwise style as you wish */
                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 'Model' checkbox is selected, we append '&model=1' or '?model=1' to the URL.
    We'll handle whether the URL already has a '?' or not.
    """
    if not model_selected:
        return url
    delimiter = "&" if "?" in url else "?"
    return f"{url}{delimiter}model=1"


# For demonstration, we add clickable nodes & edges:
#   click <nodeId> "<URL>" "_self"
# If you want edges to be clickable, you can label them as well,
# but Mermaid typically only has a 'click' property for nodes.
DEFAULT_MERMAID = """
flowchart LR
    %% Notice we have "click LLM ..." lines:
    U((User 😎)) -- "Talk 🗣️" --> LLM[LLM Agent 🤖\\nExtract Info]
    click U "/?q=User%20😎" _self
    click LLM "/?q=LLM%20Agent%20Extract%20Info" _self

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

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

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

    %% If you want an "edge click" to pass ?r= something, 
    %% Mermaid doesn't have direct 'click' for edges, 
    %% but you can define them as nodes or use linkStyle trick, etc.
"""


# ---------------------------
# Streamlit Main App
# ---------------------------
def main():
    st.set_page_config(page_title="Mermaid + Clickable Links Demo", layout="wide")

    # 1) Parse query strings on page load
    query_params = st.experimental_get_query_params()
    current_q = query_params.get("q", [""])[0]  # If present, first string
    current_r = query_params.get("r", [""])[0]

    # 2) Let user pick if they want to add the "model=1" param to clickable links
    st.sidebar.write("## Diagram Link Settings")
    model_selected = st.sidebar.checkbox("Append ?model=1 to each link?")

    # 3) Generate a dynamic Mermaid code, appending model param if user wants
    #    We'll do a simple string replace to incorporate the model param
    #    For a robust approach, parse each URL carefully, then reassemble.
    base_diagram = DEFAULT_MERMAID
    lines = base_diagram.strip().split("\n")
    new_lines = []
    for line in lines:
        if "click " in line and '"/?' in line:
            # e.g. click LLM "/?q=LLM%20Agent" _self
            # let's isolate the URL part
            parts = re.split(r'click\s+\S+\s+"([^"]+)"\s+("_self")', line)
            if len(parts) == 4:
                # parts[0] = 'click LLM '
                # parts[1] = '/?q=LLM%20Agent%20Extract%20Info'
                # parts[2] = ' _self'
                # parts[3] = '' (trailing possibly)
                url = parts[1]
                updated_url = append_model_param(url, model_selected)
                # Recombine
                new_line = f"{parts[0]}\"{updated_url}\" {parts[2]}"
                new_lines.append(new_line)
            else:
                new_lines.append(line)
        else:
            new_lines.append(line)
    mermaid_code = "\n".join(new_lines)

    # 4) Render the top-centered Mermaid diagram
    st.title("Top-Centered Mermaid Diagram with Clickable Links 🏺")
    diagram_html = generate_mermaid_html(mermaid_code)
    components.html(diagram_html, height=400, scrolling=True)

    # 5) Show what the inbound ?q / ?r was
    if current_q:
        st.markdown(f"**Detected Query**: `?q={current_q}`")
        display_content_or_image(current_q)
    if current_r:
        st.markdown(f"**Detected Relationship**: `?r={current_r}`")

    # 6) Editor Columns: 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"
        # Text area
        markdown_text = st.text_area(
            "Edit Markdown:",
            value=st.session_state["markdown_text"],
            height=300
        )
        st.session_state["markdown_text"] = markdown_text

        # Button row
        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.experimental_rerun()

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

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

        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"):
                # Rebuild the diagram
                st.session_state["current_mermaid"] = mermaid_input
                st.write("**Mermaid** diagram refreshed! 🌈")
                st.experimental_rerun()
        with colD:
            if st.button("❌ Clear Mermaid"):
                st.session_state["current_mermaid"] = ""
                st.experimental_rerun()

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

    # 7) Show Sliders & image/video 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)

    # 8) Optional "Extended" UI
    showExtendedTextInterface = False
    if showExtendedTextInterface:
        display_glossary_grid(roleplaying_glossary)
        num_columns_text = st.slider("Choose Number of Text Columns", 1, 15, 4, key="num_columns_text")
        display_buttons_with_scores(num_columns_text)
        st.markdown("Extended text interface is on...")

    # 9) Render the file sidebar
    FileSidebar()

    # 10) Random Title at 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"
    ]
    selected_title = random.choice(titles)
    st.markdown(f"**{selected_title}**")


if __name__ == "__main__":
    main()