from __future__ import annotations

import base64
import io
import json
import os
import sys
import threading
import time
from pathlib import Path
from typing import List, Optional, Tuple

import click
import requests
from PIL import Image as PILImage

from . import config as cfg

__version__ = "1.0.1"


# ---------------------------------------------------------------------------
# Spinner
# ---------------------------------------------------------------------------

class _Spinner:
    """Non-blocking TTY spinner shown while waiting on the API."""

    _FRAMES = r"|/-\\"

    def __init__(self, message: str) -> None:
        self.message = message
        self._stop = threading.Event()
        self._thread = threading.Thread(target=self._run, daemon=True)

    def _run(self) -> None:
        i = 0
        while not self._stop.is_set():
            frame = self._FRAMES[i % len(self._FRAMES)]
            click.echo(f"\r{frame}  {self.message}", nl=False, err=True)
            self._stop.wait(0.1)
            i += 1
        click.echo("\r" + " " * (len(self.message) + 4) + "\r", nl=False, err=True)

    def __enter__(self) -> "_Spinner":
        if sys.stderr.isatty():
            self._thread.start()
        return self

    def __exit__(self, *_) -> None:
        self._stop.set()
        if self._thread.is_alive():
            self._thread.join()


# ---------------------------------------------------------------------------
# Internal helpers
# ---------------------------------------------------------------------------

def _require_credentials(conf: dict) -> Tuple[str, str]:
    account_id = conf.get("account_id", "").strip()
    api_token = conf.get("api_token", "").strip()
    if not account_id or not api_token:
        click.echo(
            "Error: Cloudflare credentials are not configured.\n"
            "Run:  imageflare config",
            err=True,
        )
        sys.exit(1)
    return account_id, api_token


def _prepare_image(image_path: str) -> Tuple[bytes, Tuple[int, int], Tuple[int, int]]:
    """Open, resize to ≤512×512 (CF requirement), return (png_bytes, orig, new)."""
    img = PILImage.open(image_path)
    original = img.size
    img.thumbnail((512, 512), PILImage.LANCZOS)
    new_size = img.size
    buf = io.BytesIO()
    img.save(buf, format="PNG")
    return buf.getvalue(), original, new_size


def _call_cf_api(
    conf: dict,
    model: str,
    prompt: str,
    *,
    images: Optional[List[str]] = None,
    width: int = 1024,
    height: int = 1024,
    seed: Optional[int] = None,
) -> bytes:
    account_id, api_token = _require_credentials(conf)

    url = (
        f"https://api.cloudflare.com/client/v4/accounts/{account_id}"
        f"/ai/run/{model}"
    )
    headers = {"Authorization": f"Bearer {api_token}"}

    # All fields are sent as multipart text parts using the (None, value) form.
    # Steps is intentionally omitted — it is fixed at 4 for klein models and
    # not a standard REST parameter for the other CF models.
    files: dict = {
        "prompt": (None, prompt),
        "width":  (None, str(width)),
        "height": (None, str(height)),
    }

    if seed is not None:
        files["seed"] = (None, str(seed))

    if images:
        for idx, img_path in enumerate(images[:4]):
            img_bytes, orig, resized = _prepare_image(img_path)
            if orig != resized:
                click.echo(
                    f"  Resized image {idx}: {orig[0]}×{orig[1]} → "
                    f"{resized[0]}×{resized[1]}",
                    err=True,
                )
            files[f"input_image_{idx}"] = (Path(img_path).name, img_bytes, "image/png")

    try:
        with _Spinner("Calling Cloudflare Workers AI…"):
            response = requests.post(url, headers=headers, files=files, timeout=180)
        response.raise_for_status()
    except KeyboardInterrupt:
        click.echo("\nCancelled.", err=True)
        sys.exit(130)
    except requests.exceptions.HTTPError as exc:
        click.echo(
            f"\nAPI error {exc.response.status_code}:\n{exc.response.text[:600]}",
            err=True,
        )
        sys.exit(1)
    except requests.exceptions.ConnectionError:
        click.echo("\nConnection failed. Check your internet connection.", err=True)
        sys.exit(1)
    except requests.exceptions.Timeout:
        click.echo("\nRequest timed out. Try again or use a faster model.", err=True)
        sys.exit(1)
    except requests.exceptions.RequestException as exc:
        click.echo(f"\nRequest failed: {exc}", err=True)
        sys.exit(1)

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

    # Some CF endpoints return raw binary
    if "image" in content_type:
        return response.content

    # Others wrap it: {"result": {"image": "<base64>"}, "success": true}
    try:
        payload = response.json()
        inner = payload.get("result", payload)
        raw = (inner.get("image") or inner.get("data")) if isinstance(inner, dict) else inner
        if raw:
            return base64.b64decode(raw)
    except (json.JSONDecodeError, Exception):
        pass

    click.echo(
        f"\nUnexpected API response (Content-Type: {content_type}):"
        f"\n{response.text[:500]}",
        err=True,
    )
    sys.exit(1)


