# app.py - Mobile-First Implementation
import gradio as gr
import uuid
import logging
import traceback
from typing import Optional, Tuple, List, Dict, Any
import os

# Configure comprehensive logging
logging.basicConfig(
    level=logging.INFO,
    format='%(asctime)s - %(name)s - %(levelname)s - %(message)s',
    handlers=[
        logging.StreamHandler(),
        logging.FileHandler('app.log')
    ]
)
logger = logging.getLogger(__name__)

# Try to import orchestration components
orchestrator = None
orchestrator_available = False

try:
    logger.info("Attempting to import orchestration components...")
    import sys
    sys.path.insert(0, '.')
    sys.path.insert(0, 'src')
    
    from src.agents.intent_agent import create_intent_agent
    from src.agents.synthesis_agent import create_synthesis_agent
    from src.agents.safety_agent import create_safety_agent
    from llm_router import LLMRouter
    from orchestrator_engine import MVPOrchestrator
    from context_manager import EfficientContextManager
    from config import settings
    
    logger.info("✓ Successfully imported orchestration components")
    orchestrator_available = True
except ImportError as e:
    logger.warning(f"Could not import orchestration components: {e}")
    logger.info("Will use placeholder mode")

try:
    from spaces import GPU
    SPACES_GPU_AVAILABLE = True
    logger.info("HF Spaces GPU available")
except ImportError:
    # Not running on HF Spaces or spaces module not available
    SPACES_GPU_AVAILABLE = False
    GPU = None
    logger.info("Running without HF Spaces GPU")

