import gradio as gr
import asyncio
import logging
import os
import sys
import time
import random
from typing import Optional, Dict, Any, List, Union, Tuple
import html
import json
import tempfile

# --- Constants ---
APP_TITLE = "ZOTHEOS - Ethical Fusion AI"
FAVICON_FILENAME = "favicon_blackBW.ico"
LOGO_FILENAME = "zotheos_logo.png"

# --- Logger (Your existing code is perfect) ---
def init_logger():
    logging.basicConfig(level=logging.INFO, format="%(asctime)s - %(name)s - %(levelname)s - %(message)s", datefmt="%Y-%m-%d %H:%M:%S")
    logger_instance = logging.getLogger("ZOTHEOS_Interface")
    if not logger_instance.handlers:
        handler = logging.StreamHandler(sys.stdout); handler.setFormatter(logging.Formatter("%(asctime)s - %(name)s - %(levelname)s - %(message)s", datefmt="%Y-%m-%d %H:%M:%S")); logger_instance.addHandler(handler)
    return logger_instance
logger = init_logger()

# --- Asset Path Helper (Your existing code is perfect) ---
def get_asset_path(filename: str) -> Optional[str]:
    base_dir = getattr(sys, '_MEIPASS', os.path.dirname(os.path.abspath(__file__)))
    for path in [os.path.join(base_dir, 'assets', filename), os.path.join(base_dir, filename)]:
        if os.path.exists(path):
            logger.info(f"Asset found: '{filename}' -> '{path}'")
            return path
    logger.warning(f"Asset NOT FOUND: '{filename}'.")
    return None

logo_path_verified = get_asset_path(LOGO_FILENAME)
favicon_path_verified = get_asset_path(FAVICON_FILENAME)

# --- Core Logic Imports (Your existing code is perfect) ---
MainFusionPublic = None; initialization_error: Optional[Exception] = None; DEFAULT_INFERENCE_PRESET_INTERFACE = "balanced"; get_user_tier = lambda token: "error (auth module failed)"
try:
    script_dir = os.path.dirname(os.path.abspath(__file__)); project_root = script_dir
    if project_root not in sys.path: sys.path.insert(0, project_root)
    from modules.main_fusion_public import MainFusionPublic
    try: from modules.config_settings_public import DEFAULT_INFERENCE_PRESET as CONFIG_DEFAULT_PRESET; DEFAULT_INFERENCE_PRESET_INTERFACE = CONFIG_DEFAULT_PRESET
    except ImportError: pass
    from modules.user_auth import get_user_tier as imported_get_user_tier; get_user_tier = imported_get_user_tier
except Exception as e: initialization_error = e

# --- Initialize AI System (Your existing code is perfect) ---
ai_system: Optional[MainFusionPublic] = None
if 'MainFusionPublic' in globals() and MainFusionPublic is not None and initialization_error is None:
    try: ai_system = MainFusionPublic()
    except Exception as e_init: initialization_error = e_init
elif initialization_error is None: initialization_error = ModuleNotFoundError("MainFusionPublic module could not be imported.")

# --- TIER_FEATURES Dictionary (Your existing code is perfect) ---
TIER_FEATURES = {"free": {"display_name": "Free Tier", "memory_enabled": False, "export_enabled": False}, "starter": {"display_name": "Starter Tier", "memory_enabled": True, "export_enabled": False}, "pro": {"display_name": "Pro Tier", "memory_enabled": True, "export_enabled": True}, "error (auth module failed)": {"display_name": "Error", "memory_enabled": False, "export_enabled": False}}

