#!/usr/bin/env python3

import os
import sys
import argparse
import time
import base64
from pathlib import Path
import requests


def load_env():
    env_path = Path(__file__).parent / ".env"
    if env_path.exists():
        for line in env_path.read_text().splitlines():
            line = line.strip()
            if line and not line.startswith("#") and "=" in line:
                k, v = line.split("=", 1)
                os.environ.setdefault(k.strip(), v.strip().strip("\"'"))


load_env()

API_TOKEN = os.environ.get("CHUTES_API_TOKEN")
API_URL = "https://chutes-qwen-image-edit-2511.chutes.ai/generate"


def image_to_base64(image_path):
    if not os.path.exists(image_path):
        raise FileNotFoundError(f"Image file not found: {image_path}")

    with open(image_path, "rb") as f:
        image_bytes = f.read()

    return base64.b64encode(image_bytes).decode("utf-8")


def edit_image(
    prompt,
    image_path,
    width=1024,
    height=1024,
    steps=40,
    seed=None,
    cfg_scale=4,
    negative_prompt="",
):
    if not API_TOKEN:
        print("Error: CHUTES_API_TOKEN not set in environment", file=sys.stderr)
        sys.exit(1)

    if not os.path.exists(image_path):
        print(f"Error: Image file not found: {image_path}", file=sys.stderr)
        sys.exit(1)

    if not prompt:
        print("Error: Prompt cannot be empty", file=sys.stderr)
        sys.exit(1)

    image_b64 = image_to_base64(image_path)

    payload = {
        "seed": seed,
        "width": width,
        "height": height,
        "prompt": prompt,
        "image_b64s": [image_b64],
        "true_cfg_scale": cfg_scale,
        "negative_prompt": negative_prompt,
        "num_inference_steps": steps,
    }

    try:
        headers = {
            "Authorization": f"Bearer {API_TOKEN}",
            "Content-Type": "application/json",
        }

        response = requests.post(API_URL, headers=headers, json=payload, timeout=300)
        response.raise_for_status()

        content_type = response.headers.get("Content-Type", "")

        if "image/" in content_type:
            image_bytes = response.content
        else:
            result = response.json()
            if isinstance(result, list) and len(result) > 0:
                item = result[0]
                image_data = item.get("data", "")
                if image_data.startswith("data:image"):
                    image_bytes = base64.b64decode(image_data.split(",", 1)[1])
                else:
                    image_bytes = base64.b64decode(image_data)
            else:
                print("Error: Invalid response format", file=sys.stderr)
                sys.exit(1)

        timestamp = int(time.time())
        filename = f"edited_{timestamp}.jpg"

        with open(filename, "wb") as f:
            f.write(image_bytes)

        print(f"Image saved: {filename} [{timestamp}]")

    except requests.exceptions.RequestException as e:
        print(f"Error: API request failed - {e}", file=sys.stderr)
        sys.exit(1)
    except Exception as e:
        print(f"Error: {e}", file=sys.stderr)
        sys.exit(1)


def main():
    parser = argparse.ArgumentParser(description="Edit images with AI")
    parser.add_argument("prompt", help="Text prompt describing the edit")
    parser.add_argument("image_path", help="Path to input image file")
    parser.add_argument(
        "--width", type=int, default=1024, help="Output width (128-2048)"
    )
    parser.add_argument(
        "--height", type=int, default=1024, help="Output height (128-2048)"
    )
    parser.add_argument("--steps", type=int, default=40, help="Inference steps (5-100)")
    parser.add_argument("--seed", type=int, default=None, help="Random seed")
    parser.add_argument(
        "--cfg-scale", type=float, default=4, help="True CFG scale (0-10)"
    )
    parser.add_argument(
        "--negative-prompt", type=str, default="", help="Negative prompt"
    )

    args = parser.parse_args()

    if not (128 <= args.width <= 2048):
        print("Error: width must be between 128 and 2048", file=sys.stderr)
        sys.exit(1)
    if not (128 <= args.height <= 2048):
        print("Error: height must be between 128 and 2048", file=sys.stderr)
        sys.exit(1)
    if not (5 <= args.steps <= 100):
        print("Error: steps must be between 5 and 100", file=sys.stderr)
        sys.exit(1)
    if args.seed is not None and not (0 <= args.seed <= 4294967295):
        print("Error: seed must be between 0 and 4294967295", file=sys.stderr)
        sys.exit(1)
    if not (0 <= args.cfg_scale <= 10):
        print("Error: cfg-scale must be between 0 and 10", file=sys.stderr)
        sys.exit(1)

    edit_image(
        prompt=args.prompt,
        image_path=args.image_path,
        width=args.width,
        height=args.height,
        steps=args.steps,
        seed=args.seed,
        cfg_scale=args.cfg_scale,
        negative_prompt=args.negative_prompt,
    )


if __name__ == "__main__":
    main()