#!/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 = """Make this incredibly photorealistic.
Highly stylized, striking, highly detailed photo of a young, lithe woman.
Age her by 5 years and make her a 20 year old, taller, slimmer, more slender, fitter version of herself.
Make her 6ft tall.

Give her realistic spotless skin texture.
Her face skin must also be realistic, without any makeup.
Make the background realistic and consistent with the lighting on her.

Give her an angular, not round, narrow, elongated neck and long face shape defined by a structured, clean jawline that tapers toward a firm, slightly rounded chin.
Do not make the head larger than the body, keep them in perfect proportion.
Give her high, prominent cheekbones giving her sculpted, chieseled face that taper down to a more narrow, delicate chin.
Make her really pretty and seductive.
Preserve unnatural lip colors and wetness.

Do not change the facial expression, emotion, hand and body pose, mouth, tongue and eyelid position.
Lower her eyes slightly to give her an alluring look.
Mouth slightly agape giving her a seductive look.
Give her a sharp, cunning gaze, a very slight pleasant smirk without changing her head pose.

Remove the drool.
Keep the pose and don't turn people around."""

SYSTEM_PROMPT = """
Turn this into a heavily stylized real life photo of a woman.
Age her by 5 years and make her a 20 year old, taller, slimmer, more slender, fitter version of herself.

Apply realistic textures without grain or noise.
Skin should be smooth, flawless, without pores, clean and without blemishes.

Keep the shine.
Make her taller, slimmer, slender. Don't make the head shorter.

Do not change the pose, gaze, emotion.

Her lips are parted, creating a soft, seductive expression.
"""

SYSTEM_PROMPT = """
restyle this as an edited, finished photo taken by sony a7 iv.
Photorealistic textures.
same color grading, same colors and lighting.
Realistic, flawless skin.
White people should have a pale skin.

keep the same ethnicity and facial features.

Make the children taller and more slender with slightly elongated feminine neck and higher cheekbones, give them smaller, slimmer more feminine head and tapered face; round, smooth chin.
age them by 5 years and make them look like 20 year old adults.

keep everything else about her the same.
same facial + eye expression (keep them as open as in the original), same pose and emotion, same tongue position.
same head and eye angle.

do not turn her face around.
same hair style, color, length.
same composition and crop.
keep the same makeup, do not change the colors of the lips, keep them glossy and wet.
"""
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 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())