#!/usr/bin/env python3
# @skill: image-generation

"""
MiniMax Image Generation Script
Generate images using MiniMax API
"""

import argparse
import os
import time
import requests
from urllib.parse import urlparse


def parse_args():
    """Parse command line arguments"""
    parser = argparse.ArgumentParser(
        description="Generate images using MiniMax API",
        formatter_class=argparse.RawDescriptionHelpFormatter
    )
    
    parser.add_argument(
        "--api-key",
        type=str,
        required=False,
        help="MiniMax API key (can also be set via MINIMAX_API_KEY environment variable)"
    )
    
    parser.add_argument(
        "--prompt",
        type=str,
        required=True,
        help="Image generation prompt"
    )
    
    parser.add_argument(
        "--model",
        type=str,
        default="image-01",
        help="Image generation model (default: image-01)"
    )
    
    parser.add_argument(
        "--aspect-ratio",
        type=str,
        default="1:1",
        help="Image aspect ratio (default: 1:1, options: 16:9, 9:16, etc.)"
    )
    
    parser.add_argument(
        "--response-format",
        type=str,
        default="url",
        choices=["url", "base64"],
        help="Response format (default: url)"
    )
    
    parser.add_argument(
        "--n",
        type=int,
        default=1,
        choices=[1, 2, 3],
        help="Number of images to generate (default: 1, max: 3)"
    )
    
    parser.add_argument(
        "--prompt-optimizer",
        type=lambda x: x.lower() == "true",
        default=False,
        help="Enable prompt optimizer (default: false)"
    )
    
    # Image-to-image (subject reference) parameters
    parser.add_argument(
        "--subject-reference",
        type=str,
        default=None,
        help="Reference image URL or local file path for image-to-image generation"
    )
    
    parser.add_argument(
        "--subject-type",
        type=str,
        default="character",
        choices=["character", "product", "logo", "video_subject", "other"],
        help="Subject reference type (default: character)"
    )
    
    parser.add_argument(
        "--seed",
        type=int,
        default=None,
        help="Random seed for reproducible generation (optional)"
    )
    
    parser.add_argument(
        "--output-dir",
        type=str,
        default="./output",
        help="Output directory for images (default: ./output)"
    )
    
    parser.add_argument(
        "--api-base",
        type=str,
        default="https://api.minimaxi.com",
        help="API base URL (default: https://api.minimaxi.com)"
    )
    
    return parser.parse_args()


def download_image(url: str, output_path: str) -> bool:
    """Download image to local file"""
    try:
        response = requests.get(url, timeout=30)
        response.raise_for_status()
        
        with open(output_path, "wb") as f:
            f.write(response.content)
        
        print(f"  [OK] Saved: {output_path}")
        return True
    except Exception as e:
        print(f"  [FAIL] Download failed: {e}")
        return False


def read_local_image(file_path: str) -> str:
    """
    Read local image file and return as base64 encoded string.
    Returns the base64 string or None if failed.
    """
    try:
        with open(file_path, "rb") as f:
            image_data = f.read()
            import base64
            return base64.b64encode(image_data).decode("utf-8")
    except Exception as e:
        print(f"  [FAIL] Failed to read local image: {e}")
        return None


def build_subject_reference(subject_ref: str, subject_type: str) -> dict:
    """
    Build subject_reference object from URL or local file path.
    
    Args:
        subject_ref: URL or local file path
        subject_type: Type of subject (character, product, logo, etc.)
    
    Returns:
        dict with subject_reference structure
    """
    # Check if it's a URL or local file
    if subject_ref.startswith(("http://", "https://")):
        # It's a URL
        return {
            "type": subject_type,
            "image_file": subject_ref
        }
    else:
        # It's a local file - convert to base64
        print(f"  Processing local image: {subject_ref}")
        base64_data = read_local_image(subject_ref)
        if base64_data:
            return {
                "type": subject_type,
                "image_file": f"data:image/jpeg;base64,{base64_data}"
            }
        else:
            return None


def generate_images(args):
    """Call MiniMax API to generate images"""
    url = f"{args.api_base}/v1/image_generation"
    
    headers = {
        "Authorization": f"Bearer {args.api_key}",
        "Content-Type": "application/json"
    }
    
    payload = {
        "model": args.model,
        "prompt": args.prompt,
        "aspect_ratio": args.aspect_ratio,
        "response_format": args.response_format,
        "n": args.n,
        "prompt_optimizer": args.prompt_optimizer
    }
    
    # Add seed if provided
    if args.seed is not None:
        payload["seed"] = args.seed
    
    # Add subject_reference for image-to-image generation
    if args.subject_reference:
        subject_ref = build_subject_reference(args.subject_reference, args.subject_type)
        if subject_ref:
            payload["subject_reference"] = [subject_ref]
        else:
            print("Warning: Failed to process subject reference, continuing without it.")
    
    print(f"\n{'='*60}")
    print(f"MiniMax Image Generation")
    print(f"{'='*60}")
    print(f"Model: {args.model}")
    print(f"Prompt: {args.prompt}")
    print(f"Aspect Ratio: {args.aspect_ratio}")
    print(f"Number: {args.n}")
    print(f"Prompt Optimizer: {'Enabled' if args.prompt_optimizer else 'Disabled'}")
    if args.seed is not None:
        print(f"Seed: {args.seed}")
    if args.subject_reference:
        print(f"Subject Reference: {args.subject_type} - {args.subject_reference}")
    print(f"{'='*60}\n")
    
    try:
        print("Generating images...")
        response = requests.post(url, headers=headers, json=payload, timeout=60)
        response.raise_for_status()
        
        result = response.json()
        
        if result.get("base_resp", {}).get("status_code") != 0:
            error_msg = result.get("base_resp", {}).get("status_msg", "Unknown error")
            print(f"API Error: {error_msg}")
            return False
        
        # Create output directory
        os.makedirs(args.output_dir, exist_ok=True)
        
        # Process returned images
        image_urls = result.get("data", {}).get("image_urls", [])
        metadata = result.get("metadata", {})
        
        print(f"\nSuccessfully generated {metadata.get('success_count', len(image_urls))} images:")
        
        timestamp = int(time.time())
        saved_count = 0
        
        for i, image_url in enumerate(image_urls, 1):
            # Extract file extension from URL
            parsed = urlparse(image_url)
            path = parsed.path
            ext = os.path.splitext(path)[1] if "." in path else ".jpeg"
            
            filename = f"generated_image_{i}_{timestamp}{ext}"
            output_path = os.path.join(args.output_dir, filename)
            
            if download_image(image_url, output_path):
                saved_count += 1
        
        print(f"\n{'='*60}")
        print(f"Done! Successfully saved {saved_count}/{len(image_urls)} images")
        print(f"Output directory: {os.path.abspath(args.output_dir)}")
        print(f"{'='*60}\n")
        
        return saved_count > 0
        
    except requests.exceptions.RequestException as e:
        print(f"\nRequest Error: {e}")
        return False
    except Exception as e:
        print(f"\nUnexpected Error: {e}")
        return False


def main():
    """Main function"""
    args = parse_args()
    
    # Get API key from argument or environment variable (optional)
    if not args.api_key:
        args.api_key = os.environ.get("MINIMAX_API_KEY", "")
    
    # If still no API key, prompt user to enter it
    if not args.api_key:
        print("Error: API key not int ENV, and it is required.")
    else:
        generate_images(args)


if __name__ == "__main__":
    main()
