#!/usr/bin/env python3
import argparse
import base64
import glob
import json
import os
import subprocess
import sys
from pathlib import Path

import requests

MODEL = "gpt-5.6-sol"
API_URL = "https://api.openai.com/v1/responses"

IMAGE_ROOT = Path("/home/bowerybay/public_html/wilgus-images/WILGUS")
AUTO_ROOT = Path("/home/bowerybay/public_html/wilgus-automation")
VIDEO_ROOT = Path("/home/bowerybay/public_html/wilgus-videos")
KEY_FILE = Path("/root/.config/openai/api_key")

SHEET_SIZE = 20
SHORTLIST_SIZE = 12
FINAL_SIZE = 6

SHORTLIST_PROMPT = """These are contact sheets for one Bethany Beach vacation rental.
Images are numbered globally from left to right, top to bottom: sheet 1 contains images
1-20, sheet 2 contains 21-40, and so on (the final sheet may contain fewer).

Nominate exactly 12 strong candidate photographs for a short real-estate advertising
video. This is a SCREENING pass, so preserve strong alternatives for the final pass.
Favor strong exterior/establishing images, the best main living spaces, distinctive
features, good bedrooms when worthwhile, and strong outdoor living, pool, view, or
amenity images when available. Adapt to what the property does best.

Reject weak/dark images, bathrooms unless exceptional, laundry/utility shots, maps,
floor plans, collages, and photographs whose primary subject is readable text,
signage, or logos. Incidental decorative text in an otherwise strong room is okay.
Avoid filling the shortlist with many nearly identical angles, but do not discard a
strong candidate merely because a similar image exists; the full-size final pass will
resolve close duplicates."""

FINAL_PROMPT = """Choose exactly SIX photographs from these full-size candidates for
a short vacation-rental advertising video.

Every selected image must add materially different visual information. Treat
photographs showing substantially the same room, deck, exterior, view, or feature as
duplicates even when angle, furniture, staging, lighting, or framing differs. Keep
only the strongest member of a duplicate group.

Favor a varied sequence that sells the rental: a strong exterior/setting when
available, the best main living area, distinctive spaces/features, a bedroom when
worthwhile, and outdoor living/view/amenities when available. Do not rigidly force
categories when the property has stronger alternatives.

Before returning the final six, compare every chosen image against every other chosen
image and remove near-duplicates. Incidental decorative text is acceptable, but avoid
images whose primary subject is signage, logos, maps, floor plans, or text."""

def die(msg):
    print(f"ERROR: {msg}", file=sys.stderr)
    sys.exit(1)

def read_key():
    if not KEY_FILE.exists():
        die(f"API key not found: {KEY_FILE}")
    key = KEY_FILE.read_text().strip()
    if not key:
        die("API key file is empty")
    return key

def data_url(path):
    data = Path(path).read_bytes()
    return "data:image/jpeg;base64," + base64.b64encode(data).decode("ascii")

def response_text(j):
    if isinstance(j.get("output_text"), str) and j["output_text"]:
        return j["output_text"]
    for item in j.get("output", []):
        for content in item.get("content", []):
            if content.get("type") == "output_text" and content.get("text"):
                return content["text"]
    return None

def call_selection(key, content, count, valid_numbers, schema_name):
    schema = {
        "type": "object",
        "properties": {
            "selected": {
                "type": "array",
                "items": {"type": "integer"},
                "minItems": count,
                "maxItems": count
            }
        },
        "required": ["selected"],
        "additionalProperties": False
    }
    payload = {
        "model": MODEL,
        "input": [{"role": "user", "content": content}],
        "text": {
            "format": {
                "type": "json_schema",
                "name": schema_name,
                "strict": True,
                "schema": schema
            }
        }
    }
    r = requests.post(
        API_URL,
        headers={"Authorization": f"Bearer {key}", "Content-Type": "application/json"},
        json=payload,
        timeout=300
    )
    if r.status_code != 200:
        die(f"OpenAI API HTTP {r.status_code}: {r.text[:1000]}")
    j = r.json()
    out = response_text(j)
    if not out:
        die("OpenAI response contained no output text")
    try:
        selected = json.loads(out)["selected"]
    except Exception as e:
        die(f"Could not parse structured response: {e}; output={out!r}")

    valid_set = set(valid_numbers)
    if (
        len(selected) != count
        or len(set(selected)) != count
        or any(type(n) is not int or n not in valid_set for n in selected)
    ):
        die(f"Invalid selection returned: {selected}")
    return selected