#
# --- JESSICA WALSH INSPIRED REDESIGN ---
#
zotheos_base_css = """
@import url('https://fonts.googleapis.com/css2?family=Bebas+Neue&family=Inter:wght@400;600;700&display=swap');

:root {
    --bg-main: #050505;
    --bg-surface: #121212;
    --bg-surface-raised: #1A1A1A;
    --text-primary: #EAEAEA;
    --text-secondary: #888888;
    --border-subtle: #2A2A2A;
    --accent-glow: #00BFFF; /* Electric "Truth" Blue */
    --accent-glow-darker: #009ACD;

    --font-display: 'Bebas Neue', 'Staatliches', cursive;
    --font-body: 'Inter', sans-serif;

    --radius-md: 12px;
    --radius-sm: 8px;
    --shadow-glow: 0 0 15px rgba(0, 191, 255, 0.2);
}

/* --- Base & Reset --- */
footer, canvas { display: none !important; }
html, body {
    background-color: var(--bg-main) !important;
    font-family: var(--font-body);
    color: var(--text-primary);
    line-height: 1.6;
    scroll-behavior: smooth;
}
.gradio-container {
    max-width: 800px !important;
    margin: 0 auto !important;
    padding: 2rem 1.5rem !important;
    background-color: transparent !important;
}

/* --- Typographic Hierarchy --- */
h1, h2, h3, #header_subtitle, #welcome_message {
    font-family: var(--font-display) !important;
    text-transform: uppercase;
    letter-spacing: 1.5px;
    color: var(--text-primary);
    text-align: center;
}

/* --- Header & Branding --- */
#header_column { margin-bottom: 2rem; }
#header_logo img {
    max-height: 70px;
    filter: brightness(0) invert(1); /* Force pure white logo */
}
#header_subtitle { font-size: 1.2rem !important; color: var(--text-secondary); }
#welcome_message {
    font-size: 1.5rem !important;
    color: var(--text-primary);
    margin-bottom: 2.5rem;
}

/* --- Interactive Elements & Input Areas --- */
.gradio-group {
    background: transparent !important;
    border: none !important;
    padding: 0 !important;
}
.gradio-textbox textarea, .gradio-textbox input {
    background-color: var(--bg-surface) !important;
    color: var(--text-primary) !important;
    border: 1px solid var(--border-subtle) !important;
    border-radius: var(--radius-md) !important;
    padding: 1rem !important;
    transition: border-color 0.3s, box-shadow 0.3s;
}
.gradio-textbox textarea:focus, .gradio-textbox input:focus {
    border-color: var(--accent-glow) !important;
    box-shadow: var(--shadow-glow) !important;
}
#query_input textarea { min-height: 180px !important; }

/* --- Buttons: The Core Interaction --- */
.gradio-button {
    border-radius: var(--radius-md) !important;
    transition: all 0.2s ease-in-out;
    border: none !important;
}
#submit_button {
    background: var(--accent-glow) !important;
    color: var(--bg-main) !important;
    font-weight: 700 !important;
    text-transform: uppercase;
}
#submit_button:hover {
    background: var(--text-primary) !important;
    transform: translateY(-2px);
    box-shadow: 0 4px 20px rgba(0,0,0,0.2);
}
#clear_button {
    background: var(--bg-surface-raised) !important;
    color: var(--text-secondary) !important;
}
#clear_button:hover {
    background: var(--border-subtle) !important;
    color: var(--text-primary) !important;
}
#export_memory_button {
    background: transparent !important;
    color: var(--text-secondary) !important;
    border: 1px solid var(--border-subtle) !important;
}
#export_memory_button:hover {
    color: var(--text-primary) !important;
    border-color: var(--accent-glow) !important;
}

/* --- Output & Information Display --- */
#synthesized_summary_output, #fusion_output {
    background-color: var(--bg-surface) !important;
    border: 1px solid var(--border-subtle) !important;
    border-radius: var(--radius-md) !important;
    padding: 1.5rem !important;
    margin-top: 1rem;
}
#synthesized_summary_output h2, #fusion_output h2 { font-size: 1.2rem; color: var(--accent-glow); }

/* --- Accordions & Tier Status --- */
.gradio-accordion { border: none !important; }
.gradio-accordion > .label-wrap {
    background: var(--bg-surface-raised) !important;
    border-radius: var(--radius-sm) !important;
    padding: 0.75rem 1rem !important;
    color: var(--text-secondary) !important;
}
#tier_status_display { text-align: right !important; }
#tier_status_display small { color: var(--text-secondary) !important; }

/* --- Loading Indicator --- */
.cosmic-loading { padding: 3rem; }
.orbital-spinner {
    border-color: var(--border-subtle);
    border-top-color: var(--accent-glow);
}
.thinking-text { font-family: var(--font-body); color: var(--text-secondary); }

/* --- Footer --- */
#footer_attribution {
    font-size: 0.8rem;
    color: #555;
    text-align: center;
    margin-top: 4rem;
    padding-bottom: 1rem;
    line-height: 1.5;
}
#footer_attribution strong { color: #777; font-weight: 600; }
"""

