#!/usr/bin/env -S uv run --script
# /// script
# dependencies = ["httpx", "pillow"]
# ///
import argparse
import base64
import json
import logging
import os
import subprocess
import tempfile
import time
from concurrent.futures import ThreadPoolExecutor, as_completed
from dataclasses import dataclass
from io import BytesIO
from pathlib import Path
from typing import Any

import httpx
from PIL import Image, ImageOps

from sort_images import read_dims

MODEL_ID = "flux-2-klein-9b"
# LOCAL_MODEL_ID = "black-forest-labs/FLUX.2-klein-9B"
# MODEL_ID = "z-image-turbo"
# MODEL_ID = "qwen-image"


SYSTEM_PROMPT = """
Artistic super photorealistic conversion.
extend the background horizontally but keep the people.
# Subtle chiaroscuro lighting, not too dark.
85mm telephoto lens, f/2.8, shallow depth of field.
Soft-focus highlights, atmospheric bloom.
Kodak Portra 400 cool color palette.
#Preserve the original lighting and slight moody lighting. Not underexposed, not overexposed.
High dynamic range with a focus on rich textures.
balanced exposure.
same color grading and LUT.

remove noise and compression artifacts.
she has spotless, supple, shiny skin without splotches.
#pale skin, goth make-up.

#reduce musculature.
#keep the original facial expression, emotion, hand and body pose.

# make her prettier, give her hourglass figure, a satisfied look, voluminus hair.

#parted lips, seductive gaze, subtly lowered eyelids.
#make her face look like a 25 year old, raised cheekbones, tapered face shape, pointy chin.
#avoid making her look like a child and maintain adult proportions.
#give her a subtle mischievous, confident smile, avoid neutral face.
#make her face slimmer.
preserve the original skin color.

# do not change the colors of the lips, keep them wet.
remove all logos and watermarks.
#reflective outfit,
#glossy, reflective outfit,
clingy, skin-tight clothes.
"""


def without_comments(s: str) -> str:
    return "\n".join(
        line for line in s.splitlines() if not line.strip().startswith("#")
    )


def parse_args() -> argparse.Namespace:
    parser = argparse.ArgumentParser(
        description="Send one or more images to NanoGPT for photorealistic processing."
    )
    parser.add_argument(
        "image_paths",
        nargs="+",
        help="One or more local image paths to send",
    )
    parser.add_argument(
        "--extra-prompt",
        default="",
        help="Extra prompt text appended to the system prompt",
    )
    parser.add_argument(
        "--prompt",
        default="",
        help="Replace the system prompt with this custom prompt (instead of appending to it)",
    )
    parser.add_argument(
        "--provider",
        choices=("nanogpt", "local"),
        default=os.getenv("IMAGE_PROVIDER", "nanogpt"),
        help="Image provider backend (default from IMAGE_PROVIDER or nanogpt)",
    )
    parser.add_argument(
        "--base-url",
        default=os.getenv("LOCAL_IMAGE_API_BASE", "http://127.0.0.1:6006"),
        help="Local provider base URL (default from LOCAL_IMAGE_API_BASE)",
    )
    return parser.parse_args()


@dataclass(frozen=True)
class TransformResult:
    image_path: Path
    output_path: Path | None = None
    remaining: float | None = None


class NanoGPT:
    def __init__(self, api_key: str) -> None:
        self.api_key = api_key
        self._client = httpx.Client(
            timeout=180.0,
            headers={
                "Authorization": f"Bearer {self.api_key}",
                "Content-Type": "application/json",
            },
        )

    def transform(
        self, image_path: Path, output_path: Path, prompt: str
    ) -> TransformResult:
        if not image_path.exists():
            raise FileNotFoundError(f"Image not found: {image_path}")

        image_w, image_h = read_dims(image_path)
        aspect_ratio = image_w / image_h
        aspect_ratio = max(aspect_ratio, 2 / 3)
        target_w, target_h = 1024, int(1024 / aspect_ratio)
        size = f"{target_w}x{target_h}"

        image_data_url = _build_data_url(image_path)
        payload = {
            "model": MODEL_ID,
            "prompt": prompt,
            "imageDataUrl": image_data_url,
            "response_format": "b64_json",
            "n": 1,
            "seed": int(time.time()),
            "size": size,
        }
        response = self._client.post(
            "https://nano-gpt.com/v1/images/generations", json=payload
        )
        if not response.is_success:
            raise RuntimeError(
                f"API request failed with status {response.status_code}: {response.text}"
            )
        result = response.json()

        data = result.get("data", [])
        if not data:
            raise RuntimeError(f"No data in response for {image_path}")

        item = data[0]

        if "b64_json" in item:
            _save_b64_image(output_path, item["b64_json"])
        elif "url" in item:
            _download_image(self._client, output_path, item["url"])
        else:
            raise RuntimeError(f"Unknown response format for {image_path}")

        return TransformResult(
            image_path=image_path,
            output_path=output_path,
            remaining=result.get("remainingBalance"),
        )


