"""
Beast Mode Module for Hichi
Fast, multimodal AI assistant powered by Gemini 2.0 Flash.
Supports: Text, PDF, DOCX, and Image inputs.
"""

import click
import os
from typing import Optional, List
from pathlib import Path

# Model configuration
BEAST_MODEL = "gemini-2.0-flash-exp"

# Check for required libraries
try:
    import google.generativeai as genai
    GEMINI_AVAILABLE = True
except ImportError:
    GEMINI_AVAILABLE = False

try:
    import fitz  # PyMuPDF for PDF
    PDF_AVAILABLE = True
except ImportError:
    PDF_AVAILABLE = False

try:
    from docx import Document  # python-docx for DOCX
    DOCX_AVAILABLE = True
except ImportError:
    DOCX_AVAILABLE = False

try:
    from PIL import Image
    IMAGE_AVAILABLE = True
except ImportError:
    IMAGE_AVAILABLE = False


def extract_text_from_pdf(file_path: str) -> Optional[str]:
    """Extract text from PDF file."""
    if not PDF_AVAILABLE:
        return None
    
    try:
        text_parts = []
        with fitz.open(file_path) as doc:
            for page_num, page in enumerate(doc, 1):
                text = page.get_text()
                if text.strip():
                    text_parts.append(f"--- Page {page_num} ---\n{text}")
        return "\n\n".join(text_parts)
    except Exception as e:
        click.echo(f"⚠️  Error reading PDF: {e}")
        return None


def extract_text_from_docx(file_path: str) -> Optional[str]:
    """Extract text from DOCX file."""
    if not DOCX_AVAILABLE:
        return None
    
    try:
        doc = Document(file_path)
        paragraphs = [para.text for para in doc.paragraphs if para.text.strip()]
        return "\n\n".join(paragraphs)
    except Exception as e:
        click.echo(f"⚠️  Error reading DOCX: {e}")
        return None


def extract_text_from_txt(file_path: str) -> Optional[str]:
    """Extract text from plain text file."""
    try:
        with open(file_path, 'r', encoding='utf-8', errors='ignore') as f:
            return f.read()
    except Exception as e:
        click.echo(f"⚠️  Error reading text file: {e}")
        return None


def load_image(file_path: str):
    """Load image file for Gemini API."""
    if not IMAGE_AVAILABLE:
        return None
    
    try:
        img = Image.open(file_path)
        return img
    except Exception as e:
        click.echo(f"⚠️  Error loading image: {e}")
        return None


def process_file(file_path: str):
    """Process file and return appropriate content for Gemini API."""
    file_path = os.path.expanduser(file_path)
    
    if not os.path.exists(file_path):
        click.echo(f"❌ File not found: {file_path}")
        return None, None
    
    ext = Path(file_path).suffix.lower()
    
    # Handle images
    if ext in ['.jpg', '.jpeg', '.png', '.gif', '.bmp', '.webp']:
        click.echo(f"📸 Loading image: {file_path}")
        img = load_image(file_path)
        if img:
            return 'image', img
        else:
            if not IMAGE_AVAILABLE:
                click.echo("❌ Image support not installed. Run: pip install pillow")
            return None, None
    
    # Handle PDF
    elif ext == '.pdf':
        click.echo(f"📄 Extracting text from PDF: {file_path}")
        text = extract_text_from_pdf(file_path)
        if text:
            return 'text', text
        else:
            if not PDF_AVAILABLE:
                click.echo("❌ PDF support not installed. Run: pip install pymupdf")
            return None, None
    
    # Handle DOCX
    elif ext in ['.docx', '.doc']:
        click.echo(f"📝 Extracting text from DOCX: {file_path}")
        text = extract_text_from_docx(file_path)
        if text:
            return 'text', text
        else:
            if not DOCX_AVAILABLE:
                click.echo("❌ DOCX support not installed. Run: pip install python-docx")
            return None, None
    
    # Handle plain text
    elif ext in ['.txt', '.md', '.py', '.js', '.java', '.cpp', '.c', '.h', '.html', '.css', '.json', '.xml', '.yaml', '.yml']:
        click.echo(f"📄 Reading text file: {file_path}")
        text = extract_text_from_txt(file_path)
        if text:
            return 'text', text
        else:
            return None, None
    
    else:
        click.echo(f"⚠️  Unsupported file type: {ext}")
        click.echo("Supported: .pdf, .docx, .txt, .md, .jpg, .png, .gif, etc.")
        return None, None


