#!/usr/bin/env python3
"""
ECH0-PRIME Gradio API Interface
Enables external LLMs to navigate and use QuLab tools through ECH0-PRIME

Copyright (c) 2025 Joshua Hendricks Cole (DBA: Corporation of Light). All Rights Reserved. PATENT PENDING.
"""

import gradio as gr
import json
from typing import Dict, List, Any, Optional
import traceback
from datetime import datetime
import os
import tempfile

# Get output directory (works on HF Spaces and local)
def get_output_dir():
    """Get the output directory for saving sessions"""
    if os.path.exists("/tmp"):
        # HF Spaces and Unix-like systems
        output_dir = "/tmp/ech0_sessions"
    else:
        # Windows or other systems
        output_dir = os.path.join(tempfile.gettempdir(), "ech0_sessions")

    os.makedirs(output_dir, exist_ok=True)
    return output_dir

# Initialize ECH0-PRIME
ech0_prime = None
ech0_error = None

try:
    from ech0_prime_qulab_connector import Ech0PrimeQuLabResearcher
    ech0_prime = Ech0PrimeQuLabResearcher()
    print("✅ ECH0-PRIME initialized successfully for Gradio API")
except Exception as e:
    ech0_error = str(e)
    print(f"❌ ECH0-PRIME initialization failed: {e}")


def list_available_labs() -> str:
    """List all available labs in QuLab Infinite"""
    if not ech0_prime:
        return "❌ ECH0-PRIME not available"

    try:
        labs = ech0_prime.qulab.list_labs()

        if not labs:
            return """⚠️ **QuLab API may not be running**

Available Labs (from system prompt):
### QUANTUM LABS
- quantum_lab: Quantum computing, teleportation protocols, error correction
- quantum_mechanics_lab: Fundamental quantum simulations
- quantum_computing_lab: Gate operations, circuit design

### BIOLOGICAL LABS
- oncology_lab: Cancer simulations, tumor evolution, treatment optimization
- genetic_variant_analyzer: Genomic analysis, variant impact prediction
- cancer_metabolic_optimizer: Metabolic pathway analysis for cancer treatment
- immune_response_simulator: Immunotherapy modeling
- microbiome_optimizer: Gut microbiome analysis and optimization
- neurotransmitter_optimizer: Brain chemistry modeling
- stem_cell_predictor: Differentiation pathway prediction
- protein_folding_lab: Protein structure prediction

### CHEMISTRY LABS
- chemistry_lab: Molecular simulations, reaction kinetics
- drug_interaction_network: Drug-drug interactions, pharmacology
- materials_lab: Materials science simulations
- organic_chemistry_lab: Organic synthesis planning
- electrochemistry_lab: Electrochemical reactions

### PHYSICS LABS
- physics_engine: Classical and quantum physics simulations
- particle_physics_lab: Particle interactions
- plasma_physics_lab: Plasma dynamics
- condensed_matter_physics_lab: Solid state physics

### ENVIRONMENTAL LABS
- environmental_sim: Climate and environmental modeling
- atmospheric_chemistry_lab: Atmospheric reactions
- renewable_energy_lab: Energy system optimization

### MEDICAL LABS
- cardiovascular_plaque_lab: Atherosclerosis modeling
- metabolic_syndrome_reversal: Metabolic health optimization
- drug_design_lab: Molecular drug design
- clinical_trials_simulation: Trial outcome prediction"""

        result = "🔬 **Available QuLab Infinite Labs**\n\n"
        for lab in labs:
            result += f"• {lab}\n"

        return result

    except Exception as e:
        return f"❌ Failed to list labs: {str(e)}"


def query_ech0_prime(question: str, include_context: bool = True) -> str:
    """Query ECH0-PRIME with a scientific question"""
    if not ech0_prime:
        return "❌ ECH0-PRIME not available"

    if not question.strip():
        return "Please enter a question"

    try:
        response = ech0_prime.query_llm(question)

        result = f"🤖 **ECH0-PRIME Response**\n\n{response}\n\n"

        if include_context:
            result += f"---\n**Conversation Length**: {len(ech0_prime.conversation_history)} messages\n"
            result += f"**Experiments Run**: {len(ech0_prime.experiment_log)}\n"

        return result

    except Exception as e:
        return f"❌ Query failed: {str(e)}\n\nTraceback:\n{traceback.format_exc()}"