def _save_output(data: bytes, output: Optional[str] = None) -> Path:
    dest = Path(output) if output else Path(os.getcwd()) / f"imageflare_{int(time.time())}.png"
    dest.parent.mkdir(parents=True, exist_ok=True)
    dest.write_bytes(data)
    return dest


def _print_result(dest: Path, data: bytes) -> None:
    size_kb = len(data) / 1024
    try:
        img = PILImage.open(io.BytesIO(data))
        dims = f"{img.width}×{img.height}"
    except Exception:
        dims = "unknown"
    click.echo(f"\n  Saved  {dest}  ({dims}, {size_kb:.1f} KB)")


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

@click.group()
@click.version_option(version=__version__, prog_name="imageflare")
def cli():
    """imageflare — Image editing and generation via Cloudflare Workers AI."""


# -- edit -------------------------------------------------------------------

@cli.command()
@click.argument("image_path", type=click.Path(exists=True, dir_okay=False))
@click.option("--prompt", "-p", required=True, help="Natural-language edit instruction.")
@click.option(
    "--ref", "-r", "refs",
    multiple=True,
    type=click.Path(exists=True, dir_okay=False),
    help="Additional reference image (repeat for up to 3 extra, 4 total).",
)
@click.option("--model", "-m", default=None, help="Model override (default from config).")
@click.option("--width",  type=int, default=1024, show_default=True, help="Output width in pixels.")
@click.option("--height", type=int, default=1024, show_default=True, help="Output height in pixels.")
@click.option("--seed",   type=int, default=None, help="Seed for reproducible results.")
@click.option("--output", "-o", default=None, metavar="FILE", help="Output path (default: imageflare_<ts>.png in cwd).")
@click.option("--open",  "open_after", is_flag=True, default=False, help="Open the result after saving.")
def edit(image_path, prompt, refs, model, width, height, seed, output, open_after):
    """Edit an image guided by a text prompt.

    IMAGE_PATH is the primary image to edit.  Use --ref to supply up to 3
    additional reference images (4 total).  Reference them in your prompt as
    'image 0', 'image 1', etc.

    \b
    Examples:
      imageflare edit photo.png --prompt "change the background to a beach"
      imageflare edit photo.png --ref style.png --prompt "style image 0 like image 1"
    """
    conf = cfg.load_config()
    model = model or conf.get("model", cfg.DEFAULT_MODEL)

    if len(refs) > 3:
        click.echo("Warning: Only 4 images total are supported; extra --ref images ignored.", err=True)

    all_images = [image_path] + list(refs[:3])

    ref_note = f"  (+{len(refs)} ref)" if refs else ""
    click.echo(f"  Model   {model}")
    click.echo(f"  Input   {image_path}{ref_note}")
    click.echo(f"  Prompt  {prompt}")

    data = _call_cf_api(
        conf, model, prompt,
        images=all_images, width=width, height=height, seed=seed,
    )
    dest = _save_output(data, output)
    _print_result(dest, data)

    if open_after:
        click.launch(str(dest))


# -- generate ---------------------------------------------------------------

@cli.command()
@click.option("--prompt", "-p", required=True, help="Text prompt for image generation.")
@click.option("--model", "-m", default=None, help="Model override (default from config).")
@click.option("--width",  type=int, default=1024, show_default=True, help="Output width in pixels.")
@click.option("--height", type=int, default=1024, show_default=True, help="Output height in pixels.")
@click.option("--seed",   type=int, default=None, help="Seed for reproducible results.")
@click.option("--output", "-o", default=None, metavar="FILE", help="Output path (default: imageflare_<ts>.png in cwd).")
@click.option("--open",  "open_after", is_flag=True, default=False, help="Open the result after saving.")
def generate(prompt, model, width, height, seed, output, open_after):
    """Generate a new image from a text prompt.

    \b
    Examples:
      imageflare generate --prompt "a sunset over mountains"
      imageflare generate --prompt "a cat" --width 512 --height 512 --seed 42
    """
    conf = cfg.load_config()
    model = model or conf.get("model", cfg.DEFAULT_MODEL)

    click.echo(f"  Model   {model}")
    click.echo(f"  Prompt  {prompt}")

    data = _call_cf_api(conf, model, prompt, width=width, height=height, seed=seed)
    dest = _save_output(data, output)
    _print_result(dest, data)

    if open_after:
        click.launch(str(dest))


# -- config -----------------------------------------------------------------