def call_gemini_beast(api_key: str, prompt: str, file_contents: Optional[List] = None) -> Optional[str]:
    """Call Gemini 2.0 Flash with text and/or image inputs."""
    if not GEMINI_AVAILABLE:
        return "❌ Google Generative AI library not installed. Run: pip install google-generativeai"
    
    if not api_key:
        return "❌ Supreme API key not configured. Run: hichi s-key"
    
    try:
        genai.configure(api_key=api_key)
        
        generation_config = {
            "temperature": 0.7,
            "top_p": 0.95,
            "top_k": 40,
            "max_output_tokens": 8192,
        }
        
        model = genai.GenerativeModel(
            model_name=BEAST_MODEL,
            generation_config=generation_config,
        )
        
        # Build content list for multimodal input
        content_parts = []
        
        # Add file contents first (images or text)
        if file_contents:
            for content_type, content in file_contents:
                if content_type == 'image':
                    content_parts.append(content)
                elif content_type == 'text':
                    content_parts.append(f"[Document Content]\n{content}\n\n")
        
        # Add user prompt
        content_parts.append(prompt)
        
        response = model.generate_content(content_parts)
        return response.text
    
    except Exception as e:
        error_msg = str(e)
        if "API_KEY_INVALID" in error_msg or "401" in error_msg:
            return "❌ Invalid Gemini API key. Please update with: hichi s-key"
        elif "RATE_LIMIT" in error_msg or "429" in error_msg:
            return "⏰ Rate limit exceeded. Please try again in a moment."
        elif "NOT_FOUND" in error_msg or "404" in error_msg:
            return f"❌ Model '{BEAST_MODEL}' not found. Check model availability."
        else:
            return f"❌ API Error: {error_msg}"


def beast_mode_session(api_key: str):
    """
    Beast Mode: Interactive multimodal chat session with Gemini 2.0 Flash.
    Supports text input, file attachments (PDF, DOCX), and images.
    """
    click.echo("🦁 " + click.style("HICHI BEAST MODE", fg='red', bold=True) + " 🦁")
    click.echo("Powered by Gemini 2.0 Flash - Fast & Multimodal")
    click.echo("Supports: Text, PDF, DOCX, Images (JPG, PNG, etc.)\n")
    
    # Session context (conversation history)
    conversation_history = []
    
    click.echo("Commands:")
    click.echo("  /attach <file>  - Attach a file (PDF, DOCX, or image)")
    click.echo("  /clear          - Clear conversation history")
    click.echo("  /exit or /quit  - Exit Beast Mode")
    click.echo("-" * 60 + "\n")
    
    # Track attached files for current message
    attached_files = []
    
    while True:
        try:
            user_input = click.prompt("\n🦁 Beast", type=str, prompt_suffix=" ")
            
            if user_input.lower() in ['/exit', '/quit', 'exit', 'quit']:
                click.echo("👋 Exiting Beast Mode!")
                break
            
            # Handle /attach command
            if user_input.lower().startswith('/attach '):
                file_path = user_input[8:].strip()
                content_type, content = process_file(file_path)
                
                if content is not None:
                    attached_files.append((content_type, content))
                    if content_type == 'image':
                        click.echo(f"✅ Image attached: {Path(file_path).name}")
                    else:
                        click.echo(f"✅ Document attached: {Path(file_path).name} ({len(content)} chars)")
                continue
            
            # Handle /clear command
            if user_input.lower() == '/clear':
                conversation_history.clear()
                attached_files.clear()
                click.echo("🗑️  Conversation history cleared!")
                continue
            
            if user_input.strip() == "":
                continue
            
            # Build prompt with conversation history
            if conversation_history:
                context = "\n\n".join([
                    f"{'User' if msg['role'] == 'user' else 'Assistant'}: {msg['content']}"
                    for msg in conversation_history[-5:]  # Last 5 exchanges for context
                ])
                full_prompt = f"Previous conversation:\n{context}\n\nUser: {user_input}"
            else:
                full_prompt = user_input
            
            # Show thinking indicator
            click.echo("🤔 Thinking...", nl=False)
            
            # Call Gemini API with attached files
            response = call_gemini_beast(api_key, full_prompt, attached_files if attached_files else None)
            
            # Clear thinking indicator
            click.echo("\r" + " " * 20 + "\r", nl=False)
            
            if response and not response.startswith("❌"):
                # Add to conversation history
                conversation_history.append({"role": "user", "content": user_input})
                conversation_history.append({"role": "assistant", "content": response})
                
                # Display response
                click.echo("🤖 Beast: " + response)
                
                # Clear attached files after successful response
                if attached_files:
                    click.echo(f"\n📎 {len(attached_files)} file(s) processed")
                    attached_files.clear()
            else:
                click.echo(response)
        
        except KeyboardInterrupt:
            click.echo("\n\n👋 Beast Mode interrupted!")
            break
        except EOFError:
            click.echo("\n\n👋 Goodbye!")
            break
        except Exception as e:
            click.echo(f"\n❌ An error occurred: {str(e)}")
            continue


def beast_one_shot(api_key: str, prompt: str, files: List[str] = None):
    """
    Beast Mode: Single query with optional file attachments.
    For quick questions without entering interactive mode.
    """
    click.echo("🦁 " + click.style("HICHI BEAST MODE", fg='red', bold=True) + " 🦁\n")
    
    # Process attached files
    file_contents = []
    if files:
        for file_path in files:
            content_type, content = process_file(file_path)
            if content is not None:
                file_contents.append((content_type, content))
    
    # Show thinking indicator
    click.echo("🤔 Processing...", nl=False)
    
    # Call Gemini API
    response = call_gemini_beast(api_key, prompt, file_contents if file_contents else None)
    
    # Clear thinking indicator
    click.echo("\r" + " " * 20 + "\r", nl=False)
    
    if response:
        click.echo("🤖 Beast:\n")
        click.echo(response)
        click.echo()
    else:
        click.echo("❌ No response generated")
