#!/usr/bin/env python3
"""
Automatic model downloader for Qwen2.5-0.5B-Instruct
Downloads from HuggingFace and places in correct location
"""
import os
import sys
import urllib.request
from pathlib import Path
from typing import Optional

try:
    from tqdm import tqdm
    HAS_TQDM = True
except ImportError:
    HAS_TQDM = False


class DownloadProgressBar:
    """Progress bar for download"""
    def __init__(self):
        self.pbar = None

    def __call__(self, block_num, block_size, total_size):
        if HAS_TQDM:
            if not self.pbar:
                self.pbar = tqdm(
                    total=total_size,
                    unit='B',
                    unit_scale=True,
                    unit_divisor=1024,
                    desc='Downloading model'
                )
            downloaded = block_num * block_size
            if downloaded < total_size:
                self.pbar.update(block_size)
            else:
                self.pbar.close()
        else:
            # Simple progress without tqdm
            if total_size > 0:
                downloaded = block_num * block_size
                percent = min(100, (downloaded / total_size) * 100)
                if block_num % 100 == 0:  # Print every 100 blocks
                    print(f"Progress: {percent:.1f}%", end='\r')


def download_model(destination: Optional[str] = None):
    """
    Download Qwen2.5-0.5B-Instruct-Q4_0 GGUF model
    
    Args:
        destination: Target path for model file. 
                    Defaults to /home/luigi/rts/qwen2.5-coder-1.5b-instruct-q4_0.gguf
    
    Returns:
        str: Path to downloaded model
    """
    # Model URL from HuggingFace
    MODEL_URL = "https://huggingface.co/Qwen/Qwen2.5-Coder-1.5B-Instruct-GGUF/resolve/main/qwen2.5-coder-1.5b-instruct-q4_0.gguf"
    
    # Default destination
    if destination is None:
        destination = "/home/luigi/rts/qwen2.5-coder-1.5b-instruct-q4_0.gguf"
    
    destination_path = Path(destination)
    
    # Check if already exists
    if destination_path.exists():
        size_mb = destination_path.stat().st_size / (1024 * 1024)
        print(f"✅ Model already exists: {destination_path}")
        print(f"   Size: {size_mb:.1f} MB")
        return str(destination_path)
    
    # Create parent directory if needed
    destination_path.parent.mkdir(parents=True, exist_ok=True)
    
    print(f"📦 Downloading Qwen2.5-0.5B-Instruct model...")
    print(f"   Source: {MODEL_URL}")
    print(f"   Destination: {destination_path}")
    print(f"   Expected size: ~350 MB")
    print()
    
    try:
        # Download with progress bar
        urllib.request.urlretrieve(
            MODEL_URL,
            destination_path,
            reporthook=DownloadProgressBar()
        )
        
        # Verify download
        if destination_path.exists():
            size_mb = destination_path.stat().st_size / (1024 * 1024)
            print()
            print(f"✅ Model downloaded successfully!")
            print(f"   Location: {destination_path}")
            print(f"   Size: {size_mb:.1f} MB")
            return str(destination_path)
        else:
            print("❌ Download failed: File not created")
            return None
            
    except Exception as e:
        print(f"❌ Download error: {e}")
        # Clean up partial download
        if destination_path.exists():
            destination_path.unlink()
        return None


def main():
    """Main entry point"""
    import argparse
    
    parser = argparse.ArgumentParser(
        description='Download Qwen2.5-0.5B-Instruct model for tactical analysis'
    )
    parser.add_argument(
        '--destination',
        '-d',
        type=str,
        default=None,
        help='Destination path (default: /home/luigi/rts/qwen2.5-coder-1.5b-instruct-q4_0.gguf)'
    )
    
    args = parser.parse_args()
    
    print("=" * 60)
    print("  Qwen2.5-0.5B-Instruct Model Downloader")
    print("=" * 60)
    print()
    
    model_path = download_model(args.destination)
    
    if model_path:
        print()
        print("🎉 Setup complete! The AI tactical analysis system is ready.")
        print()
        print("To use:")
        print("  1. Restart the web server: python3 web/app.py")
        print("  2. The model will be automatically loaded")
        print("  3. Tactical analysis will run every 30 seconds")
        sys.exit(0)
    else:
        print()
        print("❌ Setup failed. Please try again or download manually from:")
        print("   https://huggingface.co/Qwen/Qwen2.5-Coder-1.5B-Instruct-GGUF")
        sys.exit(1)


if __name__ == '__main__':
    main()