def _get_models(conf: dict, *, quiet: bool = False) -> list:
    """Return live model list from CF API, falling back to the static list."""
    account_id = conf.get("account_id", "").strip()
    api_token = conf.get("api_token", "").strip()
    if account_id and api_token:
        if not quiet:
            click.echo("  Fetching available models from Cloudflare…", err=True)
        models = cfg.fetch_image_models(account_id, api_token)
        if models:
            return models
        if not quiet:
            click.echo("  (Could not reach API — showing cached list)", err=True)
    # Fallback: static list
    return [{"id": m, "description": ""} for m in cfg.AVAILABLE_MODELS]


def _interactive_config() -> None:
    conf = cfg.load_config()
    config_path = cfg.get_config_path()

    click.echo("\n╔══════════════════════════════════════╗")
    click.echo("║     imageflare  ·  configuration     ║")
    click.echo("╚══════════════════════════════════════╝\n")
    click.echo(f"  Config file: {config_path}")
    click.echo("  Press Enter to keep the current value.\n")

    # Account ID
    current_id = conf.get("account_id", "")
    account_id = click.prompt(
        "Cloudflare Account ID",
        default=current_id or "",
        show_default=bool(current_id),
        prompt_suffix=": ",
    ).strip() or current_id

    # API Token (hidden input)
    current_token = conf.get("api_token", "")
    masked = ("*" * max(0, len(current_token) - 4)) + current_token[-4:] if current_token else ""
    token_label = f"Cloudflare API Token [{masked}]" if masked else "Cloudflare API Token"
    raw_token = click.prompt(
        token_label,
        default="",
        hide_input=True,
        show_default=False,
        prompt_suffix=": ",
    ).strip()
    api_token = raw_token or current_token

    # Model selection — fetch live list using the (possibly just-entered) creds
    live_conf = {"account_id": account_id, "api_token": api_token}
    models = _get_models(live_conf)
    current_model = conf.get("model", cfg.DEFAULT_MODEL)

    click.echo("\nAvailable models:\n")
    click.echo(f"  {'':3}Model")
    click.echo(f"  {'':3}{'-'*50}")
    for i, m in enumerate(models, 1):
        current_marker = "  ◀" if m["id"] == current_model else ""
        click.echo(f"  {i:2}.  {m['id']}{current_marker}")
    click.echo("")

    current_idx = next(
        (i for i, m in enumerate(models, 1) if m["id"] == current_model), 1
    )
    while True:
        raw = click.prompt("Select model", default=str(current_idx), prompt_suffix=": ").strip()
        if raw.isdigit() and 1 <= int(raw) <= len(models):
            chosen_model = models[int(raw) - 1]["id"]
            break
        click.echo(f"  Enter a number between 1 and {len(models)}.")

    cfg.save_config(
        account_id=account_id or None,
        api_token=api_token or None,
        model=chosen_model,
    )

    token_display = ("*" * max(0, len(api_token) - 4)) + api_token[-4:] if api_token else "(not set)"
    click.echo(f"\n  Saved to {config_path}")
    click.echo(f"  Account ID : {account_id or '(not set)'}")
    click.echo(f"  API Token  : {token_display}")
    click.echo(f"  Model      : {chosen_model}\n")


@cli.group(invoke_without_command=True)
@click.pass_context
def config(ctx):
    """Interactive config wizard (no subcommand), or use a subcommand directly."""
    if ctx.invoked_subcommand is None:
        _interactive_config()


@config.command("set")
@click.option("--account-id", default=None, help="Cloudflare Account ID.")
@click.option("--token",      default=None, help="Cloudflare API Token.")
@click.option("--model",      default=None, help="Default model.")
def config_set(account_id, token, model):
    """Set one or more configuration values non-interactively."""
    if not any([account_id, token, model]):
        click.echo("Nothing to update. Pass --account-id, --token, and/or --model.")
        return
    cfg.save_config(account_id=account_id, api_token=token, model=model)
    click.echo(f"  Saved to {cfg.get_config_path()}")


@config.command("show")
def config_show():
    """Display the current configuration."""
    conf = cfg.load_config()
    token = conf.get("api_token", "")
    masked = ("*" * max(0, len(token) - 4)) + token[-4:] if token else "(not set)"
    click.echo(f"  Config file : {cfg.get_config_path()}")
    click.echo(f"  Account ID  : {conf.get('account_id', '(not set)')}")
    click.echo(f"  API Token   : {masked}")
    click.echo(f"  Model       : {conf.get('model', cfg.DEFAULT_MODEL)}")


@config.command("models")
def config_models():
    """List available Cloudflare Workers AI Text-to-Image models."""
    conf = cfg.load_config()
    models = _get_models(conf)
    current = conf.get("model", cfg.DEFAULT_MODEL)
    click.echo("\nAvailable Text-to-Image models:\n")
    click.echo(f"  {'':3}Model")
    click.echo(f"  {'':3}{'-'*50}")
    for m in models:
        marker = "  *" if m["id"] == current else "   "
        click.echo(f"{marker}   {m['id']}")
    click.echo(f"\n  {len(models)} model(s)  |  * = current default\n")
