#!/usr/bin/env python3
"""
Enhanced MiniMax CLI - With Advanced Terminal Display Widgets
"""

import os
import sys
import argparse
import json
import subprocess
from dotenv import load_dotenv
from pathlib import Path

# Load environment variables
load_dotenv()

# Add the project root to Python path for imports
project_root = Path(__file__).parent
sys.path.insert(0, str(project_root))

def main():
    parser = argparse.ArgumentParser(
        description="MiniMax AI - Enhanced CLI with Widget Display",
        formatter_class=argparse.RawDescriptionHelpFormatter,
        epilog="""
Examples:
  python simple_minimax_enhanced.py "What is Python?"
  python simple_minimax_enhanced.py "Write a function" --ui widgets
  python simple_minimax_enhanced.py --ui widgets --stream
  python simple_minimax_enhanced.py --help
        """
    )
    
    parser.add_argument(
        "message", 
        nargs="?",
        help="Message to send to MiniMax AI"
    )
    
    parser.add_argument(
        "--model", 
        default="MiniMaxAI/MiniMax-M1-80k",
        help="Model to use (default: MiniMaxAI/MiniMax-M1-80k)"
    )
    
    parser.add_argument(
        "--stream", 
        action="store_true",
        help="Enable streaming output"
    )
    
    parser.add_argument(
        "--ui",
        choices=["basic", "widgets", "rich"],
        default="basic",
        help="UI mode: basic (default), widgets (advanced), rich (colored)"
    )
    
    parser.add_argument(
        "--theme",
        choices=["dark", "light", "neon"],
        default="dark",
        help="Color theme for widget mode"
    )
    
    parser.add_argument(
        "--version", 
        action="version", 
        version="MiniMax Enhanced CLI v2.0.0"
    )
    
    args = parser.parse_args()
    
    # Check if message provided
    if not args.message:
        if args.ui == "widgets":
            run_interactive_widget_mode(args)
        else:
            run_interactive_basic_mode(args)
    else:
        send_message(args.message, args.model, args.stream, args.ui, args.theme)

def run_interactive_basic_mode(args):
    """Run basic interactive mode."""
    print("🤖 MiniMax AI - Simple CLI")
    print("Enter your message (or 'quit' to exit):")
    
    while True:
        try:
            message = input("\n> ").strip()
            if message.lower() in ['quit', 'exit', 'q']:
                print("Goodbye! 👋")
                break
            if message:
                send_message(message, args.model, args.stream, args.ui, args.theme)
        except KeyboardInterrupt:
            print("\nGoodbye! 👋")
            break
        except EOFError:
            break

def run_interactive_widget_mode(args):
    """Run interactive mode with widgets using Node.js."""
    # Create a temporary Node.js script that uses our widgets
    widget_script = """
const { InteractiveWidgetManager } = require('./lib/widgets');
const readline = require('readline');
const chalk = require('chalk');

const manager = new InteractiveWidgetManager({
    theme: process.argv[2] || 'dark',
    layout: 'vertical',
    interactive: true
});

const rl = readline.createInterface({
    input: process.stdin,
    output: process.stdout
});

console.clear();
manager.showBanner();

const chatWidget = manager.createWidget('chat', 'streaming', {
    title: '💬 MiniMax AI Chat',
    showCursor: true
});

const promptWidget = manager.createWidget('prompt', 'toolResult', {
    title: 'Input',
    toolType: 'generic'
});

promptWidget.setResult({ status: 'Ready for your message...' });

manager.refresh();
console.log(chalk.gray('\nPress Tab to navigate, F1 for help, Ctrl+Q to quit'));

function prompt() {
    rl.question(chalk.cyan('You> '), async (message) => {
        if (message.toLowerCase() === 'quit' || message.toLowerCase() === 'exit') {
            console.log(chalk.yellow('\\nGoodbye! 👋'));
            rl.close();
            process.exit(0);
        }
        
        // Update prompt widget
        promptWidget.setResult({ message: message });
        
        // Start streaming response
        chatWidget.clear();
        chatWidget.startStreaming('Thinking...');
        
        // Call Python to get response
        const { spawn } = require('child_process');
        const python = spawn('python3', [
            'simple_minimax_enhanced.py',
            message,
            '--ui', 'json',
            '--stream'
        ]);
        
        let response = '';
        python.stdout.on('data', (data) => {
            const text = data.toString();
            if (text.startsWith('{') && text.includes('"response"')) {
                try {
                    const json = JSON.parse(text);
                    response = json.response;
                    chatWidget.clear();
                    chatWidget.startStreaming(response);
                } catch (e) {
                    response += text;
                }
            }
        });
        
        python.on('close', () => {
            setTimeout(() => {
                console.clear();
                manager.refresh();
                prompt();
            }, 2000);
        });
    });
}

prompt();
"""
    
    # Write temporary script
    temp_script = Path("temp_widget_chat.js")
    temp_script.write_text(widget_script)
    
    try:
        # Run the Node.js widget interface
        subprocess.run(["node", str(temp_script), args.theme])
    finally:
        # Clean up
        if temp_script.exists():
            temp_script.unlink()

