#!/usr/bin/env python3
"""
DALL-E Image Generator
Generates images from text prompts and saves to Etsy Designs folder.
Tracks spending in dalle-spend.json
"""

import argparse
import os
import sys
import json
import urllib.request
import urllib.error
from datetime import datetime
from pathlib import Path

# Configuration
# API key from environment variable (set in ~/.zshrc or provide via command line)
API_KEY = os.environ.get("OPENAI_DALLE_API_KEY", "")
OUTPUT_DIR = Path.home() / "Documents" / "Etsy Designs" / "01-Source-PNG"
SPEND_FILE = Path.home() / ".openclaw" / "workspace" / "data" / "dalle-spend.json"
API_URL = "https://api.openai.com/v1/images/generations"

# Pricing (DALL-E 3)
PRICING = {
    ("dall-e-3", "1024x1024", "standard"): 0.04,
    ("dall-e-3", "1024x1024", "hd"): 0.08,
    ("dall-e-3", "1792x1024", "standard"): 0.08,
    ("dall-e-3", "1792x1024", "hd"): 0.12,
    ("dall-e-3", "1024x1792", "standard"): 0.08,
    ("dall-e-3", "1024x1792", "hd"): 0.12,
    ("dall-e-2", "256x256", "standard"): 0.016,
    ("dall-e-2", "512x512", "standard"): 0.018,
    ("dall-e-2", "1024x1024", "standard"): 0.02,
}

def load_spend_data():
    """Load or initialize spend tracking data."""
    SPEND_FILE.parent.mkdir(parents=True, exist_ok=True)
    if SPEND_FILE.exists():
        with open(SPEND_FILE) as f:
            return json.load(f)
    return {
        "totalSpend": 0.0,
        "totalImages": 0,
        "dailySpend": {},
        "monthlySpend": {},
        "history": []
    }

def save_spend_data(data):
    """Save spend tracking data."""
    with open(SPEND_FILE, 'w') as f:
        json.dump(data, f, indent=2)

def record_spend(cost: float, prompt: str, filepath: str, model: str, size: str, quality: str):
    """Record a generation in the spend tracker."""
    data = load_spend_data()
    
    now = datetime.now()
    today = now.strftime("%Y-%m-%d")
    month = now.strftime("%Y-%m")
    
    # Update totals
    data["totalSpend"] = round(data["totalSpend"] + cost, 4)
    data["totalImages"] += 1
    
    # Update daily
    if today not in data["dailySpend"]:
        data["dailySpend"][today] = {"spend": 0.0, "images": 0}
    data["dailySpend"][today]["spend"] = round(data["dailySpend"][today]["spend"] + cost, 4)
    data["dailySpend"][today]["images"] += 1
    
    # Update monthly
    if month not in data["monthlySpend"]:
        data["monthlySpend"][month] = {"spend": 0.0, "images": 0}
    data["monthlySpend"][month]["spend"] = round(data["monthlySpend"][month]["spend"] + cost, 4)
    data["monthlySpend"][month]["images"] += 1
    
    # Add to history (keep last 100)
    data["history"].append({
        "timestamp": now.isoformat(),
        "cost": cost,
        "model": model,
        "size": size,
        "quality": quality,
        "prompt": prompt[:100] + "..." if len(prompt) > 100 else prompt,
        "file": str(filepath)
    })
    data["history"] = data["history"][-100:]
    
    # Update last modified
    data["lastUpdated"] = now.isoformat()
    
    save_spend_data(data)
    return data