def create_mobile_optimized_interface():
    """Create the mobile-optimized Gradio interface and return demo with components"""
    
    # Store components for wiring
    interface_components = {}
    
    with gr.Blocks(
        title="AI Research Assistant MVP",
        theme=gr.themes.Soft(
            primary_hue="blue",
            secondary_hue="gray",
            font=("Inter", "system-ui", "sans-serif")
        ),
        css="""
        /* Mobile-first responsive CSS */
        .mobile-container {
            max-width: 100vw;
            margin: 0 auto;
            padding: 0 12px;
        }
        
        /* Touch-friendly button sizing */
        .gradio-button {
            min-height: 44px !important;
            min-width: 44px !important;
            font-size: 16px !important; /* Prevents zoom on iOS */
        }
        
        /* Mobile-optimized chat interface */
        .chatbot-container {
            height: 60vh !important;
            max-height: 60vh !important;
            overflow-y: auto !important;
            -webkit-overflow-scrolling: touch !important;
        }
        
        /* Mobile input enhancements */
        .textbox-input {
            font-size: 16px !important; /* Prevents zoom */
            min-height: 44px !important;
            padding: 12px !important;
        }
        
        /* Responsive grid adjustments */
        @media (max-width: 768px) {
            .gradio-row {
                flex-direction: column !important;
                gap: 8px !important;
            }
            
            .gradio-column {
                width: 100% !important;
            }
            
            .chatbot-container {
                height: 50vh !important;
            }
        }
        
        /* Dark mode support */
        @media (prefers-color-scheme: dark) {
            body {
                background: #1a1a1a;
                color: #ffffff;
            }
        }
        
        /* Hide scrollbars but maintain functionality */
        .chatbot-container::-webkit-scrollbar {
            width: 4px;
        }
        
        /* Loading states */
        .loading-indicator {
            display: flex;
            align-items: center;
            justify-content: center;
            padding: 20px;
        }
        
        /* Mobile menu enhancements */
        .accordion-content {
            max-height: 200px !important;
            overflow-y: auto !important;
        }
        """
    ) as demo:
        
        # Session Management (Mobile-Optimized)
        with gr.Column(elem_classes="mobile-container"):
            gr.Markdown("""
            # 🧠 Research Assistant
            *Academic AI with transparent reasoning*
            """)
            
            # Session Header Bar (Mobile-Friendly)
            with gr.Row():
                session_info = gr.Textbox(
                    label="Session ID",
                    value=str(uuid.uuid4())[:8],  # Shortened for mobile
                    max_lines=1,
                    show_label=False,
                    container=False,
                    scale=3
                )
                
                new_session_btn = gr.Button(
                    "🔄 New",
                    size="sm",
                    variant="secondary",
                    scale=1,
                    min_width=60
                )
                
                menu_toggle = gr.Button(
                    "⚙️",
                    size="sm",
                    variant="secondary",
                    scale=1,
                    min_width=60
                )
            
            # Main Chat Area (Mobile-Optimized)
            with gr.Tabs() as main_tabs:
                with gr.TabItem("💬 Chat", id="chat_tab"):
                    chatbot = gr.Chatbot(
                        label="",
                        show_label=False,
                        height="60vh",
                        elem_classes="chatbot-container",
                        type="messages"
                    )
                    interface_components['chatbot'] = chatbot
                    
                    # Mobile Input Area
                    with gr.Row():
                        message_input = gr.Textbox(
                            placeholder="Ask me anything...",
                            show_label=False,
                            max_lines=3,
                            container=False,
                            scale=4,
                            autofocus=True
                        )
                        interface_components['message_input'] = message_input
                        
                        send_btn = gr.Button(
                            "↑ Send",
                            variant="primary",
                            scale=1,
                            min_width=80
                        )
                        interface_components['send_btn'] = send_btn
                
                # Technical Details Tab (Collapsible for Mobile)
                with gr.TabItem("🔍 Details", id="details_tab"):
                    with gr.Accordion("Reasoning Chain", open=False):
                        reasoning_display = gr.JSON(
                            label="",
                            show_label=False
                        )
                    
                    with gr.Accordion("Agent Performance", open=False):
                        performance_display = gr.JSON(
                            label="",
                            show_label=False
                        )
                    
                    with gr.Accordion("Session Context", open=False):
                        context_display = gr.JSON(
                            label="",
                            show_label=False
                        )
            
            # Mobile Bottom Navigation
            with gr.Row(visible=False, elem_id="mobile_nav") as mobile_navigation:
                chat_nav_btn = gr.Button("💬 Chat", variant="secondary", size="sm", min_width=0)
                details_nav_btn = gr.Button("🔍 Details", variant="secondary", size="sm", min_width=0)
                settings_nav_btn = gr.Button("⚙️ Settings", variant="secondary", size="sm", min_width=0)
        
        # Settings Panel (Modal for Mobile)
        with gr.Column(visible=False, elem_id="settings_panel") as settings:
            with gr.Accordion("Display Options", open=True):
                show_reasoning = gr.Checkbox(
                    label="Show reasoning chain",
                    value=True,
                    info="Display step-by-step reasoning"
                )
                
                show_agent_trace = gr.Checkbox(
                    label="Show agent execution trace",
                    value=False,
                    info="Display which agents processed your request"
                )
                
                compact_mode = gr.Checkbox(
                    label="Compact mode",
                    value=False,
                    info="Optimize for smaller screens"
                )
            
            with gr.Accordion("Performance Options", open=False):
                response_speed = gr.Radio(
                    choices=["Fast", "Balanced", "Thorough"],
                    value="Balanced",
                    label="Response Speed Preference"
                )
                
                cache_enabled = gr.Checkbox(
                    label="Enable context caching",
                    value=True,
                    info="Faster responses using session memory"
                )
            
            gr.Button("Save Preferences", variant="primary")
        
        # Wire up the submit handler INSIDE the gr.Blocks context
        if 'send_btn' in interface_components and 'message_input' in interface_components and 'chatbot' in interface_components:
            # Connect the submit handler with the GPU-decorated function
            interface_components['send_btn'].click(
                fn=chat_handler_fn,
                inputs=[interface_components['message_input'], interface_components['chatbot']],
                outputs=[interface_components['chatbot'], interface_components['message_input']]
            )
    
    return demo, interface_components

def setup_event_handlers(demo, event_handlers):
    """Setup event handlers for the interface"""
    
    # Find components by their labels or types
    components = {}
    for block in demo.blocks:
        if hasattr(block, 'label'):
            if block.label == 'Session ID':
                components['session_info'] = block
            elif hasattr(block, 'value') and 'session' in str(block.value).lower():
                components['session_id'] = block
    
    # Setup message submission handler
    try:
        # This is a simplified version - you'll need to adapt based on your actual component structure
        if hasattr(demo, 'submit'):
            demo.submit(
                fn=event_handlers.handle_message_submit,
                inputs=[components.get('message_input'), components.get('chatbot')],
                outputs=[components.get('message_input'), components.get('chatbot')]
            )
    except Exception as e:
        print(f"Could not setup event handlers: {e}")
        # Fallback to basic functionality
    
    return demo