def run_single_experiment(lab: str, experiment_type: str, parameters: str) -> str:
    """Run a single experiment on QuLab Infinite"""
    if not ech0_prime:
        return "❌ ECH0-PRIME not available"

    try:
        # Parse parameters
        if parameters.strip():
            try:
                params = json.loads(parameters)
            except json.JSONDecodeError:
                return "❌ Invalid JSON parameters. Please check your input format."
        else:
            params = {}

        # Create experiment dict
        experiment = {
            "lab": lab,
            "experiment_type": experiment_type,
            "parameters": params
        }

        # Run experiment
        results = ech0_prime.run_experiment(experiment)

        # Format result
        result_text = f"🧪 **Experiment Result**\n\n"
        result_text += f"**Lab**: {lab}\n"
        result_text += f"**Type**: {experiment_type}\n\n"

        if "error" in results:
            result_text += f"❌ **Error**: {results['error']}\n\n"
            if "detail" in results:
                result_text += f"**Details**: {results['detail']}"
        else:
            result_text += f"✅ **Success!**\n\n```json\n{json.dumps(results, indent=2)}\n```"

        # Log to ECH0-PRIME
        ech0_prime.experiment_log.append({
            "experiment": experiment,
            "results": results,
            "timestamp": datetime.now().isoformat()
        })

        return result_text

    except Exception as e:
        return f"❌ Experiment failed: {str(e)}\n\nTraceback:\n{traceback.format_exc()}"


def start_research_session(research_question: str, max_iterations: int = 5) -> str:
    """Start an autonomous research session"""
    if not ech0_prime:
        return "❌ ECH0-PRIME not available"

    if not research_question.strip():
        return "Please enter a research question"

    try:
        # Run research loop
        results = ech0_prime.research_loop(
            research_question,
            max_iterations=max_iterations,
            verbose=False
        )

        # Format results
        output = f"🔬 **Research Session Complete**\n\n"
        output += f"**Question**: {research_question}\n"
        output += f"**Iterations**: {results['iterations']}\n\n"

        output += "### Experiments Run:\n"
        for i, exp in enumerate(results['experiments'], 1):
            exp_data = exp['experiment']
            output += f"\n**{i}. {exp_data.get('lab')}** - {exp_data.get('experiment_type')}\n"

            # Show result status
            if 'error' in exp['results']:
                output += f"   ❌ Error: {exp['results']['error']}\n"
            else:
                output += f"   ✅ Success\n"

        output += f"\n### Final Summary:\n\n{results['final_summary']}\n\n"

        # Save session
        timestamp = datetime.now().strftime("%Y%m%d_%H%M%S")
        output_dir = get_output_dir()
        filepath = os.path.join(output_dir, f"ech0_research_gradio_{timestamp}.json")
        ech0_prime.save_session(filepath)
        output += f"---\n**Session saved to**: `{filepath}`"

        return output

    except Exception as e:
        return f"❌ Research session failed: {str(e)}\n\nTraceback:\n{traceback.format_exc()}"


def get_conversation_history() -> str:
    """Get the current conversation history"""
    if not ech0_prime:
        return "❌ ECH0-PRIME not available"

    try:
        if not ech0_prime.conversation_history:
            return "No conversation history yet. Start by asking a question!"

        output = f"💬 **Conversation History** ({len(ech0_prime.conversation_history)} messages)\n\n"

        for i, msg in enumerate(ech0_prime.conversation_history[-10:], 1):  # Last 10 messages
            role = msg['role'].upper()
            content = msg['content'][:200] + "..." if len(msg['content']) > 200 else msg['content']
            output += f"**{role}**: {content}\n\n"

        if len(ech0_prime.conversation_history) > 10:
            output += f"_(Showing last 10 of {len(ech0_prime.conversation_history)} messages)_"

        return output

    except Exception as e:
        return f"❌ Failed to get history: {str(e)}"


def get_experiment_log() -> str:
    """Get the experiment log"""
    if not ech0_prime:
        return "❌ ECH0-PRIME not available"

    try:
        if not ech0_prime.experiment_log:
            return "No experiments run yet."

        output = f"📊 **Experiment Log** ({len(ech0_prime.experiment_log)} experiments)\n\n"

        for i, exp in enumerate(ech0_prime.experiment_log, 1):
            exp_data = exp['experiment']
            output += f"**{i}. {exp_data.get('lab')}** - {exp_data.get('experiment_type')}\n"
            output += f"   Timestamp: {exp.get('timestamp', 'N/A')}\n"

            if 'error' in exp['results']:
                output += f"   ❌ Error: {exp['results']['error']}\n"
            else:
                output += f"   ✅ Success\n"
            output += "\n"

        return output

    except Exception as e:
        return f"❌ Failed to get log: {str(e)}"