def make_contact_sheets(images, workdir):
    sheets_dir = workdir / "sheets"
    sheets_dir.mkdir(parents=True, exist_ok=True)
    for old in sheets_dir.glob("sheet-*.jpg"):
        old.unlink()

    sheets = []
    for start in range(0, len(images), SHEET_SIZE):
        chunk = images[start:start + SHEET_SIZE]
        sheet_num = len(sheets) + 1
        outfile = sheets_dir / f"sheet-{sheet_num}.jpg"
        cmd = [
            "montage",
            *[str(p) for p in chunk],
            "-thumbnail", "480x343",
            "-tile", "4x5",
            "-geometry", "+10+10",
            "-background", "white",
            str(outfile)
        ]
        subprocess.run(cmd, check=True)
        sheets.append(outfile)
    return sheets

def shortlist(key, sheets, total_images):
    content = [{"type": "input_text", "text": SHORTLIST_PROMPT}]
    for i, sheet in enumerate(sheets, 1):
        first = (i - 1) * SHEET_SIZE + 1
        last = min(i * SHEET_SIZE, total_images)
        content.append({
            "type": "input_text",
            "text": f"SHEET {i}: master images {first}-{last}"
        })
        content.append({
            "type": "input_image",
            "image_url": data_url(sheet),
            "detail": "high"
        })

    count = min(SHORTLIST_SIZE, total_images)
    return call_selection(
        key, content, count, range(1, total_images + 1), "wilgus_shortlist"
    )

def final_selection(key, images, candidates):
    content = [{"type": "input_text", "text": FINAL_PROMPT}]
    for n in candidates:
        content.append({"type": "input_text", "text": f"MASTER IMAGE {n}"})
        content.append({
            "type": "input_image",
            "image_url": data_url(images[n - 1]),
            "detail": "high"
        })
    return call_selection(
        key, content, FINAL_SIZE, candidates, "wilgus_final_selection"
    )

