#!/usr/bin/env python3

import argparse
import base64
from concurrent.futures import ThreadPoolExecutor, as_completed
import hashlib
import json
import os
from pathlib import Path
import random
import time
import urllib.error
import urllib.request


API_KEY_FILE = Path("src/main/scala/net/eagle0/common/llm_integration/api_keys.txt")
IMAGE_ENDPOINT = "https://api.openai.com/v1/images/generations"

APPEARANCE_PROFILES = [
    "West African features, very dark brown skin",
    "East African features, deep brown skin",
    "North African features, olive-brown skin",
    "Arabian Peninsula features, warm medium-brown skin",
    "Persian features, olive skin",
    "Turkic Central Asian features, tan skin",
    "Mongolian features, golden-tan skin",
    "North Indian features, medium-brown skin",
    "South Indian features, dark brown skin",
    "Bengali features, warm brown skin",
    "Southeast Asian features, medium tan skin",
    "Vietnamese features, light golden skin",
    "Chinese features, light olive skin",
    "Korean features, fair skin",
    "Japanese features, fair-to-light olive skin",
    "Indigenous Andean features, copper-brown skin",
    "Mesoamerican features, medium brown skin",
    "Native North American features, warm brown skin",
    "Maori or Polynesian features, medium brown skin",
    "Aboriginal Australian features, dark brown skin",
    "Afro-Caribbean features, dark brown skin",
    "Afro-Latin features, medium-dark brown skin",
    "Brazilian mixed-heritage features, medium brown skin",
    "Mexican mestizo features, tan skin",
    "Mediterranean European features, olive skin",
    "Iberian features, light olive skin",
    "Italian features, olive skin",
    "Greek features, olive skin",
    "French features, light skin",
    "Irish features, fair skin with freckles",
    "Scottish features, ruddy fair skin",
    "Nordic features, very fair skin",
    "Slavic features, fair-to-light olive skin",
    "Ashkenazi Jewish features, light olive skin",
    "Armenian features, olive skin",
    "Ethiopian highland features, medium-dark brown skin",
    "Moroccan Amazigh features, tan olive skin",
    "Egyptian features, warm tan skin",
    "Filipino features, medium golden-brown skin",
    "Pacific Islander features, warm brown skin",
]

AGE_PROFILES = [
    "young adult",
    "adult",
    "seasoned middle-aged",
    "older veteran",
]

FACE_PROFILES = [
    "broad, sturdy face",
    "narrow, severe face",
    "round, watchful face",
    "angular, weathered face",
    "scarred but composed face",
    "plain, practical face",
    "commanding, heavy-browed face",
    "tired but alert face",
]

HAIR_PROFILES = [
    "cropped dark hair",
    "close-cropped gray hair",
    "braided black hair",
    "shaved head",
    "curly dark hair pulled back",
    "straight black hair tied back",
    "auburn hair streaked with gray",
    "white hair cut short",
    "dark beard and close hair",
    "neat graying beard",
    "coiled natural hair",
    "long dark hair bound under a cap",
]


def api_key():
    env_value = os.environ.get("OPENAI_API_KEY")
    if env_value:
        return env_value

    if API_KEY_FILE.exists():
        for line in API_KEY_FILE.read_text().splitlines():
            line = line.strip()
            if not line or line.startswith("#") or "=" not in line:
                continue
            key, value = line.split("=", 1)
            if key.strip() == "openai_api_key":
                value = value.strip()
                if value and not value.startswith("your_"):
                    return value

    raise RuntimeError(
        "OpenAI API key not found. Set OPENAI_API_KEY or add openai_api_key to "
        f"{API_KEY_FILE}."
    )


def parse_args():
    parser = argparse.ArgumentParser(
        description="Generate Warden hero headshots and update generated_heroes.tsv image paths."
    )
    parser.add_argument(
        "--hero-tsv",
        default="src/main/resources/net/eagle0/eagle/generated_heroes.tsv",
        help="Generated heroes TSV to read and update.",
    )
    parser.add_argument(
        "--base-image-dir",
        default=str(Path.home() / "Documents/headshots"),
        help="Local headshot root that sync_headshots.sh uploads to S3.",
    )
    parser.add_argument(
        "--model",
        default="gpt-image-2",
        help="OpenAI image generation model.",
    )
    parser.add_argument(
        "--quality",
        default="low",
        choices=("low", "medium", "high", "auto"),
        help="Image generation quality.",
    )
    parser.add_argument(
        "--size",
        default="1024x1024",
        help="Generated image size.",
    )
    parser.add_argument(
        "--limit",
        type=int,
        default=0,
        help="Generate at most this many missing images; 0 means all missing images.",
    )
    parser.add_argument(
        "--dry-run",
        action="store_true",
        help="Print planned paths without calling the API or writing files.",
    )
    parser.add_argument(
        "--overwrite",
        action="store_true",
        help="Regenerate images even if the target PNG already exists.",
    )
    parser.add_argument(
        "--concurrency",
        type=int,
        default=4,
        help="Number of parallel image generation requests.",
    )
    return parser.parse_args()


def read_rows(path):
    lines = path.read_text().splitlines()
    fieldnames = lines[0].split("\t")
    rows = []
    for line in lines[1:]:
        fields = line.split("\t")
        row = dict(zip(fieldnames, fields))
        row["_fields"] = fields
        rows.append(row)
    return fieldnames, rows


def write_rows(path, fieldnames, rows):
    image_path_index = fieldnames.index("image_path")
    lines = ["\t".join(fieldnames)]
    for row in rows:
        fields = list(row["_fields"])
        fields[image_path_index] = row["image_path"]
        lines.append("\t".join(fields))
    path.write_text("\n".join(lines) + "\n")