async def process_message_async(message: str, history: Optional[List], session_id: str) -> Tuple[List, str]:
    """
    Process message with full orchestration system
    Returns (updated_history, empty_string)
    """
    global orchestrator
    
    try:
        logger.info(f"Processing message: {message[:100]}")
        logger.info(f"Session ID: {session_id}")
        
        if not message or not message.strip():
            logger.debug("Empty message received")
            return history if history else [], ""
        
        if history is None:
            history = []
        
        new_history = list(history) if isinstance(history, list) else []
        
        # Add user message
        new_history.append({"role": "user", "content": message.strip()})
        
        # Try to use orchestrator if available
        if orchestrator is not None:
            try:
                logger.info("Attempting full orchestration...")
                # Use orchestrator to process
                result = await orchestrator.process_request(
                    session_id=session_id,
                    user_input=message.strip()
                )
                
                # Extract response from result
                response = result.get('response', result.get('final_response', str(result)))
                logger.info(f"Orchestrator returned response: {response[:100]}")
                
            except Exception as orch_error:
                logger.error(f"Orchestrator error: {orch_error}", exc_info=True)
                response = f"[Orchestrator Error] {str(orch_error)}"
        else:
            # Fallback placeholder
            logger.info("Using placeholder response")
            response = f"I received your message: {message}\n\nThis is a placeholder response. The orchestrator system is {'' if orchestrator_available else 'not'} available."
        
        # Add assistant response
        new_history.append({"role": "assistant", "content": response})
        logger.info("Message processing complete")
        
        return new_history, ""
        
    except Exception as e:
        logger.error(f"Error in process_message_async: {e}", exc_info=True)
        error_history = list(history) if history else []
        error_history.append({"role": "user", "content": message})
        error_history.append({"role": "assistant", "content": f"I encountered an error: {str(e)}"})
        return error_history, ""

def process_message(message: str, history: Optional[List]) -> Tuple[List, str]:
    """
    Synchronous wrapper for async processing
    """
    import asyncio
    
    # Generate a session ID
    session_id = str(uuid.uuid4())[:8]
    
    try:
        # Run async processing
        loop = asyncio.new_event_loop()
        asyncio.set_event_loop(loop)
        result = loop.run_until_complete(process_message_async(message, history, session_id))
        return result
    except Exception as e:
        logger.error(f"Error in process_message: {e}", exc_info=True)
        error_history = list(history) if history else []
        error_history.append({"role": "user", "content": message})
        error_history.append({"role": "assistant", "content": f"Error: {str(e)}"})
        return error_history, ""

# Decorate the chat handler with GPU if available
if SPACES_GPU_AVAILABLE and GPU is not None:
    @GPU  # This decorator is detected by HF Spaces for ZeroGPU allocation
    def gpu_chat_handler(message, history):
        """Handle chat messages with GPU support"""
        return process_message(message, history)
    chat_handler_fn = gpu_chat_handler
else:
    chat_handler_fn = process_message

# Initialize orchestrator on module load
def initialize_orchestrator():
    """Initialize the orchestration system with logging"""
    global orchestrator
    
    if not orchestrator_available:
        logger.info("Orchestrator components not available, skipping initialization")
        return
    
    try:
        logger.info("=" * 60)
        logger.info("INITIALIZING ORCHESTRATION SYSTEM")
        logger.info("=" * 60)
        
        # Get HF token
        hf_token = os.getenv('HF_TOKEN', '')
        if not hf_token:
            logger.warning("HF_TOKEN not found in environment")
        
        # Initialize LLM Router
        logger.info("Step 1/6: Initializing LLM Router...")
        llm_router = LLMRouter(hf_token)
        logger.info("✓ LLM Router initialized")
        
        # Initialize Agents
        logger.info("Step 2/6: Initializing Agents...")
        agents = {
            'intent_recognition': create_intent_agent(llm_router),
            'response_synthesis': create_synthesis_agent(llm_router),
            'safety_check': create_safety_agent(llm_router)
        }
        logger.info(f"✓ Initialized {len(agents)} agents")
        
        # Initialize Context Manager
        logger.info("Step 3/6: Initializing Context Manager...")
        context_manager = EfficientContextManager()
        logger.info("✓ Context Manager initialized")
        
        # Initialize Orchestrator
        logger.info("Step 4/6: Initializing Orchestrator...")
        orchestrator = MVPOrchestrator(llm_router, context_manager, agents)
        logger.info("✓ Orchestrator initialized")
        
        logger.info("=" * 60)
        logger.info("ORCHESTRATION SYSTEM READY")
        logger.info("=" * 60)
        
    except Exception as e:
        logger.error(f"Failed to initialize orchestrator: {e}", exc_info=True)
        orchestrator = None

# Try to initialize orchestrator
initialize_orchestrator()

if __name__ == "__main__":
    logger.info("=" * 60)
    logger.info("STARTING APP")
    logger.info("=" * 60)
    
    demo, components = create_mobile_optimized_interface()
    
    logger.info("✓ Interface created")
    logger.info(f"Orchestrator available: {orchestrator is not None}")
    
    # Launch the app
    logger.info("=" * 60)
    logger.info("LAUNCHING GRADIO APP")
    logger.info("=" * 60)
    demo.launch(
        server_name="0.0.0.0",
        server_port=7860,
        share=False
    )