import gradio as gr
import torch
from unsloth import FastLanguageModel
import time

# Global variables for model
model = None
tokenizer = None

def load_model():
    """Load the Ideon Audio support model"""
    global model, tokenizer
    
    if model is None:
        print("🔄 Loading Ideon Audio Support Model...")
        
        try:
            model, tokenizer = FastLanguageModel.from_pretrained(
                model_name="ainovatronsec/ideoaudio",
                max_seq_length=2048,
                dtype=None,
                load_in_4bit=True,
                trust_remote_code=False,
            )
            
            # Enable fast inference
            FastLanguageModel.for_inference(model)
            print("✅ Model loaded successfully!")
            
        except Exception as e:
            print(f"❌ Error loading model: {e}")
            return False
    
    return True

def get_ideon_response(message, history, temperature=0.7, max_tokens=256):
    """Generate response from Ideon Audio support model"""
    
    # Load model if not already loaded
    if not load_model():
        return "❌ Sorry, the model failed to load. Please try again later."
    
    try:
        # Format message for the model
        messages = [{"role": "user", "content": message}]
        
        # Apply chat template
        inputs = tokenizer.apply_chat_template(
            messages,
            tokenize=True,
            add_generation_prompt=True,
            return_tensors="pt"
        )
        
        # Move to GPU if available
        if torch.cuda.is_available():
            inputs = inputs.to("cuda")
        
        # Generate response
        with torch.no_grad():
            outputs = model.generate(
                input_ids=inputs,
                max_new_tokens=max_tokens,
                temperature=temperature,
                top_p=0.9,
                do_sample=True,
                use_cache=True,
                pad_token_id=tokenizer.eos_token_id
            )
        
        # Decode the response
        response = tokenizer.decode(outputs[0][inputs.shape[-1]:], skip_special_tokens=True)
        
        # Clean up the response
        response = response.strip()
        if not response:
            response = "I apologize, but I couldn't generate a proper response. Could you please rephrase your question about Ideon Audio products?"
        
        return response
        
    except Exception as e:
        print(f"Error generating response: {e}")
        return f"❌ I encountered an error while processing your question. Please try again. Error details: {str(e)}"

def create_interface():
    """Create the Gradio interface"""
    
    # Custom CSS for better styling
    custom_css = """
    .gradio-container {
        max-width: 1000px !important;
    }
    .header {
        text-align: center;
        background: linear-gradient(90deg, #1e3a8a, #3b82f6);
        color: white;
        padding: 20px;
        border-radius: 10px;
        margin-bottom: 20px;
    }
    .examples-row {
        margin-top: 20px;
    }
    """
    
    # Create the chat interface
    with gr.Blocks(css=custom_css, title="Ideon Audio Support Assistant") as demo:
        
        # Header
        gr.HTML("""
        <div class="header">
            <h1>🎵 Ideon Audio Technical Support Assistant</h1>
            <p>Expert knowledge for high-end audio equipment</p>
        </div>
        """)
        
        # Description
        gr.Markdown("""
        Welcome to the **Ideon Audio Technical Support Assistant**! I'm here to help you with:
        
        🔹 **Product Information** - Specifications, features, and capabilities  
        🔹 **Setup & Installation** - Connection procedures and configuration  
        🔹 **Troubleshooting** - Diagnosing and resolving technical issues  
        🔹 **Warranty Support** - Coverage details and service procedures  
        🔹 **Technical Questions** - Performance metrics and compatibility  
        
        **Supported Products**: Absolute ε DAC, ΙΩΝ DAC, eos DAC, Absolute Stream, USB Re-clockers, and more!
        """)
        
        # Chat interface with advanced settings
        with gr.Row():
            with gr.Column(scale=3):
                chatbot = gr.ChatInterface(
                    fn=get_ideon_response,
                    chatbot=gr.Chatbot(
                        height=500,
                        placeholder="Ask me anything about Ideon Audio products...",
                        avatar_images=(None, "🎵")
                    ),
                    textbox=gr.Textbox(
                        placeholder="Type your question about Ideon Audio products here...",
                        container=False,
                        scale=7
                    ),
                    submit_btn="Ask Question",
                    retry_btn="🔄 Retry",
                    undo_btn="↩️ Undo",
                    clear_btn="🗑️ Clear Chat",
                )
            
            with gr.Column(scale=1):
                gr.Markdown("### ⚙️ Settings")
                
                temperature = gr.Slider(
                    minimum=0.1,
                    maximum=1.0,
                    value=0.7,
                    step=0.1,
                    label="Response Creativity",
                    info="Higher = more creative, Lower = more focused"
                )
                
                max_tokens = gr.Slider(
                    minimum=50,
                    maximum=500,
                    value=256,
                    step=50,
                    label="Response Length",
                    info="Maximum tokens in response"
                )
                
                # Update the chat function with new parameters
                def update_chat_fn(message, history):
                    return get_ideon_response(message, history, temperature.value, max_tokens.value)
                
                chatbot.fn = update_chat_fn
        
        # Example questions
        gr.Markdown("### 💡 Example Questions")
        
        example_questions = [
            "What is the recommended burn-in period for the Absolute ε DAC?",
            "How do I connect the Absolute DAC to my audio system?",
            "My DAC won't lock onto the digital signal. What should I do?",
            "What are the main technical specifications of the Absolute ε DAC?",
            "What's covered under the Ideon Audio warranty?",
            "How do I navigate to the General Settings screen?",
            "What digital filters are available on the DAC?",
            "What should I do if my DAC malfunctions?",
            "Can I connect the DAC directly to a power amplifier?",
            "What are the dimensions and weight of the Absolute E?"
        ]
        
        with gr.Row():
            for i in range(0, len(example_questions), 2):
                with gr.Column():
                    if i < len(example_questions):
                        gr.Examples(
                            examples=[[example_questions[i]]],
                            inputs=chatbot.textbox,
                            label=None
                        )
                    if i + 1 < len(example_questions):
                        gr.Examples(
                            examples=[[example_questions[i + 1]]],
                            inputs=chatbot.textbox,
                            label=None
                        )
        
        # Footer
        gr.Markdown("""
        ---
        
        **About this Assistant**: This AI model has been fine-tuned specifically on Ideon Audio product documentation 
        to provide expert technical support. It covers the complete product line including DACs, streamers, 
        USB re-clockers, and network optimizers.
        
        **Note**: For warranty claims or complex technical issues, please contact your authorized Ideon Audio dealer 
        or email info@ideonaudio.com directly.
        
        *Powered by Meta Llama 3.1 8B + Unsloth fine-tuning*
        """)
    
    return demo

# Create and launch the interface
if __name__ == "__main__":
    print("🚀 Starting Ideon Audio Support Assistant...")
    
    # Pre-load the model (optional, for faster first response)
    print("🔄 Pre-loading model for faster responses...")
    load_model()
    
    # Create and launch the interface
    demo = create_interface()
    demo.launch(
        share=False,
        server_name="0.0.0.0",
        server_port=7860,
        show_error=True
    )