def image_gender(row, warden_index):
    gender = row["gender"]
    if gender in ("male", "female"):
        return gender

    digest = hashlib.sha256(row["name"].encode("utf-8")).digest()
    return "male" if (digest[0] + warden_index) % 2 == 0 else "female"


def image_filename(row, warden_index):
    digest = hashlib.sha256(row["name"].encode("utf-8")).hexdigest()[:16]
    return f"{warden_index:04d}_{digest}.png"


def choose_profile(items, row, warden_index, salt):
    digest = hashlib.sha256(f"{salt}:{warden_index}:{row['name']}".encode("utf-8")).digest()
    return items[int.from_bytes(digest[:4], "big") % len(items)]


def prompt_for(row, warden_index):
    prime_notes = [
        f"strength {row['strength']}",
        f"agility {row['agility']}",
        f"constitution {row['constitution']}",
        f"wisdom {row['wisdom']}",
        f"charisma {row['charisma']}",
    ]
    appearance = choose_profile(APPEARANCE_PROFILES, row, warden_index, "appearance")
    age = choose_profile(AGE_PROFILES, row, warden_index, "age")
    face = choose_profile(FACE_PROFILES, row, warden_index, "face")
    hair = choose_profile(HAIR_PROFILES, row, warden_index, "hair")
    return (
        "Fantasy strategy game hero portrait, square headshot composition. "
        f"Create a distinctive medieval Warden named {row['name']}. "
        f"Gender presentation: {row['gender']}. "
        f"Appearance: {age}, {appearance}, {face}, {hair}. "
        f"Personality traits: {row['personality']}. "
        "The character is a jailer, keeper of prisoners, and battlefield custodian: "
        "stern keys, austere authority, prison-watch uniform, no modern clothing. "
        "Show a resilient, high-constitution figure with weathered endurance and moral judgment. "
        f"Stat cues: {', '.join(prime_notes)}. "
        "Painted realistic fantasy style, expressive face, shoulders visible, neutral background, "
        "soft dramatic lighting, no text, no logo, no frame, no watermark."
    )


def request_image(key, model, quality, size, prompt):
    body = {
        "model": model,
        "prompt": prompt,
        "size": size,
        "quality": quality,
        "n": 1,
        "output_format": "png",
    }
    request = urllib.request.Request(
        IMAGE_ENDPOINT,
        data=json.dumps(body).encode("utf-8"),
        headers={
            "Authorization": f"Bearer {key}",
            "Content-Type": "application/json",
        },
        method="POST",
    )
    with urllib.request.urlopen(request, timeout=300) as response:
        response_body = json.loads(response.read().decode("utf-8"))

    data = response_body["data"][0]
    if "b64_json" in data:
        return base64.b64decode(data["b64_json"])

    if "url" in data:
        with urllib.request.urlopen(data["url"], timeout=300) as response:
            return response.read()

    raise RuntimeError("Image response did not include b64_json or url")


def generate_with_retries(key, args, prompt, attempts=5):
    for attempt in range(1, attempts + 1):
        try:
            return request_image(key, args.model, args.quality, args.size, prompt)
        except urllib.error.HTTPError as error:
            body = error.read().decode("utf-8", errors="replace")
            retryable = error.code in (408, 409, 429, 500, 502, 503, 504)
            if not retryable or attempt == attempts:
                raise RuntimeError(f"OpenAI image request failed: {error.code} {body}") from error
        except (TimeoutError, urllib.error.URLError):
            if attempt == attempts:
                raise

        time.sleep(min(60, 2**attempt) + random.random())


def main():
    args = parse_args()
    hero_tsv = Path(args.hero_tsv)
    base_image_dir = Path(args.base_image_dir).expanduser()
    fieldnames, rows = read_rows(hero_tsv)

    key = None if args.dry_run else api_key()
    warden_index = 0
    generation_tasks = []

    for row in rows:
        if row["profession"] != "warden":
            continue

        warden_index += 1
        gender = image_gender(row, warden_index)
        rel_path = f"warden/{gender}/{image_filename(row, warden_index)}"
        target_path = base_image_dir / rel_path
        row["image_path"] = rel_path

        if target_path.exists() and not args.overwrite:
            print(f"exists {rel_path}", flush=True)
            continue

        generation_tasks.append(
            (row["name"], rel_path, target_path, prompt_for(row, warden_index))
        )

    if args.dry_run:
        print(f"Dry run planned {warden_index} Warden image paths", flush=True)
    else:
        if args.limit:
            generation_tasks = generation_tasks[: args.limit]

        for _, _, target_path, _ in generation_tasks:
            target_path.parent.mkdir(parents=True, exist_ok=True)

        def generate_task(task):
            name, rel_path, target_path, prompt = task
            print(f"generate {rel_path} for {name}", flush=True)
            image_bytes = generate_with_retries(key, args, prompt)
            if not image_bytes.startswith(b"\x89PNG\r\n\x1a\n"):
                raise RuntimeError(f"Generated image for {name} is not a PNG")
            target_path.write_bytes(image_bytes)
            return rel_path

        generated_count = 0
        if generation_tasks:
            with ThreadPoolExecutor(max_workers=max(1, args.concurrency)) as executor:
                futures = [executor.submit(generate_task, task) for task in generation_tasks]
                for future in as_completed(futures):
                    rel_path = future.result()
                    generated_count += 1
                    print(
                        f"wrote {rel_path} ({generated_count}/{len(generation_tasks)})",
                        flush=True,
                    )

        write_rows(hero_tsv, fieldnames, rows)
        print(f"Updated {hero_tsv}; generated {generated_count} missing images", flush=True)


if __name__ == "__main__":
    main()
