#!/usr/bin/env python3
"""Capture the speed-change (aux) lanes ANCHORED to the parts themselves.

Hardcoded camera offsets are how you end up photographing empty air after the
geometry moves. Every shot here is framed from the LIVE position of the aux
lane / slip parts queried out of the DataModel at capture time, so the camera
cannot drift away from the thing under test.
"""
import shutil
import sys
import time

sys.path.insert(0, "/tmp")
import rbx  # noqa: E402

OUT = "/home/jesse/roblox/tokyo-wangan-racing/tools/shots"

# Ask Studio where the aux lanes actually are, then frame from that.
# Returns "x y z | x y z ..." anchors: west aux mid, east aux mid, and the
# two slip mouths.
ANCHOR_CODE = """
local map = workspace:FindFirstChild("Map")
if not map then return "NOMAP" end
local aux = {}
local mouths = {}
for _, d in ipairs(map:GetDescendants()) do
    if d:IsA("BasePart") then
        local ai = d.Name:match("^AuxLane_(%d+)$")
        if ai then table.insert(aux, {i = tonumber(ai), p = d.Position}) end
        for _, s in ipairs({"PAWestSlip", "PAEastSlip"}) do
            local k = d.Name:match("^" .. s .. "_Deck_(%d+)$")
            if k and tonumber(k) <= 1 then mouths[s] = d.Position end
        end
    end
end
if #aux == 0 then return "NOAUX" end
table.sort(aux, function(a, b) return a.i < b.i end)
-- Split the aux segments into two clusters (west run and east run) by index gap.
local groups, cur = {}, {aux[1]}
for k = 2, #aux do
    if aux[k].i - aux[k - 1].i > 3 then
        table.insert(groups, cur); cur = {}
    end
    table.insert(cur, aux[k])
end
table.insert(groups, cur)
local parts = {}
for _, g in ipairs(groups) do
    local mid = g[math.max(1, math.floor(#g / 2))].p
    table.insert(parts, ("%.1f %.1f %.1f"):format(mid.X, mid.Y, mid.Z))
end
for _, s in ipairs({"PAWestSlip", "PAEastSlip"}) do
    local m = mouths[s]
    if m then table.insert(parts, ("%.1f %.1f %.1f"):format(m.X, m.Y, m.Z)) end
end
return table.concat(parts, " | ")
"""

CAM_TPL = """
local cam = workspace.CurrentCamera
cam.CameraType = Enum.CameraType.Scriptable
cam.FieldOfView = {fov}
cam.CFrame = CFrame.lookAt(
    Vector3.new({px}, {py}, {pz}),
    Vector3.new({tx}, {ty}, {tz})
)
return "cam ok"
"""


def grab(name, px, py, pz, tx, ty, tz, fov=55):
    code = CAM_TPL.format(fov=fov, px=px, py=py, pz=pz, tx=tx, ty=ty, tz=tz)
    r = rbx.call("execute_luau", {"datamodel_type": "Edit", "code": code})
    if isinstance(r, dict) and r.get("isError"):
        print(f"{name}: CAM FAIL {r}")
        return None
    time.sleep(1.5)
    cap = rbx.call("screen_capture", {"capture_id": "ScreenCapture_1"}, wait=22, timeout=300)
    imgs = cap.get("images") if isinstance(cap, dict) else None
    if not imgs:
        print(f"{name}: CAPTURE FAIL {str(cap)[:160]}")
        return None
    dest = f"{OUT}/{name}.jpg"
    shutil.copy(imgs[0], dest)
    print(f"{name}: OK -> {dest}")
    return dest


def main():
    a = rbx.call("execute_luau", {"datamodel_type": "Edit", "code": ANCHOR_CODE},
                 wait=40, timeout=300)
    txt = a if isinstance(a, str) else str(a)
    if "NOMAP" in txt or "NOAUX" in txt:
        print("ANCHOR FAIL: " + txt[:200])
        return
    # Pull the coordinate triples out of whatever wrapper the bridge returns.
    import re
    trip = re.findall(r"(-?\d+\.\d+) (-?\d+\.\d+) (-?\d+\.\d+)", txt)
    if not trip:
        print("ANCHOR PARSE FAIL: " + txt[:300])
        return
    pts = [tuple(float(v) for v in t) for t in trip]
    print("anchors: " + str(pts))

    made = []
    for n, (x, y, z) in enumerate(pts[:2], start=1):
        # Oblique from outside the deck, looking back at the merge.
        d = grab(f"aux{n}_oblique", x + 260, y + 150, z + 260, x, y, z, fov=50)
        if d:
            made.append(d)
        # Top-down: the merge taper reads clearest from directly above.
        d = grab(f"aux{n}_topdown", x + 3, y + 430, z + 3, x, y, z, fov=42)
        if d:
            made.append(d)
        # Driver's-eye down the lane.
        d = grab(f"aux{n}_driver", x - 190, y + 16, z - 60, x + 120, y + 8, z + 40, fov=68)
        if d:
            made.append(d)

    print("\nMADE " + str(len(made)) + " shots")
    for m in made:
        print(m)


if __name__ == "__main__":
    main()