import click
import getpass
import json
import os
import base64
from .config import get_config, save_config, DEFAULT_MODEL
from .crypto import encrypt, decrypt
from .chat import start_chat_session

HARDCODED_PASSWORD = "jai_shree_ram"

def verify_password():
    """Verify the hardcoded password."""
    password = getpass.getpass("Enter password: ")
    if password != HARDCODED_PASSWORD:
        click.echo("❌ Incorrect password!")
        exit(1)
    return True

def get_or_setup_api_key():
    """Get API key from config or prompt user to enter it."""
    config = get_config()
    
    if config.get("encrypted_api_key"):
        try:
            # Decrypt existing API key
            encrypted_key = base64.b64decode(config["encrypted_api_key"])
            api_key = decrypt(encrypted_key)
            return api_key
        except Exception:
            click.echo("⚠️  Error decrypting API key. Please re-enter it.")
    
    # First time or error - ask for API key
    click.echo("🔑 Setting up Groq API key...")
    api_key = getpass.getpass("Enter your Groq API key: ")
    
    # Encrypt and save the API key
    encrypted_key = encrypt(api_key)
    config["encrypted_api_key"] = base64.b64encode(encrypted_key).decode()
    save_config(config)
    
    click.echo("✅ API key saved securely!")
    return api_key

def get_or_setup_supreme_key():
    """Get Supreme (Gemini) API key from config or prompt user to enter it."""
    config = get_config()
    
    if config.get("encrypted_supreme_key"):
        try:
            # Decrypt existing Supreme API key
            encrypted_key = base64.b64decode(config["encrypted_supreme_key"])
            api_key = decrypt(encrypted_key)
            return api_key
        except Exception:
            click.echo("⚠️  Error decrypting Supreme API key. Please re-enter it.")
    
    # First time or error - ask for Supreme API key
    click.echo("🌟 Setting up Supreme (Gemini) API key...")
    click.echo("💡 Get your key from: https://makersuite.google.com/app/apikey")
    api_key = getpass.getpass("Enter your Gemini API key: ")
    
    # Encrypt and save the API key
    encrypted_key = encrypt(api_key)
    config["encrypted_supreme_key"] = base64.b64encode(encrypted_key).decode()
    save_config(config)
    
    click.echo("✅ Supreme API key saved securely!")
    return api_key

@click.group()
def main():
    """
    Hichi - A CLI tool for interacting with Groq AI.
    """
    pass

@main.command()
def hi():
    """Initialize hichi with password verification and API key setup."""
    verify_password()
    click.echo("🙏 Welcome to Hichi!")
    
    # Setup or verify API key
    api_key = get_or_setup_api_key()
    config = get_config()
    
    click.echo(f"📋 Current model: {config.get('model', DEFAULT_MODEL)}")
    click.echo("✨ Hichi is ready to use!")
    click.echo("\nAvailable commands:")
    click.echo("  hichi chat    - Start a chat session (Groq)")
    click.echo("  hichi key     - Update Groq API key")
    click.echo("  hichi s-key   - Update Supreme (Gemini) API key")
    click.echo("  hichi model   - Change model")
    click.echo("  hichi beast   - Fast multimodal AI (text, PDF, DOCX, images) 🦁")
    click.echo("  hichi supreme - Supreme agent for CP/DSA & DBMS problems 🌟")

@main.command()
def key():
    """Update the Groq API key."""
    verify_password()
    
    click.echo("🔑 Updating Groq API key...")
    api_key = getpass.getpass("Enter your new Groq API key: ")
    
    # Encrypt and save the new API key
    config = get_config()
    encrypted_key = encrypt(api_key)
    config["encrypted_api_key"] = base64.b64encode(encrypted_key).decode()
    save_config(config)
    
    click.echo("✅ API key updated successfully!")

@main.command("s-key")
def supreme_key():
    """Update the Supreme (Gemini) API key for hichi supreme feature."""
    verify_password()
    
    click.echo("🌟 Updating Supreme (Gemini) API key...")
    click.echo("💡 Get your key from: https://makersuite.google.com/app/apikey")
    api_key = getpass.getpass("Enter your new Gemini API key: ")
    
    # Encrypt and save the Supreme API key
    config = get_config()
    encrypted_key = encrypt(api_key)
    config["encrypted_supreme_key"] = base64.b64encode(encrypted_key).decode()
    save_config(config)
    
    click.echo("✅ Supreme API key updated successfully!")
    click.echo("🎯 You can now use: hichi supreme")

@main.command()
@click.argument('model_name', required=False)
def model(model_name):
    """Update the model. If no model name provided, shows current model."""
    verify_password()
    
    config = get_config()
    
    if not model_name:
        click.echo(f"📋 Current model: {config.get('model', DEFAULT_MODEL)}")
        click.echo("\nTo change model, use: hichi model <model_name>")
        click.echo("Example: hichi model 'Qwen3 480b Coder'")
        return
    
    config["model"] = model_name
    save_config(config)
    click.echo(f"✅ Model updated to: {model_name}")

@main.command()
def chat():
    """Start an interactive chat session with the AI."""
    verify_password()
    
    # Get API key and model
    api_key = get_or_setup_api_key()
    config = get_config()
    model = config.get('model', DEFAULT_MODEL)
    
    click.echo(f"💬 Starting chat session with {model}")
    click.echo("🔄 Session-only memory (no chat history saved)")
    click.echo("Type 'exit' or 'quit' to end the session")
    click.echo("-" * 50)
    
    start_chat_session(api_key, model)

@main.command()
def supreme():
    """
    Supreme Agent: Multi-stage AI pipeline for CP/DSA and DBMS problem solving.
    
    Stage 1: Extract and summarize problem (Gemini 2.0 Flash)
    Stage 2: Generate optimized code solution (Gemini 2.5 Pro)
    
    Includes human-in-the-loop validation and feedback.
    """
    verify_password()
    
    # Get or setup Supreme API key
    supreme_api_key = get_or_setup_supreme_key()
    
    from .supreme import supreme_agent_pipeline
    supreme_agent_pipeline(supreme_api_key)

@main.command()
@click.option('--prompt', '-p', help='Quick prompt without entering interactive mode')
@click.option('--file', '-f', 'files', multiple=True, help='Attach files (PDF, DOCX, images)')
def beast(prompt, files):
    """
    Beast Mode: Fast multimodal AI chat powered by Gemini 2.0 Flash.
    
    Supports: Text, PDF, DOCX, and Images (JPG, PNG, etc.)
    
    Interactive mode (default): Just run 'hichi beast'
    One-shot mode: Use --prompt with optional --file attachments
    
    Examples:
      hichi beast                           # Interactive chat
      hichi beast -p "Explain this image" -f photo.jpg
      hichi beast -p "Summarize" -f doc.pdf -f report.docx
    """
    verify_password()
    
    # Get or setup Supreme API key
    supreme_api_key = get_or_setup_supreme_key()
    
    from .beast import beast_mode_session, beast_one_shot
    
    # One-shot mode if prompt provided
    if prompt:
        beast_one_shot(supreme_api_key, prompt, list(files) if files else None)
    else:
        # Interactive mode
        beast_mode_session(supreme_api_key)

if __name__ == "__main__":
    main()