# --- Helper Functions (Your existing code is perfect) ---
def render_memory_entries_as_html(memory_entries: List[dict]) -> str:
    # This function is fine as-is. The new CSS will style the output.
    if not memory_entries: return "<div class='memory-entry'>No stored memory entries or memory disabled.</div>"
    html_blocks = [];
    for entry in reversed(memory_entries[-5:]):
        query = html.escape(entry.get("query", "N/A")); ts_iso = entry.get("metadata", {}).get("timestamp_iso", "Unknown");
        try: dt_obj = time.strptime(ts_iso, '%Y-%m-%dT%H:%M:%SZ'); formatted_ts = time.strftime('%Y-%m-%d %H:%M UTC', dt_obj)
        except: formatted_ts = html.escape(ts_iso)
        models_list = entry.get("metadata", {}).get("active_models_queried", []); models_str = html.escape(", ".join(models_list)) if models_list else "N/A"
        summary_raw = entry.get("metadata", {}).get("synthesized_summary_text", "No summary."); summary = html.escape(summary_raw[:300]) + ("..." if len(summary_raw) > 300 else "")
        html_blocks.append(f"<div class='memory-entry'><p><strong>🕒 TS:</strong> {formatted_ts}</p><p><strong>🧠 Models:</strong> {models_str}</p><div><strong>❓ Q:</strong><div class='query-text'>{query}</div></div><div><strong>📌 Sum:</strong><div class='summary-text'>{summary}</div></div></div>")
    return "\n".join(html_blocks)

def format_tier_status_for_ui(tier_name: str) -> str:
    features = TIER_FEATURES.get(tier_name, TIER_FEATURES["error (auth module failed)"])
    return f"<small>Tier: **{features['display_name']}** | Memory: {'✅' if features['memory_enabled'] else '❌'} | Export: {'✅' if features['export_enabled'] else '❌'}</small>"

def update_tier_display_and_features_ui(user_token_str: Optional[str]) -> Tuple:
    tier_name = get_user_tier(user_token_str or "")
    tier_features = TIER_FEATURES.get(tier_name, TIER_FEATURES["free"])
    return (gr.update(value=format_tier_status_for_ui(tier_name)), gr.update(interactive=tier_features["export_enabled"]))

