"""
Main Dash application for HalluNox Dashboard with streaming support
"""

import dash
from dash import dcc, html, Input, Output, State, callback, clientside_callback, MATCH, ALL
import dash_bootstrap_components as dbc
import json
from datetime import datetime
import threading
import time
import uuid

from ..detector import HallucinationDetector
from .utils import ConversationManager, format_confidence_message, generate_conversation_title


def create_app(debug=False):
    """Create and configure the Dash app."""
    
    # Initialize Dash app with Bootstrap theme
    app = dash.Dash(
        __name__,
        external_stylesheets=[dbc.themes.BOOTSTRAP, dbc.icons.FONT_AWESOME],
        suppress_callback_exceptions=True,
        title="Hallunox Dashboard"
    )
    
    # Initialize conversation manager
    conv_manager = ConversationManager()
    
    # Initialize HalluNox detector (will be loaded on first use)
    detector = None
    detector_lock = threading.Lock()
    
    # Streaming state management
    streaming_state = {
        'is_streaming': False,
        'current_conv_id': None,
        'should_stop': False,
        'current_response': '',
        'message_id': None
    }
    streaming_lock = threading.Lock()
    
    def get_detector():
        """Lazy load the detector."""
        nonlocal detector
        with detector_lock:
            if detector is None:
                print("Loading HalluNox detector...")
                detector = HallucinationDetector(
                    llm_model_id="unsloth/Llama-3.2-3B-Instruct",  # Default Llama
                    enable_response_generation=True,
                    enable_inference=True,
                    confidence_threshold=0.6,
                    verbose=True
                )
                print("HalluNox detector loaded successfully!")
            return detector
    
    def stop_streaming():
        """Stop current streaming."""
        with streaming_lock:
            streaming_state['should_stop'] = True
    
    # App layout
    app.layout = dbc.Container([
        dcc.Store(id='current-conversation-id', data=None),
        dcc.Store(id='conversation-list-store', data=[]),
        dcc.Store(id='streaming-state-store', data={'is_streaming': False, 'message_id': None}),
        dcc.Store(id='model-settings-store', data={
            'model_type': 'llama',
            'confidence_threshold': 0.6,
            'temperature': 0.7,
            'max_tokens': 512
        }),
        dcc.Interval(id='refresh-interval', interval=1000, n_intervals=0),  # Refresh every 1s for streaming
        dcc.Interval(id='streaming-interval', interval=200, n_intervals=0, disabled=True),  # Fast refresh for streaming
        
        # Header
        dbc.Row([
            dbc.Col([
                html.H1("🤖 Hallunox", className="text-light mb-0"),
                html.P("Confidence-Aware AI Assistant", className="text-muted small mb-0")
            ], width=4),
            dbc.Col([
                dbc.ButtonGroup([
                    dbc.Button("🔄 New Chat", id="new-chat-btn", color="primary", size="sm"),
                    dbc.Button("⚙️ Settings", id="settings-btn", color="secondary", size="sm"),
                ], className="float-end")
            ], width=8)
        ], className="bg-dark p-3 mb-0"),
        
        # Main content area
        dbc.Row([
            # Sidebar for conversations
            dbc.Col([
                dbc.Card([
                    dbc.CardHeader([
                        html.H6("💬 Conversations", className="mb-0"),
                        dbc.Button("🔽", id="sidebar-toggle", size="sm", color="link", className="text-muted float-end")
                    ]),
                    dbc.Collapse([
                        dbc.CardBody([
                            html.Div(id="conversation-list", children=[])
                        ], style={"max-height": "400px", "overflow-y": "auto"})
                    ], id="sidebar-collapse", is_open=True)
                ], className="mb-3")
            ], width=3, id="sidebar-col"),
            
            # Chat area
            dbc.Col([
                # Chat messages area
                dbc.Card([
                    dbc.CardBody([
                        html.Div(id="chat-messages", 
                                children=[
                                    html.Div([
                                        html.P("👋 Welcome to Hallunox! I'm your confidence-aware AI assistant.", 
                                              className="text-muted text-center")
                                    ], className="text-center p-4")
                                ],
                                style={"height": "500px", "overflow-y": "auto", "padding": "10px"})
                    ], className="p-0")
                ], className="mb-3"),
                
                # Input area
                dbc.Card([
                    dbc.CardBody([
                        dbc.InputGroup([
                            dbc.Textarea(
                                id="message-input",
                                placeholder="Type your message here... (Press Ctrl+Enter to send)",
                                rows=2,
                                style={"resize": "none"}
                            ),
                            dbc.Button("➤", id="send-btn", color="primary", disabled=False),
                            dbc.Button("⏹", id="stop-btn", color="danger", size="sm", style={"display": "none"})
                        ])
                    ])
                ])
            ], width=6),
            
            # Settings panel
            dbc.Col([
                dbc.Card([
                    dbc.CardHeader([
                        html.H6("⚙️ Model Settings", className="mb-0")
                    ]),
                    dbc.CardBody([
                        # Model Selection
                        html.Label("Model Type", className="form-label"),
                        dcc.Dropdown(
                            id="model-type-dropdown",
                            options=[
                                {"label": "🦙 Llama-3.2-3B (Default)", "value": "llama"},
                                {"label": "🏥 MedGemma-4B (Medical)", "value": "medgemma"}
                            ],
                            value="llama",
                            className="mb-3"
                        ),
                        
                        # Confidence Threshold
                        html.Label("Confidence Threshold", className="form-label"),
                        dcc.Slider(
                            id="confidence-threshold-slider",
                            min=0.1,
                            max=0.9,
                            step=0.05,
                            value=0.6,
                            marks={i/10: f"{i/10:.1f}" for i in range(1, 10, 2)},
                            tooltip={"placement": "bottom", "always_visible": True},
                            className="mb-3"
                        ),
                        
                        # Temperature
                        html.Label("Temperature", className="form-label"),
                        dcc.Slider(
                            id="temperature-slider",
                            min=0.1,
                            max=1.0,
                            step=0.1,
                            value=0.7,
                            marks={i/10: f"{i/10:.1f}" for i in range(1, 11, 2)},
                            tooltip={"placement": "bottom", "always_visible": True},
                            className="mb-3"
                        ),
                        
                        # Max Tokens
                        html.Label("Max Tokens", className="form-label"),
                        dcc.Slider(
                            id="max-tokens-slider",
                            min=128,
                            max=2048,
                            step=128,
                            value=512,
                            marks={i: str(i) for i in range(128, 2049, 384)},
                            tooltip={"placement": "bottom", "always_visible": True},
                            className="mb-3"
                        ),
                        
                        html.Hr(),
                        html.Small("💡 Lower confidence threshold = more cautious responses", 
                                 className="text-muted")
                    ])
                ])
            ], width=3)
        ])
    ], fluid=True, className="bg-light min-vh-100")
    
    # Helper function to create message component
    def create_message_component(message, msg_index, is_streaming=False):
        """Create a message component with action buttons."""
        timestamp = datetime.fromisoformat(message['timestamp']).strftime('%H:%M')
        
        if message['role'] == 'user':
            return dbc.Card([
                dbc.CardBody([
                    html.Div([
                        html.Strong("You", className="text-primary"),
                        html.Span(f" • {timestamp}", className="text-muted small float-end"),
                        dbc.ButtonGroup([
                            dbc.Button("🔄", id={"type": "regenerate-user-btn", "index": msg_index}, 
                                     size="sm", color="link", title="Regenerate", className="p-1"),
                            dbc.Button("✏️", id={"type": "edit-user-btn", "index": msg_index}, 
                                     size="sm", color="link", title="Edit", className="p-1"),
                            dbc.Button("🗑️", id={"type": "delete-user-btn", "index": msg_index}, 
                                     size="sm", color="link", title="Delete", className="p-1 text-danger")
                        ], size="sm", className="float-end me-2")
                    ], className="mb-2"),
                    html.P(message['content'], className="mb-0", id={"type": "user-content", "index": msg_index})
                ], className="p-3")
            ], className="mb-3 ms-5", color="primary", outline=True)
        else:
            # Assistant message
            confidence = message.get('confidence_score', 1.0)
            meets_threshold = message.get('metadata', {}).get('meets_threshold', True)
            is_error = message.get('metadata', {}).get('error', False)
            
            # Choose color based on confidence
            if is_error:
                card_color = "danger"
            elif confidence < 0.6:
                card_color = "warning"
            elif confidence < 0.8:
                card_color = "info"
            else:
                card_color = "success"
            
            confidence_badge = dbc.Badge(
                f"Confidence: {confidence:.2f}",
                color=card_color,
                className="float-end"
            ) if not is_error else dbc.Badge("Error", color="danger", className="float-end")
            
            # Content with streaming effect
            content = message['content']
            if is_streaming:
                content += "▍"  # Cursor for streaming effect
            
            return dbc.Card([
                dbc.CardBody([
                    html.Div([
                        html.Strong("🤖 Hallunox", className="text-success"),
                        html.Span(f" • {timestamp}", className="text-muted small"),
                        confidence_badge
                    ], className="mb-2"),
                    html.P(content, className="mb-0", id={"type": "assistant-content", "index": msg_index}),
                    html.Div([
                        dbc.ButtonGroup([
                            dbc.Button("🔄", id={"type": "regenerate-btn", "index": msg_index}, 
                                     size="sm", color="link", title="Regenerate", className="p-1"),
                            dbc.Button("📋", id={"type": "copy-msg-btn", "index": msg_index}, 
                                     size="sm", color="link", title="Copy", className="p-1"),
                            dbc.Button("✏️", id={"type": "edit-msg-btn", "index": msg_index}, 
                                     size="sm", color="link", title="Edit", className="p-1"),
                            dbc.Button("⏹", id={"type": "stop-streaming-btn", "index": msg_index}, 
                                     size="sm", color="danger", title="Stop", className="p-1",
                                     style={"display": "block" if is_streaming else "none"})
                        ], size="sm", className="mt-2")
                    ])
                ], className="p-3")
            ], className="mb-3 me-5", color=card_color, outline=True)
    
    # Callbacks
    
    @app.callback(
        Output('conversation-list', 'children'),
        Input('refresh-interval', 'n_intervals'),
        State('current-conversation-id', 'data')
    )
    def update_conversation_list(n_intervals, current_conv_id):
        """Update the conversation list."""
        conversations = conv_manager.get_conversation_list()
        
        if not conversations:
            return [html.P("No conversations yet", className="text-muted text-center small")]
        
        items = []
        for conv in conversations:
            is_active = conv['id'] == current_conv_id
            card_class = "border-primary" if is_active else "border-light"
            
            items.append(
                dbc.Card([
                    dbc.CardBody([
                        html.H6(conv['title'], className="card-title mb-1", 
                               style={"font-size": "0.9rem"}),
                        html.Small(f"{conv['message_count']} messages", 
                                 className="text-muted"),
                        dbc.ButtonGroup([
                            dbc.Button("📝", id={"type": "edit-conv-btn", "index": conv['id']}, 
                                     size="sm", color="link", className="p-0"),
                            dbc.Button("📋", id={"type": "copy-conv-btn", "index": conv['id']}, 
                                     size="sm", color="link", className="p-0"),
                            dbc.Button("🗑️", id={"type": "delete-conv-btn", "index": conv['id']}, 
                                     size="sm", color="link", className="p-0 text-danger")
                        ], size="sm", className="float-end")
                    ], className="p-2 cursor-pointer", 
                       id={"type": "conv-item", "index": conv['id']})
                ], className=f"mb-2 {card_class}", style={"cursor": "pointer"})
            )
        
        return items
    
    @app.callback(
        Output('current-conversation-id', 'data'),
        Input('new-chat-btn', 'n_clicks'),
        Input({'type': 'conv-item', 'index': ALL}, 'n_clicks'),
        State('current-conversation-id', 'data'),
        prevent_initial_call=True
    )
    def handle_conversation_selection(new_chat_clicks, conv_clicks, current_conv_id):
        """Handle conversation selection and new chat creation."""
        ctx = dash.callback_context
        
        if not ctx.triggered:
            return current_conv_id
        
        trigger_id = ctx.triggered[0]['prop_id'].split('.')[0]
        
        if trigger_id == 'new-chat-btn':
            # Create new conversation
            conv_id = conv_manager.create_conversation()
            return conv_id
        else:
            # Parse the trigger to get conversation ID
            try:
                trigger_data = json.loads(trigger_id)
                if trigger_data['type'] == 'conv-item':
                    return trigger_data['index']
            except:
                pass
        
        return current_conv_id
    
    @app.callback(
        [Output('chat-messages', 'children'),
         Output('streaming-state-store', 'data'),
         Output('streaming-interval', 'disabled')],
        [Input('current-conversation-id', 'data'),
         Input('send-btn', 'n_clicks'),
         Input('message-input', 'n_submit'),
         Input('streaming-interval', 'n_intervals'),
         Input({'type': 'stop-streaming-btn', 'index': ALL}, 'n_clicks')],
        [State('message-input', 'value'),
         State('model-settings-store', 'data'),
         State('streaming-state-store', 'data')],
        prevent_initial_call=True
    )
    def handle_chat_messages(conv_id, send_clicks, input_submit, streaming_intervals, 
                           stop_clicks, message, model_settings, streaming_state_data):
        """Handle chat messages with streaming support."""
        ctx = dash.callback_context
        
        # If no conversation selected, show welcome message
        if not conv_id:
            return ([html.Div([
                html.P("👋 Welcome to Hallunox! Create a new chat to get started.", 
                      className="text-muted text-center")
            ], className="text-center p-4")], 
            {'is_streaming': False, 'message_id': None}, True)
        
        conversation = conv_manager.get_conversation(conv_id)
        if not conversation:
            return ([html.P("Conversation not found", className="text-danger text-center")], 
                   {'is_streaming': False, 'message_id': None}, True)
        
        # Check trigger
        trigger_id = ctx.triggered[0]['prop_id'].split('.')[0] if ctx.triggered else None
        
        # Handle stop streaming
        if any(stop_clicks) and any(c for c in stop_clicks if c):
            stop_streaming()
            return dash.no_update, {'is_streaming': False, 'message_id': None}, True
        
        # Check if this is a new message
        is_new_message = trigger_id in ['send-btn', 'message-input'] and message and message.strip()
        
        # Handle new user message
        if is_new_message:
            # Add user message immediately
            user_message = {
                "role": "user",
                "content": message.strip(),
                "timestamp": datetime.now().isoformat()
            }
            conv_manager.add_message(conv_id, user_message)
            
            # Update conversation title if this is the first message
            if len(conversation['messages']) == 0:
                title = generate_conversation_title(message.strip())
                conv_manager.update_conversation_title(conv_id, title)
            
            # Start streaming response
            with streaming_lock:
                streaming_state['is_streaming'] = True
                streaming_state['current_conv_id'] = conv_id
                streaming_state['should_stop'] = False
                streaming_state['current_response'] = ''
                streaming_state['message_id'] = str(uuid.uuid4())
            
            # Start AI response generation in background thread
            def generate_response():
                try:
                    detector_instance = get_detector()
                    
                    # Update detector settings
                    if model_settings.get('model_type') == 'medgemma':
                        # TODO: Switch to MedGemma model if needed
                        pass
                    
                    detector_instance.confidence_threshold = model_settings.get('confidence_threshold', 0.6)
                    
                    # Generate response with confidence checking
                    result = detector_instance.generate_response(
                        message.strip(),
                        max_length=model_settings.get('max_tokens', 512),
                        check_confidence=True,
                        force_generate=False
                    )
                    
                    if isinstance(result, dict):
                        confidence_score = result.get('confidence_score', 0.0)
                        response_text = result.get('response', '')
                        
                        if not result.get('should_generate', False):
                            # Low confidence response
                            response_text = format_confidence_message(
                                confidence_score, 
                                model_settings.get('confidence_threshold', 0.6)
                            )
                        
                        ai_message = {
                            "role": "assistant",
                            "content": response_text,
                            "timestamp": datetime.now().isoformat(),
                            "confidence_score": confidence_score,
                            "metadata": {
                                "model_type": model_settings.get('model_type', 'llama'),
                                "threshold": model_settings.get('confidence_threshold', 0.6),
                                "meets_threshold": result.get('meets_threshold', False)
                            }
                        }
                    else:
                        # Simple string response
                        ai_message = {
                            "role": "assistant",
                            "content": str(result),
                            "timestamp": datetime.now().isoformat(),
                            "confidence_score": 1.0,
                            "metadata": {
                                "model_type": model_settings.get('model_type', 'llama'),
                                "threshold": model_settings.get('confidence_threshold', 0.6),
                                "meets_threshold": True
                            }
                        }
                    
                    # Simulate streaming by updating response gradually
                    full_response = ai_message['content']
                    words = full_response.split()
                    
                    for i in range(len(words)):
                        if streaming_state['should_stop']:
                            break
                        
                        partial_response = ' '.join(words[:i+1])
                        with streaming_lock:
                            streaming_state['current_response'] = partial_response
                        
                        time.sleep(0.1)  # Simulate typing delay
                    
                    # Finalize message
                    ai_message['content'] = streaming_state['current_response']
                    conv_manager.add_message(conv_id, ai_message)
                    
                    with streaming_lock:
                        streaming_state['is_streaming'] = False
                        streaming_state['current_response'] = ''
                        streaming_state['message_id'] = None
                
                except Exception as e:
                    # Error handling
                    error_message = {
                        "role": "assistant",
                        "content": f"Sorry, I encountered an error: {str(e)}",
                        "timestamp": datetime.now().isoformat(),
                        "confidence_score": 0.0,
                        "metadata": {"error": True}
                    }
                    conv_manager.add_message(conv_id, error_message)
                    
                    with streaming_lock:
                        streaming_state['is_streaming'] = False
                        streaming_state['current_response'] = ''
                        streaming_state['message_id'] = None
            
            # Start generation thread
            thread = threading.Thread(target=generate_response)
            thread.daemon = True
            thread.start()
            
            return dash.no_update, {'is_streaming': True, 'message_id': streaming_state['message_id']}, False
        
        # Handle streaming updates
        if trigger_id == 'streaming-interval' and streaming_state['is_streaming']:
            # Continue streaming without changing the data
            return dash.no_update, streaming_state_data, False
        
        # Render messages
        conversation = conv_manager.get_conversation(conv_id)  # Refresh conversation
        messages = []
        
        for i, msg in enumerate(conversation['messages']):
            messages.append(create_message_component(msg, i))
        
        # Add streaming message if currently streaming
        if streaming_state['is_streaming'] and streaming_state['current_response']:
            streaming_msg = {
                "role": "assistant",
                "content": streaming_state['current_response'],
                "timestamp": datetime.now().isoformat(),
                "confidence_score": 0.0,
                "metadata": {}
            }
            messages.append(create_message_component(streaming_msg, len(messages), is_streaming=True))
        
        return messages, streaming_state_data, not streaming_state['is_streaming']
    
    @app.callback(
        Output('message-input', 'value'),
        [Input('send-btn', 'n_clicks'),
         Input('message-input', 'n_submit')],
        State('message-input', 'value'),
        prevent_initial_call=True
    )
    def clear_input(send_clicks, input_submit, message):
        """Clear input after sending message."""
        if message and message.strip():
            return ""
        return message
    
    @app.callback(
        Output('model-settings-store', 'data'),
        [Input('model-type-dropdown', 'value'),
         Input('confidence-threshold-slider', 'value'),
         Input('temperature-slider', 'value'),
         Input('max-tokens-slider', 'value')],
        State('model-settings-store', 'data'),
        prevent_initial_call=True
    )
    def update_model_settings(model_type, confidence_threshold, temperature, max_tokens, current_settings):
        """Update model settings."""
        return {
            'model_type': model_type,
            'confidence_threshold': confidence_threshold,
            'temperature': temperature,
            'max_tokens': max_tokens
        }
    
    @app.callback(
        Output('sidebar-collapse', 'is_open'),
        Input('sidebar-toggle', 'n_clicks'),
        State('sidebar-collapse', 'is_open'),
        prevent_initial_call=True
    )
    def toggle_sidebar(n_clicks, is_open):
        """Toggle sidebar collapse."""
        return not is_open
    
    # Client-side callback for Enter key handling
    app.clientside_callback(
        """
        function(id) {
            document.addEventListener('keydown', function(event) {
                if (event.ctrlKey && event.key === 'Enter') {
                    const sendBtn = document.getElementById('send-btn');
                    if (sendBtn && !sendBtn.disabled) {
                        sendBtn.click();
                    }
                }
            });
            return window.dash_clientside.no_update;
        }
        """,
        Output('message-input', 'id'),
        Input('message-input', 'id')
    )
    
    return app


def run_dashboard(host='127.0.0.1', port=8050, debug=False):
    """Run the dashboard server."""
    app = create_app(debug=debug)
    print(f"🚀 Starting Hallunox Dashboard at http://{host}:{port}")
    app.run(host=host, port=port, debug=debug)


if __name__ == '__main__':
    run_dashboard(debug=True)