#!/usr/bin/env python3
"""
Batch upload files to PyGEAI with progress tracking and error handling.
"""

import argparse
from pathlib import Path
from typing import List
from concurrent.futures import ThreadPoolExecutor, as_completed
from tqdm import tqdm
from pygeai import Client


def validate_file(file_path: Path, max_size_mb: int = 512) -> tuple[bool, str]:
    """Validate file before upload."""
    if not file_path.exists():
        return False, "File does not exist"
    
    size_mb = file_path.stat().st_size / 1024 / 1024
    if size_mb > max_size_mb:
        return False, f"File too large: {size_mb:.2f}MB"
    
    allowed_extensions = {".pdf", ".txt", ".docx", ".csv", ".json", ".jsonl", ".md"}
    if file_path.suffix.lower() not in allowed_extensions:
        return False, f"Unsupported extension: {file_path.suffix}"
    
    return True, ""


def upload_single_file(client: Client, file_path: Path, purpose: str):
    """Upload a single file with error handling."""
    try:
        is_valid, error = validate_file(file_path)
        if not is_valid:
            return {"success": False, "file": file_path.name, "error": error}
        
        file = client.files.upload(file=file_path, purpose=purpose)
        
        return {
            "success": True,
            "file": file_path.name,
            "file_id": file.id,
            "bytes": file.bytes
        }
    except Exception as e:
        return {"success": False, "file": file_path.name, "error": str(e)}


def batch_upload(
    client: Client,
    files: List[Path],
    purpose: str = "assistants",
    max_workers: int = 5
):
    """Upload multiple files in parallel."""
    results = {"success": [], "failed": []}
    
    with ThreadPoolExecutor(max_workers=max_workers) as executor:
        futures = {
            executor.submit(upload_single_file, client, f, purpose): f
            for f in files
        }
        
        with tqdm(total=len(files), desc="Uploading files") as pbar:
            for future in as_completed(futures):
                result = future.result()
                
                if result["success"]:
                    results["success"].append(result)
                else:
                    results["failed"].append(result)
                
                pbar.update(1)
    
    return results


def main():
    parser = argparse.ArgumentParser(description="Batch upload files to PyGEAI")
    parser.add_argument("path", type=Path, help="File or directory to upload")
    parser.add_argument("--purpose", default="assistants", help="File purpose")
    parser.add_argument("--workers", type=int, default=5, help="Parallel workers")
    parser.add_argument("--recursive", action="store_true", help="Recursive directory scan")
    
    args = parser.parse_args()
    
    # Collect files
    if args.path.is_file():
        files = [args.path]
    elif args.path.is_dir():
        if args.recursive:
            files = list(args.path.rglob("*"))
        else:
            files = list(args.path.glob("*"))
        files = [f for f in files if f.is_file()]
    else:
        print(f"Error: {args.path} not found")
        return
    
    print(f"📁 Found {len(files)} files to upload\n")
    
    # Upload
    client = Client()
    results = batch_upload(client, files, args.purpose, args.workers)
    
    # Report
    print(f"\n✅ Success: {len(results['success'])}")
    print(f"❌ Failed: {len(results['failed'])}")
    
    if results["failed"]:
        print("\nFailed uploads:")
        for item in results["failed"]:
            print(f"  - {item['file']}: {item['error']}")
    
    total_bytes = sum(r["bytes"] for r in results["success"])
    print(f"\n📊 Total uploaded: {total_bytes / 1024 / 1024:.2f} MB")


if __name__ == "__main__":
    main()