def generate_image(prompt: str, size: str = "1024x1024", quality: str = "standard", model: str = "dall-e-3") -> str:
    """Generate an image from a prompt and save it."""
    
    # Check for API key
    if not API_KEY:
        print("Error: OPENAI_DALLE_API_KEY environment variable not set.", file=sys.stderr)
        print("Set it with: export OPENAI_DALLE_API_KEY='sk-...'", file=sys.stderr)
        sys.exit(1)
    
    # Ensure output directory exists
    OUTPUT_DIR.mkdir(parents=True, exist_ok=True)
    
    # Get cost for this generation
    cost = PRICING.get((model, size, quality), 0.04)
    
    # Prepare request
    headers = {
        "Authorization": f"Bearer {API_KEY}",
        "Content-Type": "application/json"
    }
    
    data = json.dumps({
        "model": model,
        "prompt": prompt,
        "n": 1,
        "size": size,
        "quality": quality,
        "response_format": "url"
    }).encode('utf-8')
    
    # Make request
    req = urllib.request.Request(API_URL, data=data, headers=headers, method="POST")
    
    try:
        with urllib.request.urlopen(req, timeout=120) as response:
            result = json.loads(response.read().decode('utf-8'))
    except urllib.error.HTTPError as e:
        error_body = e.read().decode('utf-8')
        print(f"API Error {e.code}: {error_body}", file=sys.stderr)
        sys.exit(1)
    except Exception as e:
        print(f"Error: {e}", file=sys.stderr)
        sys.exit(1)
    
    # Get image URL
    image_url = result["data"][0]["url"]
    revised_prompt = result["data"][0].get("revised_prompt", prompt)
    
    # Download image with descriptive name
    date_prefix = datetime.now().strftime("%y%m%d")
    # Create slug from prompt (first 40 chars, alphanumeric + spaces → dashes)
    slug = prompt[:40].strip()
    slug = ''.join(c if c.isalnum() or c == ' ' else '' for c in slug)
    slug = '-'.join(slug.split()).lower()
    if not slug:
        slug = "design"
    filename = f"{date_prefix}-{slug}.png"
    filepath = OUTPUT_DIR / filename
    
    # Handle duplicates
    counter = 1
    while filepath.exists():
        filename = f"{date_prefix}-{slug}-{counter}.png"
        filepath = OUTPUT_DIR / filename
        counter += 1
    
    try:
        urllib.request.urlretrieve(image_url, filepath)
    except Exception as e:
        print(f"Download error: {e}", file=sys.stderr)
        sys.exit(1)
    
    # Record spend
    spend_data = record_spend(cost, prompt, filepath, model, size, quality)
    
    # Output result
    print(f"✓ Image saved: {filepath}")
    print(f"  Cost: ${cost:.2f} | Today: ${spend_data['dailySpend'].get(datetime.now().strftime('%Y-%m-%d'), {}).get('spend', 0):.2f} | Total: ${spend_data['totalSpend']:.2f}")
    print(f"  Revised prompt: {revised_prompt[:80]}..." if len(revised_prompt) > 80 else f"  Revised prompt: {revised_prompt}")
    
    return str(filepath)

def show_stats():
    """Show current spend statistics."""
    data = load_spend_data()
    today = datetime.now().strftime("%Y-%m-%d")
    month = datetime.now().strftime("%Y-%m")
    
    print("=== DALL-E Spend Statistics ===")
    print(f"Total: ${data['totalSpend']:.2f} ({data['totalImages']} images)")
    
    daily = data['dailySpend'].get(today, {"spend": 0, "images": 0})
    print(f"Today: ${daily['spend']:.2f} ({daily['images']} images)")
    
    monthly = data['monthlySpend'].get(month, {"spend": 0, "images": 0})
    print(f"This month: ${monthly['spend']:.2f} ({monthly['images']} images)")

def main():
    parser = argparse.ArgumentParser(description="Generate images with DALL-E")
    parser.add_argument("prompt", nargs="?", help="Image prompt")
    parser.add_argument("--size", default="1024x1024", choices=["1024x1024", "1792x1024", "1024x1792"], help="Image size")
    parser.add_argument("--quality", default="standard", choices=["standard", "hd"], help="Image quality")
    parser.add_argument("--model", default="dall-e-3", choices=["dall-e-2", "dall-e-3"], help="Model to use")
    parser.add_argument("--stats", action="store_true", help="Show spend statistics")
    
    args = parser.parse_args()
    
    if args.stats:
        show_stats()
        return
    
    # Read prompt from stdin if not provided
    if args.prompt:
        prompt = args.prompt
    elif not sys.stdin.isatty():
        prompt = sys.stdin.read().strip()
    else:
        print("Usage: generate-image.py 'your prompt here'")
        print("   or: echo 'your prompt' | generate-image.py")
        print("   or: generate-image.py --stats")
        sys.exit(1)
    
    if not prompt:
        print("Error: No prompt provided", file=sys.stderr)
        sys.exit(1)
    
    generate_image(prompt, args.size, args.quality, args.model)

if __name__ == "__main__":
    main()

# TONY-APPROVED: 2026-03-01 | sha:9d2705b4