class LocalFlux:
    def __init__(self, base_url: str) -> None:
        self.base_url = base_url.rstrip("/")
        self._client = httpx.Client(timeout=240.0)

    def transform(
        self, image_path: Path, output_path: Path, prompt: str
    ) -> TransformResult:
        if not image_path.exists():
            raise FileNotFoundError(f"Image not found: {image_path}")

        image_w, image_h = read_dims(image_path)
        aspect_ratio = image_w / image_h
        target_w, target_h = 1024, int(1024 / aspect_ratio)
        size = f"{target_w}x{target_h}"

        image_data_url = _build_data_url(image_path)
        payload = {
            "model": LOCAL_MODEL_ID,
            "prompt": prompt,
            "imageDataUrl": image_data_url,
            "response_format": "b64_json",
            "guidance_scale": 1.0,
            "n": 1,
            "seed": int(time.time()),
            "size": size,
        }
        response = self._client.post(
            f"{self.base_url}/v1/images/generations", json=payload
        )
        if not response.is_success:
            raise RuntimeError(
                f"Local API request failed with status {response.status_code}: {response.text}"
            )
        result = response.json()

        data = result.get("data", [])
        if not data:
            raise RuntimeError(f"No data in local response for {image_path}")

        item = data[0]
        if "b64_json" in item:
            _save_b64_image(output_path, item["b64_json"])
        elif "url" in item:
            _download_image(self._client, output_path, item["url"])
        else:
            raise RuntimeError(f"Unknown local response format for {image_path}")

        return TransformResult(
            image_path=image_path,
            output_path=output_path,
            remaining=result.get("remainingBalance"),
        )


def _build_data_url(image_path: Path) -> str:
    with Image.open(image_path) as image:
        image = ImageOps.exif_transpose(image)
        image.thumbnail((3000, 3000), Image.Resampling.LANCZOS)
        if image.mode not in {"RGB", "L"}:
            image = image.convert("RGB")
        buffer = BytesIO()
        image.save(buffer, format="JPEG", quality=95, optimize=True)
    encoded = base64.b64encode(buffer.getvalue()).decode("utf-8")
    return f"data:image/jpeg;base64,{encoded}"


def _save_b64_image(out_path: Path, b64_data: str) -> None:
    out_path.write_bytes(base64.b64decode(b64_data))


def _download_image(client: httpx.Client, out_path: Path, url: str) -> None:
    response = client.get(url, follow_redirects=True)
    response.raise_for_status()
    out_path.write_bytes(response.content)


def _next_output_path(image_path: Path) -> Path:
    base = image_path.with_name(f"{image_path.stem}-edit.jpg")
    if not base.exists():
        return base
    counter = 2
    while True:
        candidate = image_path.with_name(f"{image_path.stem}-edit{counter}.jpg")
        if not candidate.exists():
            return candidate
        counter += 1


def find_custom_prompts(file_paths: list[Path]) -> dict[Path, str]:
    custom_prompt_path = Path(r"~/Downloads/prompts.jsonl").expanduser()
    if not custom_prompt_path.exists():
        return {}

    stem_to_path: dict[str, Path] = {p.stem: p for p in file_paths}
    prompts: dict[Path, str] = {}
    with custom_prompt_path.open() as f:
        for line in f:
            line = line.strip()
            if not line:
                continue
            try:
                entry = json.loads(line)
                filename = entry.get("filename")
                prompt = entry.get("prompt")
                if filename and prompt and filename in stem_to_path:
                    prompts[stem_to_path[filename]] = prompt
            except json.JSONDecodeError:
                pass

    return prompts


def main() -> int:
    args = parse_args()
    provider = args.provider

    extra_prompt = args.extra_prompt.strip()
    if extra_prompt:
        prompt = f"{SYSTEM_PROMPT}\n\n{extra_prompt}"
    elif args.prompt.strip():
        prompt = args.prompt.strip()
    else:
        prompt = SYSTEM_PROMPT

    if provider == "nanogpt":
        api_key = os.getenv("NANOGPT_API_KEY")
        if not api_key:
            raise SystemExit("NANOGPT_API_KEY is not set")
        client = NanoGPT(api_key)
    else:
        client = LocalFlux(args.base_url)
    image_paths = [Path(p) for p in args.image_paths]
    custom_prompts = find_custom_prompts(image_paths)
    got_error = False
    with ThreadPoolExecutor() as executor:

        def _submit(image_path: Path) -> TransformResult:
            effective_prompt = without_comments(prompt)

            custom_prompt = custom_prompts.get(image_path)
            if custom_prompt:
                print(f"Using custom prompt for {image_path}")
                effective_prompt = custom_prompt

            with tempfile.TemporaryDirectory() as temp_dir:
                temp_path = Path(temp_dir) / image_path.name
                result = client.transform(
                    image_path,
                    temp_path,
                    prompt=effective_prompt,
                )

                output_path = _next_output_path(image_path)
                os.replace(temp_path, output_path)
                return TransformResult(
                    image_path=result.image_path,
                    output_path=output_path,
                    remaining=result.remaining,
                )

        futures = {
            executor.submit(_submit, image_path): image_path
            for image_path in image_paths
        }
        for future in as_completed(futures):
            try:
                outcome = future.result()
            except Exception as exc:
                image_path = futures[future]
                print(f"[error] {image_path} ({exc})")
                got_error = True
                continue

            image_path = outcome.image_path
            output_path = outcome.output_path
            print(f"[ok] {image_path} -> {output_path}")
            remaining = outcome.remaining
            if remaining is not None:
                print(f"[balance] {remaining}")

    return 1 if got_error else 0


if __name__ == "__main__":
    raise SystemExit(main())