def send_message(message: str, model: str, stream: bool, ui_mode: str, theme: str):
    """Send a message to MiniMax AI and display the response."""
    try:
        # Check for API key
        api_key = os.environ.get('HF_TOKEN')
        if not api_key:
            if ui_mode == "json":
                print(json.dumps({"error": "HF_TOKEN not found"}))
            else:
                print("❌ Error: HF_TOKEN not found in environment variables.")
                print("Please set your Hugging Face token in the .env file.")
                print("Get your token from: https://huggingface.co/settings/tokens")
            return

        if ui_mode == "widgets":
            # Use widget display
            display_with_widgets(message, model, stream, theme)
        elif ui_mode == "rich":
            # Use rich display
            display_with_rich(message, model, stream)
        elif ui_mode == "json":
            # Return JSON for widget integration
            display_as_json(message, model, stream)
        else:
            # Use basic display
            display_basic(message, model, stream)
            
    except Exception as e:
        if ui_mode == "json":
            print(json.dumps({"error": str(e)}))
        else:
            print(f"❌ Error: {e}")
            print("\nTroubleshooting:")
            print("1. Make sure your HF_TOKEN is set in .env")
            print("2. Check your internet connection")
            print("3. Verify the model name is correct")

def display_basic(message: str, model: str, stream: bool):
    """Basic display mode - original implementation."""
    print(f"🤖 Sending to {model}...")
    print("💭 Thinking...")

    from huggingface_hub import InferenceClient
    client = InferenceClient(token=os.environ.get('HF_TOKEN'))
    
    messages = [{"role": "user", "content": message}]

    if stream:
        print("\n" + "="*50)
        print("🤖 MiniMax AI Response:")
        print("="*50)

        response_text = ""
        for chunk in client.chat_completion(
            messages=messages,
            model=model,
            stream=True,
            max_tokens=2000
        ):
            if chunk.choices[0].delta.content:
                content = chunk.choices[0].delta.content
                print(content, end="", flush=True)
                response_text += content
        print("\n" + "="*50)
    else:
        response_obj = client.chat_completion(
            messages=messages,
            model=model,
            stream=False,
            max_tokens=2000
        )
        response = response_obj.choices[0].message.content
        
        print("\n" + "="*50)
        print("🤖 MiniMax AI Response:")
        print("="*50)
        print(response)
        print("="*50)

def display_with_widgets(message: str, model: str, stream: bool, theme: str):
    """Display using advanced widgets via Node.js."""
    # Create a Node.js script that uses our widgets
    widget_display_script = f"""
const {{ WidgetManager }} = require('./lib/widgets');

const manager = new WidgetManager({{
    theme: '{theme}',
    layout: 'vertical'
}});

// Show user message
const userWidget = manager.createWidget('user', 'toolResult', {{
    title: '👤 Your Message',
    toolType: 'generic'
}});
userWidget.setResult({{ message: `{message.replace('`', '\\`')}` }});

// Show AI response
const responseWidget = manager.createWidget('response', 'streaming', {{
    title: '🤖 MiniMax AI Response',
    showCursor: true,
    typingSpeed: {20 if stream else 0}
}});

// Get the actual response
const getResponse = require('child_process').execSync;
const result = getResponse('python3 simple_minimax_enhanced.py "{message.replace('"', '\\"')}" --ui json', {{
    encoding: 'utf8'
}});

try {{
    const data = JSON.parse(result);
    if (data.response) {{
        console.clear();
        console.log(manager.render());
        
        if ({str(stream).lower()}) {{
            responseWidget.startStreaming(data.response);
        }} else {{
            responseWidget.addContent(data.response);
            console.log(manager.render());
        }}
    }}
}} catch (e) {{
    console.error('Error:', e);
}}
"""
    
    # Write and execute temporary script
    temp_script = Path("temp_display.js")
    temp_script.write_text(widget_display_script)
    
    try:
        subprocess.run(["node", str(temp_script)])
    finally:
        if temp_script.exists():
            temp_script.unlink()

def display_with_rich(message: str, model: str, stream: bool):
    """Display using rich formatting."""
    try:
        from rich.console import Console
        from rich.panel import Panel
        from rich.live import Live
        from rich.markdown import Markdown
        
        console = Console()
        
        console.print(Panel(f"[bold cyan]🤖 Sending to {model}...[/bold cyan]"))
        
        from huggingface_hub import InferenceClient
        client = InferenceClient(token=os.environ.get('HF_TOKEN'))
        
        messages = [{"role": "user", "content": message}]
        
        if stream:
            response_text = ""
            with Live("", console=console, refresh_per_second=10) as live:
                for chunk in client.chat_completion(
                    messages=messages,
                    model=model,
                    stream=True,
                    max_tokens=2000
                ):
                    if chunk.choices[0].delta.content:
                        response_text += chunk.choices[0].delta.content
                        live.update(Panel(Markdown(response_text), title="🤖 MiniMax AI Response", border_style="cyan"))
        else:
            response_obj = client.chat_completion(
                messages=messages,
                model=model,
                stream=False,
                max_tokens=2000
            )
            response = response_obj.choices[0].message.content
            console.print(Panel(Markdown(response), title="🤖 MiniMax AI Response", border_style="cyan"))
            
    except ImportError:
        print("⚠️  Rich library not available. Install with: pip install rich")
        display_basic(message, model, stream)

def display_as_json(message: str, model: str, stream: bool):
    """Return response as JSON for widget integration."""
    from huggingface_hub import InferenceClient
    client = InferenceClient(token=os.environ.get('HF_TOKEN'))
    
    messages = [{"role": "user", "content": message}]
    
    try:
        if stream:
            response_text = ""
            for chunk in client.chat_completion(
                messages=messages,
                model=model,
                stream=True,
                max_tokens=2000
            ):
                if chunk.choices[0].delta.content:
                    response_text += chunk.choices[0].delta.content
            response = response_text
        else:
            response_obj = client.chat_completion(
                messages=messages,
                model=model,
                stream=False,
                max_tokens=2000
            )
            response = response_obj.choices[0].message.content
            
        print(json.dumps({"response": response, "model": model}))
        
    except Exception as e:
        print(json.dumps({"error": str(e)}))

if __name__ == "__main__":
    main()