#!/usr/bin/env python3
"""
CE Image Editor — fal.ai (primary) + Google Nano Banana Pro (fallback)

Usage:
  # Edit an existing image with a text instruction
  python3 img-edit.py edit <image_path> "<instruction>" [--output <path>] [--provider fal|google]

  # Generate a new image from a text prompt
  python3 img-edit.py gen "<prompt>" [--output <path>] [--provider fal|google]

  # Compare: run the same edit/gen on both providers side by side
  python3 img-edit.py compare edit <image_path> "<instruction>" [--output-dir <dir>]
  python3 img-edit.py compare gen "<prompt>" [--output-dir <dir>]

Outputs PNG to --output or auto-names based on timestamp.
"""

import argparse
import base64
import httpx
import json
import os
import sys
import time
from datetime import datetime
from pathlib import Path

# Load env from tools/.env
ENV_FILE = Path(__file__).parent / ".env"
if ENV_FILE.exists():
    for line in ENV_FILE.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())

FAL_KEY = os.environ.get("FAL_KEY", "")
GEMINI_API_KEY = os.environ.get("GEMINI_API_KEY", "")

TIMEOUT = 120  # seconds


# ---------------------------------------------------------------------------
# fal.ai provider
# ---------------------------------------------------------------------------

def fal_edit(image_path: str, instruction: str) -> bytes:
    """Edit image via fal.ai Nano Banana Pro endpoint."""
    img_b64 = _img_to_b64(image_path)
    mime = _mime_type(image_path)

    resp = httpx.post(
        "https://fal.run/fal-ai/nano-banana-pro/edit",
        headers={"Authorization": f"Key {FAL_KEY}", "Content-Type": "application/json"},
        json={
            "image_url": f"data:{mime};base64,{img_b64}",
            "prompt": instruction,
        },
        timeout=TIMEOUT,
    )
    resp.raise_for_status()
    result = resp.json()

    # fal returns either a URL or base64 — handle both
    return _extract_image(result)


def fal_generate(prompt: str) -> bytes:
    """Generate image via fal.ai Nano Banana Pro endpoint."""
    resp = httpx.post(
        "https://fal.run/fal-ai/nano-banana-pro",
        headers={"Authorization": f"Key {FAL_KEY}", "Content-Type": "application/json"},
        json={
            "prompt": prompt,
            "image_size": {"width": 1080, "height": 1350},  # 4:5 social ratio
        },
        timeout=TIMEOUT,
    )
    resp.raise_for_status()
    result = resp.json()
    return _extract_image(result)


# ---------------------------------------------------------------------------
# Google Gemini / Nano Banana Pro provider
# ---------------------------------------------------------------------------

def google_edit(image_path: str, instruction: str) -> bytes:
    """Edit image via Google Gemini Nano Banana Pro."""
    from google import genai
    from google.genai import types

    client = genai.Client(api_key=GEMINI_API_KEY)
    img_bytes = Path(image_path).read_bytes()
    mime = _mime_type(image_path)

    response = client.models.generate_content(
        model="nano-banana-pro-preview",
        contents=[
            types.Part.from_bytes(data=img_bytes, mime_type=mime),
            instruction,
        ],
        config=types.GenerateContentConfig(
            response_modalities=["IMAGE", "TEXT"],
        ),
    )

    # Extract image from response
    for part in response.candidates[0].content.parts:
        if part.inline_data and part.inline_data.mime_type.startswith("image/"):
            return part.inline_data.data

    raise RuntimeError("Google returned no image in response")


def google_generate(prompt: str) -> bytes:
    """Generate image via Google Gemini Nano Banana Pro."""
    from google import genai
    from google.genai import types

    client = genai.Client(api_key=GEMINI_API_KEY)

    response = client.models.generate_content(
        model="nano-banana-pro-preview",
        contents=prompt,
        config=types.GenerateContentConfig(
            response_modalities=["IMAGE", "TEXT"],
        ),
    )

    for part in response.candidates[0].content.parts:
        if part.inline_data and part.inline_data.mime_type.startswith("image/"):
            return part.inline_data.data

    raise RuntimeError("Google returned no image in response")


# ---------------------------------------------------------------------------
# Helpers
# ---------------------------------------------------------------------------

def _img_to_b64(path: str) -> str:
    return base64.b64encode(Path(path).read_bytes()).decode()


def _mime_type(path: str) -> str:
    ext = Path(path).suffix.lower()
    return {
        ".png": "image/png",
        ".jpg": "image/jpeg",
        ".jpeg": "image/jpeg",
        ".webp": "image/webp",
    }.get(ext, "image/png")


def _extract_image(result: dict) -> bytes:
    """Extract image bytes from fal.ai response."""
    # fal typically returns {"images": [{"url": "..."}]} or {"image": {"url": "..."}}
    img_data = None

    if "images" in result and result["images"]:
        img_data = result["images"][0]
    elif "image" in result:
        img_data = result["image"]

    if img_data:
        if "url" in img_data:
            url = img_data["url"]
            if url.startswith("data:"):
                # data URI
                _, encoded = url.split(",", 1)
                return base64.b64decode(encoded)
            else:
                # HTTP URL — download it
                resp = httpx.get(url, timeout=TIMEOUT)
                resp.raise_for_status()
                return resp.content
        elif "content" in img_data:
            return base64.b64decode(img_data["content"])

    raise RuntimeError(f"Could not extract image from fal response: {json.dumps(result)[:500]}")