def make_video(selected_paths, output):
    if len(selected_paths) != 6:
        die("Video generator requires exactly six images")

    inputs = []
    for p in selected_paths:
        inputs += ["-i", str(p)]

    filters = [
        "[0:v]scale=1408:1006,zoompan=z='min(1+0.0015*on,1.09)':x='iw/2-(iw/zoom/2)':y='ih/2-(ih/zoom/2)':d=60:s=1280x914:fps=30,trim=duration=2,setpts=PTS-STARTPTS[v0]",
        "[1:v]scale=1408:1006,zoompan=z='max(1.09-0.0015*on,1)':x='iw/2-(iw/zoom/2)':y='ih/2-(ih/zoom/2)':d=60:s=1280x914:fps=30,trim=duration=2,setpts=PTS-STARTPTS[v1]",
        "[2:v]scale=1408:1006,zoompan=z='min(1+0.0012*on,1.072)':x='iw/2-(iw/zoom/2)':y='ih/2-(ih/zoom/2)':d=60:s=1280x914:fps=30,trim=duration=2,setpts=PTS-STARTPTS[v2]",
        "[3:v]scale=1408:1006,zoompan=z='max(1.072-0.0012*on,1)':x='iw/2-(iw/zoom/2)':y='ih/2-(ih/zoom/2)':d=60:s=1280x914:fps=30,trim=duration=2,setpts=PTS-STARTPTS[v3]",
        "[4:v]scale=1408:1006,zoompan=z='min(1+0.0012*on,1.072)':x='iw/2-(iw/zoom/2)':y='ih/2-(ih/zoom/2)':d=60:s=1280x914:fps=30,trim=duration=2,setpts=PTS-STARTPTS[v4]",
        "[5:v]scale=1408:1006,zoompan=z='max(1.09-0.0015*on,1)':x='iw/2-(iw/zoom/2)':y='ih/2-(ih/zoom/2)':d=60:s=1280x914:fps=30,trim=duration=2,setpts=PTS-STARTPTS[v5]",
        "[v0][v1]xfade=transition=fade:duration=0.35:offset=1.65[x1]",
        "[x1][v2]xfade=transition=fade:duration=0.35:offset=3.30[x2]",
        "[x2][v3]xfade=transition=fade:duration=0.35:offset=4.95[x3]",
        "[x3][v4]xfade=transition=fade:duration=0.35:offset=6.60[x4]",
        "[x4][v5]xfade=transition=fade:duration=0.35:offset=8.25[outv]"
    ]

    VIDEO_ROOT.mkdir(parents=True, exist_ok=True)
    cmd = [
        "ffmpeg", "-y",
        *inputs,
        "-filter_complex", ";".join(filters),
        "-map", "[outv]",
        "-t", "10.25",
        "-r", "30",
        "-c:v", "libx264",
        "-crf", "20",
        "-preset", "medium",
        "-pix_fmt", "yuv420p",
        "-movflags", "+faststart",
        str(output)
    ]
    subprocess.run(cmd, check=True)

def main():
    parser = argparse.ArgumentParser(
        description="Create one Wilgus property video using two-stage AI photo selection."
    )
    parser.add_argument("property_id")
    parser.add_argument("--no-video", action="store_true",
                        help="Run selection only; do not create MP4")
    args = parser.parse_args()

    prop = args.property_id.strip()
    source = IMAGE_ROOT / prop
    if not source.is_dir():
        die(f"Property folder not found: {source}")

    images = sorted(
        [p for p in source.iterdir()
         if p.is_file() and p.suffix.lower() in {".jpg", ".jpeg"}],
        key=lambda p: p.name.lower()
    )
    if len(images) < FINAL_SIZE:
        die(f"Only {len(images)} JPGs found; need at least {FINAL_SIZE}")

    workdir = AUTO_ROOT / "properties" / prop
    workdir.mkdir(parents=True, exist_ok=True)

    print(f"Property {prop}: {len(images)} source JPGs")
    print("1/4 Creating contact sheets...")
    sheets = make_contact_sheets(images, workdir)

    key = read_key()

    print(f"2/4 AI screening pass ({len(sheets)} contact sheet(s))...")
    candidates = shortlist(key, sheets, len(images))
    print("Shortlist:", ", ".join(map(str, candidates)))

    print("3/4 AI full-size final pass...")
    final = final_selection(key, images, candidates)
    print("Final:", ", ".join(map(str, final)))

    selected_files = [images[n - 1].name for n in final]
    record = {
        "property_id": prop,
        "model": MODEL,
        "source_count": len(images),
        "shortlist_master_numbers": candidates,
        "final_master_numbers": final,
        "final_files": selected_files
    }
    record_path = workdir / "selection.json"
    record_path.write_text(json.dumps(record, indent=2) + "\n")
    print("Files:")
    for n, name in zip(final, selected_files):
        print(f"  {n}: {name}")
    print(f"Selection record: {record_path}")

    if args.no_video:
        print("4/4 Video skipped (--no-video)")
        return

    output = VIDEO_ROOT / f"{prop}.mp4"
    print("4/4 Creating MP4...")
    make_video([images[n - 1] for n in final], output)
    print(f"DONE: {output}")

if __name__ == "__main__":
    main()