def reset_session() -> str:
    """Reset the ECH0-PRIME session"""
    if not ech0_prime:
        return "❌ ECH0-PRIME not available"

    try:
        # Save current session before reset
        if ech0_prime.experiment_log or ech0_prime.conversation_history:
            timestamp = datetime.now().strftime("%Y%m%d_%H%M%S")
            output_dir = get_output_dir()
            filepath = os.path.join(output_dir, f"ech0_session_backup_{timestamp}.json")
            ech0_prime.save_session(filepath)
            backup_msg = f"\n✅ Previous session backed up to: `{filepath}`"
        else:
            backup_msg = ""

        # Reset
        ech0_prime.conversation_history = []
        ech0_prime.experiment_log = []

        return f"🔄 **Session Reset**\n\nConversation history and experiment log cleared.{backup_msg}"

    except Exception as e:
        return f"❌ Reset failed: {str(e)}"


# Create the Gradio interface
def create_interface():
    """Create the ECH0-PRIME Gradio API interface"""

    with gr.Blocks(title="ECH0-PRIME × QuLab Infinite API",
                   theme=gr.themes.Soft()) as interface:

        gr.Markdown("""
        # 🤖 ECH0-PRIME × QuLab Infinite
        ## Autonomous Scientific Research Agent API

        **ECH0-PRIME** (Kimi-K2 1046B) connected to **QuLab Infinite** with 1,532+ scientific tools across 220+ laboratories.

        This API enables external LLMs to:
        - Query ECH0-PRIME for scientific insights
        - Run experiments on QuLab Infinite
        - Navigate available labs and tools
        - Start autonomous research sessions
        """)

        if ech0_error:
            gr.Markdown(f"⚠️ **Warning**: ECH0-PRIME initialization failed: {ech0_error}")

        with gr.Tabs():

            # Overview Tab
            with gr.TabItem("🏠 Overview"):
                with gr.Row():
                    with gr.Column():
                        gr.Markdown("""
                        ## 📊 System Capabilities
                        - **LLM**: Kimi-K2 1046B (Moonshot AI)
                        - **Labs**: 220+ specialized laboratories
                        - **Tools**: 1,532+ scientific tools
                        - **Domains**: Quantum, Biology, Chemistry, Physics, Medical, Environmental

                        ## 🔗 API Endpoints
                        - `/query` - Ask scientific questions
                        - `/run_experiment` - Execute experiments
                        - `/list_labs` - Browse available labs
                        - `/research` - Start research session
                        - `/history` - View conversation history
                        - `/experiments` - View experiment log
                        """)

                    with gr.Column():
                        status_indicator = "🟢 **Fully Operational**" if ech0_prime else "🟡 **Limited Functionality**"
                        gr.Markdown(f"### System Status\n{status_indicator}")

                        # Lab list button
                        list_labs_btn = gr.Button("🔬 List Available Labs", variant="secondary")
                        labs_display = gr.Textbox(
                            label="Available Labs",
                            lines=15,
                            interactive=False
                        )

                        list_labs_btn.click(
                            fn=list_available_labs,
                            inputs=[],
                            outputs=labs_display,
                            api_name="list_labs"
                        )

            # Query Tab
            with gr.TabItem("💬 Query ECH0-PRIME"):
                with gr.Row():
                    with gr.Column():
                        query_input = gr.Textbox(
                            label="🔍 Scientific Question",
                            placeholder="e.g., 'Design an experiment to test quantum teleportation fidelity at 100km'",
                            lines=3
                        )
                        include_context = gr.Checkbox(
                            label="Include conversation context",
                            value=True
                        )
                        query_btn = gr.Button("🤖 Ask ECH0-PRIME", variant="primary", size="lg")

                    with gr.Column():
                        query_output = gr.Textbox(
                            label="📝 Response",
                            lines=20,
                            interactive=False
                        )

                query_btn.click(
                    fn=query_ech0_prime,
                    inputs=[query_input, include_context],
                    outputs=query_output,
                    api_name="query"
                )

                gr.Markdown("""
                ### 💡 Usage Tips
                - Ask specific scientific questions
                - Request experiment designs
                - Get analysis of hypothetical scenarios
                - ECH0-PRIME has full context of all previous queries in the session
                """)

            # Experiment Execution Tab
            with gr.TabItem("🧪 Run Experiments"):
                with gr.Row():
                    with gr.Column():
                        lab_name = gr.Textbox(
                            label="🔬 Lab Name",
                            placeholder="e.g., quantum_lab, oncology_lab",
                            lines=1
                        )
                        exp_type = gr.Textbox(
                            label="🧬 Experiment Type",
                            placeholder="e.g., teleportation_fidelity, tumor_evolution",
                            lines=1
                        )
                        exp_params = gr.Textbox(
                            label="⚙️ Parameters (JSON)",
                            placeholder='{"temperature": 25, "concentration": 0.1}',
                            lines=5
                        )
                        run_exp_btn = gr.Button("🚀 Run Experiment", variant="primary", size="lg")

                    with gr.Column():
                        exp_output = gr.Textbox(
                            label="📊 Experiment Results",
                            lines=20,
                            interactive=False
                        )

                run_exp_btn.click(
                    fn=run_single_experiment,
                    inputs=[lab_name, exp_type, exp_params],
                    outputs=exp_output,
                    api_name="run_experiment"
                )

                gr.Markdown("""
                ### 📚 Example Experiments

                **Quantum Lab**:
                ```json
                Lab: quantum_lab
                Type: teleportation_fidelity
                Parameters: {"protocol": "GHZ", "num_qubits": 5, "distance_km": 100, "error_rate": 0.001}
                ```

                **Oncology Lab**:
                ```json
                Lab: oncology_lab
                Type: tumor_evolution
                Parameters: {"tumor_type": "NSCLC", "initial_size_mm": 15, "treatment": "pembrolizumab"}
                ```
                """)

            # Research Session Tab
            with gr.TabItem("🔬 Research Sessions"):
                with gr.Row():
                    with gr.Column():
                        research_q = gr.Textbox(
                            label="🎯 Research Question",
                            placeholder="e.g., 'What drug combinations are most effective for treating NSCLC with EGFR mutations?'",
                            lines=3
                        )
                        max_iter = gr.Slider(
                            minimum=1,
                            maximum=10,
                            value=5,
                            step=1,
                            label="Max Iterations"
                        )
                        research_btn = gr.Button("🔬 Start Research Session", variant="primary", size="lg")

                    with gr.Column():
                        research_output = gr.Textbox(
                            label="📊 Research Results",
                            lines=25,
                            interactive=False
                        )

                research_btn.click(
                    fn=start_research_session,
                    inputs=[research_q, max_iter],
                    outputs=research_output,
                    api_name="research"
                )

                gr.Markdown("""
                ### 🔄 Autonomous Research Loop

                ECH0-PRIME will:
                1. Design experiments based on your question
                2. Run them on QuLab Infinite
                3. Analyze results
                4. Iterate and refine
                5. Provide final conclusions

                This can take several minutes depending on complexity.
                """)

            # Session Management Tab
            with gr.TabItem("📋 Session Management"):
                with gr.Row():
                    with gr.Column():
                        gr.Markdown("### 💬 Conversation History")
                        history_btn = gr.Button("📜 View History", variant="secondary")
                        history_output = gr.Textbox(
                            label="Conversation History",
                            lines=15,
                            interactive=False
                        )

                        history_btn.click(
                            fn=get_conversation_history,
                            inputs=[],
                            outputs=history_output,
                            api_name="history"
                        )

                    with gr.Column():
                        gr.Markdown("### 📊 Experiment Log")
                        log_btn = gr.Button("📋 View Log", variant="secondary")
                        log_output = gr.Textbox(
                            label="Experiment Log",
                            lines=15,
                            interactive=False
                        )

                        log_btn.click(
                            fn=get_experiment_log,
                            inputs=[],
                            outputs=log_output,
                            api_name="experiments"
                        )

                with gr.Row():
                    reset_btn = gr.Button("🔄 Reset Session", variant="stop")
                    reset_output = gr.Textbox(
                        label="Reset Status",
                        lines=3,
                        interactive=False
                    )

                reset_btn.click(
                    fn=reset_session,
                    inputs=[],
                    outputs=reset_output,
                    api_name="reset"
                )

    return interface


# Launch the interface
if __name__ == "__main__":
    interface = create_interface()
    interface.launch(
        server_name="0.0.0.0",
        server_port=7861,  # Different port from main QuLab GUI
        show_api=True,     # Enable Gradio API for external LLMs
        share=False
    )