def _output_path(prefix: str, output: str | None, suffix: str = "") -> str:
    if output:
        return output
    ts = datetime.now().strftime("%Y-%m-%d-%H-%M-%S")
    return f"{prefix}{suffix}-{ts}.png"


def _save(data: bytes, path: str) -> str:
    Path(path).parent.mkdir(parents=True, exist_ok=True)
    Path(path).write_bytes(data)
    return str(Path(path).resolve())


# ---------------------------------------------------------------------------
# Runner with fallback
# ---------------------------------------------------------------------------

def run_with_fallback(primary_fn, fallback_fn, provider: str | None, *args):
    """Run primary provider, fall back to secondary on failure."""
    providers = {
        "fal": primary_fn,
        "google": fallback_fn,
    }

    if provider:
        # Explicit provider requested
        try:
            return providers[provider](*args), provider
        except Exception as e:
            print(f"[{provider}] failed: {e}", file=sys.stderr)
            raise

    # Default: Google NBP first, fal.ai fallback
    try:
        print("[google/nbp] attempting...", file=sys.stderr)
        return fallback_fn(*args), "google"
    except Exception as e:
        print(f"[google/nbp] failed: {e}", file=sys.stderr)
        print("[fal.ai] falling back...", file=sys.stderr)
        try:
            return primary_fn(*args), "fal"
        except Exception as e2:
            print(f"[fal.ai] also failed: {e2}", file=sys.stderr)
            raise RuntimeError(f"Both providers failed. google: {e} | fal: {e2}")


# ---------------------------------------------------------------------------
# CLI
# ---------------------------------------------------------------------------

def main():
    parser = argparse.ArgumentParser(description="CE Image Editor")
    sub = parser.add_subparsers(dest="command")

    # edit
    p_edit = sub.add_parser("edit", help="Edit an existing image")
    p_edit.add_argument("image", help="Path to source image")
    p_edit.add_argument("instruction", help="Edit instruction")
    p_edit.add_argument("--output", "-o", help="Output path")
    p_edit.add_argument("--provider", "-p", choices=["fal", "google"], help="Force provider")

    # gen
    p_gen = sub.add_parser("gen", help="Generate a new image")
    p_gen.add_argument("prompt", help="Text prompt")
    p_gen.add_argument("--output", "-o", help="Output path")
    p_gen.add_argument("--provider", "-p", choices=["fal", "google"], help="Force provider")

    # compare
    p_cmp = sub.add_parser("compare", help="Run on both providers")
    p_cmp.add_argument("mode", choices=["edit", "gen"])
    p_cmp.add_argument("args", nargs="+", help="image + instruction (edit) or prompt (gen)")
    p_cmp.add_argument("--output-dir", "-d", default=".")

    args = parser.parse_args()

    if args.command == "edit":
        img_data, provider = run_with_fallback(
            fal_edit, google_edit, args.provider, args.image, args.instruction
        )
        out = _save(img_data, _output_path("edit", args.output))
        print(json.dumps({"status": "ok", "provider": provider, "output": out}))

    elif args.command == "gen":
        img_data, provider = run_with_fallback(
            fal_generate, google_generate, args.provider, args.prompt
        )
        out = _save(img_data, _output_path("gen", args.output))
        print(json.dumps({"status": "ok", "provider": provider, "output": out}))

    elif args.command == "compare":
        if args.mode == "edit":
            image_path, instruction = args.args[0], " ".join(args.args[1:])
            results = {}
            for name, fn in [("fal", fal_edit), ("google", google_edit)]:
                try:
                    data = fn(image_path, instruction)
                    out = _save(data, str(Path(args.output_dir) / f"compare-{name}-{datetime.now().strftime('%H-%M-%S')}.png"))
                    results[name] = {"status": "ok", "output": out}
                    print(f"[{name}] done -> {out}", file=sys.stderr)
                except Exception as e:
                    results[name] = {"status": "error", "error": str(e)}
                    print(f"[{name}] failed: {e}", file=sys.stderr)
            print(json.dumps(results))

        elif args.mode == "gen":
            prompt = " ".join(args.args)
            results = {}
            for name, fn in [("fal", fal_generate), ("google", google_generate)]:
                try:
                    data = fn(prompt)
                    out = _save(data, str(Path(args.output_dir) / f"compare-{name}-{datetime.now().strftime('%H-%M-%S')}.png"))
                    results[name] = {"status": "ok", "output": out}
                    print(f"[{name}] done -> {out}", file=sys.stderr)
                except Exception as e:
                    results[name] = {"status": "error", "error": str(e)}
                    print(f"[{name}] failed: {e}", file=sys.stderr)
            print(json.dumps(results))

    else:
        parser.print_help()
        sys.exit(1)


if __name__ == "__main__":
    main()