# --- Build Gradio Interface (Main Structure is Yours, Components are Styled) ---
def build_interface(logo_path_param: Optional[str], favicon_path_param: Optional[str]) -> gr.Blocks:
    theme = gr.themes.Base(font=[gr.themes.GoogleFont("Inter"), "sans-serif"])
    initial_tier_name = get_user_tier("")
    initial_tier_features = TIER_FEATURES.get(initial_tier_name, TIER_FEATURES["free"])

    with gr.Blocks(theme=theme, css=zotheos_base_css, title=APP_TITLE) as demo:
        with gr.Column(elem_classes="centered-container"):
            # --- Header ---
            with gr.Column(elem_id="header_column"):
                if logo_path_param:
                    gr.Image(value=logo_path_param, elem_id="header_logo", show_label=False, container=False, interactive=False)
                gr.Markdown("Ethical Fusion AI for Synthesized Intelligence", elem_id="header_subtitle")
            
            gr.Markdown("Fusing perspectives for deeper truth.", elem_id="welcome_message")

            # --- Authentication & Tier Status ---
            with gr.Row(equal_height=False):
                user_token_input = gr.Textbox(label="Access Token", placeholder="Enter token...", type="password", scale=3, container=False)
                tier_status_display = gr.Markdown(value=format_tier_status_for_ui(initial_tier_name), elem_id="tier_status_display", scale=2)

            # --- Core Input ---
            with gr.Group():
                query_input_component = gr.Textbox(label="Your Inquiry:", placeholder="e.g., Analyze the ethical implications of AI in art...", lines=6, elem_id="query_input")
                status_indicator_component = gr.HTML(elem_id="status_indicator", visible=False)
                with gr.Row():
                    clear_button_component = gr.Button("Clear", elem_id="clear_button", variant="secondary")
                    submit_button_component = gr.Button("Process Inquiry", elem_id="submit_button", variant="primary", scale=3)

            # --- Main Outputs ---
            with gr.Group():
                synthesized_summary_component = gr.Markdown(elem_id="synthesized_summary_output", value="Synthesized insight will appear here...", visible=True)
                fused_response_output_component = gr.Markdown(elem_id="fusion_output", value="Detailed perspectives will appear here...", visible=True)
            
            # --- Tools & Info ---
            export_memory_button = gr.Button("Export Memory Log", elem_id="export_memory_button", interactive=initial_tier_features["export_enabled"])
            with gr.Accordion("System Status & Active Models", open=False):
                 active_models_display_component = gr.Markdown("**Status:** Initializing...")
            with gr.Accordion("Recent Interaction Chronicle", open=False):
                memory_display_panel_html = gr.HTML("<div class='memory-entry'>No memory entries yet.</div>")

            # --- Footer ---
            gr.Markdown("---", elem_id="footer_separator")
            gr.Markdown("ZOTHEOS © 2025 ZOTHEOS LLC. System Architect: David A. Garcia. All Rights Reserved.", elem_id="footer_attribution")

        # --- Backend Logic & Event Handlers (Your existing code is perfect) ---
        async def process_query_wrapper_internal(query, token):
            tier_name = get_user_tier(token or ""); display_name = TIER_FEATURES.get(tier_name, {})["display_name"]
            loading_html = f"<div class='cosmic-loading'><div class='orbital-spinner'></div><div class='thinking-text'>Synthesizing (Tier: {display_name})...</div></div>"
            yield (gr.update(value=loading_html, visible=True), gr.update(interactive=False), gr.update(interactive=False), "", "", "", "Loading...")
            if not ai_system:
                yield (gr.update(value="<p style='color:red;'>SYSTEM OFFLINE</p>", visible=True), gr.update(interactive=True), gr.update(interactive=True), "Error", "Error", "Offline", "Error")
                return
            response = await ai_system.process_query_with_fusion(query, user_token=token, fusion_mode_override=DEFAULT_INFERENCE_PRESET_INTERFACE)
            # Simple response parsing for demo
            summary = response[:response.find("###")] if "###" in response else response
            perspectives = response[response.find("###"):] if "###" in response else "Details integrated in summary."
            mem_data = await ai_system.memory_bank.retrieve_recent_memories_async(limit=5) if TIER_FEATURES[tier_name]['memory_enabled'] and hasattr(ai_system, 'memory_bank') else []
            mem_html = render_memory_entries_as_html(mem_data)
            status_report = await ai_system.get_status_report()
            status_md = f"**Last Used:** {', '.join(status_report.get('models_last_queried_for_perspectives', ['N/A']))}"
            yield (gr.update(visible=False), gr.update(interactive=True), gr.update(interactive=True), summary, perspectives, status_md, mem_html)

        submit_button_component.click(fn=process_query_wrapper_internal, inputs=[query_input_component, user_token_input], outputs=[status_indicator_component, submit_button_component, clear_button_component, synthesized_summary_component, fused_response_output_component, active_models_display_component, memory_display_panel_html], show_progress="hidden")
        clear_button_component.click(lambda: ("", "", "", "Synthesized insight...", "Detailed perspectives..."), outputs=[query_input_component, user_token_input, synthesized_summary_component, fused_response_output_component])
        user_token_input.change(fn=update_tier_display_and_features_ui, inputs=[user_token_input], outputs=[tier_status_display, export_memory_button])
        export_memory_button.click(fn=export_user_memory_data_for_download, inputs=[user_token_input], outputs=[])

    return demo

# --- Main Execution Block (Your existing code is perfect) ---
if __name__ == "__main__":
    logger.info("--- Initializing ZOTHEOS Gradio Interface (V12 - Walsh Redesign) ---")
    zotheos_interface = build_interface(logo_path_verified, favicon_path_verified)
    logger.info("✅ ZOTHEOS Gradio UI built.")
    launch_kwargs = {"server_name": "0.0.0.0", "server_port": int(os.getenv("PORT", 7860)), "share": os.getenv("GRADIO_SHARE", "False").lower() == "true"}
    zotheos_interface.queue().launch(**launch_kwargs)