#!/usr/bin/env python3
"""Bigme B6 bench rig server.

Serves the Bench/ folder over the LAN so the Bigme's browser can open it,
and exposes a tiny listing API so the pages auto-discover whatever files
you drop into art/ and seq/.

Run:  python3 serve.py
Then on the Bigme browser open:  http://<mac-lan-ip>:8000/
"""
import http.server
import json
import os
import re
import socket
import socketserver
import struct
import time
import urllib.request
from urllib.parse import urlparse, parse_qs

# PORT is overridable so a platform can inject one (App Runner, ECS, Elastic Beanstalk
# all do). Defaults to 8000 for the laptop-and-tablet setup.
PORT = int(os.environ.get("PORT", "8000"))
ROOT = os.path.dirname(os.path.abspath(__file__))
IMG_EXT = (".png", ".jpg", ".jpeg", ".gif", ".webp", ".svg", ".avif")
AUD_EXT = (".mp3", ".wav", ".m4a", ".ogg")

# Live voice runs on ElevenLabs Agents (STT + LLM + TTS in one session).
# The agent ID lives in voice.json; the API key, if used, comes from the
# environment only — never hardcode it here and never commit it.
ELEVEN_KEY = os.environ.get("ELEVENLABS_API_KEY", "")

# When set, /api/list loads zoo + motion catalogs from S3 on every request
# (indexes uploaded separately; art binaries stay on S3/CloudFront).
# Unset → scan local art/ and motion/ (laptop LAN demo).
ASSETS_BUCKET = os.environ.get("ASSETS_BUCKET", "").strip()
ASSETS_PREFIX = os.environ.get("ASSETS_PREFIX", "zoo").strip().strip("/")
AWS_REGION = (
    os.environ.get("AWS_REGION")
    or os.environ.get("AWS_DEFAULT_REGION")
    or "us-east-1"
)


ZOO_DIR = os.path.join(ROOT, "art", "CCZOO")
# Drop-anything folder: the filename IS the subject ("Great_White_Shark.png").
# Lets art from any other series be pulled in without the CCZOO naming convention.
ZOO_EXTRA_DIR = os.path.join(ROOT, "art", "zoo-extra")

# Files dropped into zoo-extra usually keep their production names
# ("SHRWCM09_14_Ebony_IceCream_Delivery_FSD_CCLING11_Shiba_Inu_Dog.png"), so the
# subject is recovered by scanning for a known animal. The LAST match wins,
# because the layer name at the end of these names describes the actual asset —
# "..._Duck_Horse_Park_Field_FSD_horse5" is a horse, not a duck.
ANIMAL_WORDS = [
    "dog", "puppy", "corgi", "shiba", "retriever", "poodle", "beagle",
    "cat", "kitten", "kitty", "tabby",
    "horse", "pony", "donkey", "mule", "cow", "cattle", "bull", "calf",
    "ox", "oxen", "pig", "piglet", "hog", "sheep", "lamb", "goat",
    "chicken", "hen", "rooster", "duck", "duckling", "goose", "turkey",
    "rabbit", "bunny", "hare", "mouse", "mice", "rat", "hamster",
    "squirrel", "chipmunk", "fox", "wolf", "deer", "moose", "elk",
    "tiger", "leopard", "cheetah", "giraffe", "zebra", "rhino", "monkey",
    "gorilla", "chimp", "kangaroo", "koala", "panda", "sloth", "camel",
    "penguin", "owl", "eagle", "hawk", "parrot", "flamingo", "swan", "crow",
    "snake", "python", "cobra", "lizard", "gecko", "iguana", "frog", "toad",
    "shark", "orca", "whale", "seahorse", "starfish", "stingray", "eel",
    "dinosaur", "tyrannosaurus", "trex", "rex", "triceratops", "raptor",
    "stegosaurus", "steg", "brachiosaurus", "coelophysis", "dino", "fossil",
    "butterfly", "dragonfly", "ant", "snail", "slug", "bat", "hedgehog",
    "puffin", "platypus", "worm", "otter", "badger", "raccoon", "walrus",
]
# a few need a friendlier display name than the word found in the filename
ANIMAL_CANON = {
    "puppy": "Dog", "corgi": "Dog", "shiba": "Dog", "retriever": "Dog",
    "poodle": "Dog", "beagle": "Dog", "kitten": "Cat", "kitty": "Cat",
    "tabby": "Cat", "mice": "Mouse", "bunny": "Rabbit", "oxen": "Ox",
    "cattle": "Cow", "calf": "Cow", "piglet": "Pig", "hog": "Pig",
    "lamb": "Sheep", "hen": "Chicken", "rooster": "Chicken",
    "duckling": "Duck", "chimp": "Chimpanzee", "trex": "Tyrannosaurus",
    # "T-Rex" normalises to the two words "t rex", so "rex" is the token that hits
    "rex": "Tyrannosaurus", "steg": "Stegosaurus", "dino": "Dinosaur",
    "fossil": "Dinosaur_Fossil",
}


def animal_from_filename(stem):
    low = " " + re.sub(r"[^a-z]+", " ", stem.lower()) + " "
    best, best_at = None, -1
    for w in ANIMAL_WORDS:
        # whole words only, or "fox" also matches the "ox" inside it
        at = max((m.start() for m in re.finditer(r"\b" + w + r"\b", low)), default=-1)
        if at > best_at:
            best, best_at = w, at
    if best is None:
        return None
    return ANIMAL_CANON.get(best, best.capitalize())
_ZOO_CACHE = None
_ZOO_STAMP = None


def _dir_stamp():
    """Cheap change detector so newly dropped art appears without a restart."""
    out = []
    for d in (ZOO_DIR, ZOO_EXTRA_DIR):
        try:
            out.append(os.stat(d).st_mtime)
        except OSError:
            out.append(0)
    return tuple(out)


def _png_size(path):
    try:
        with open(path, "rb") as f:
            head = f.read(26)
        if head[:8] != b"\x89PNG\r\n\x1a\n":
            return None
        return struct.unpack(">II", head[16:24])
    except Exception:
        return None


def list_zoo():
    """Index the CCZOO corpus by subject.

    The corpus is mostly layer fragments (textures, vignettes, backgrounds). A
    second CCZOO token in the filename marks a *named subject* illustration —
    e.g. ..._DIDDY_CCZOO_Axolotl.ai.png — which is what a page can be built
    from. Returns {"Axolotl": ["art/CCZOO/....png", ...]}, largest file first.
    """
    global _ZOO_CACHE, _ZOO_STAMP
    stamp = _dir_stamp()
    if _ZOO_CACHE is not None and stamp == _ZOO_STAMP:
        return _ZOO_CACHE
    _ZOO_STAMP = stamp
    subjects = {}
    if os.path.isdir(ZOO_DIR):
        pat = re.compile(r"_CCZOO\d*_(.+?)(?:\.ai)?\.png$", re.I)
        junk = re.compile(r"(MULTIPLY|OVERLAY|texture|shadow|vingette|grad|particle)", re.I)
        for f in os.listdir(ZOO_DIR):
            if not f.lower().endswith(".png"):
                continue
            m = pat.search(f)
            if not m or junk.search(m.group(1)):
                continue
            size = _png_size(os.path.join(ZOO_DIR, f))
            if not size or size[0] < 250 or size[1] < 250:
                continue
            name = m.group(1).replace(".ai-2", "").replace(",ai", "").strip("_")
            # DIDDY/OSA are the isolated-element cuts — cleaner as a page subject
            # than a layer lifted out of a full scene, so rank them first.
            rank = 1 if re.search(r"_(DIDDY|OSA)_", f) else 0
            subjects.setdefault(name, []).append((rank, size[0] * size[1], "art/CCZOO/" + f))

    # zoo-extra: filename is the subject, so anything from any series can be added.
    # Ranked above CCZOO so a deliberately-added asset always wins.
    if os.path.isdir(ZOO_EXTRA_DIR):
        for f in os.listdir(ZOO_EXTRA_DIR):
            if not f.lower().endswith((".png", ".svg", ".webp")):
                continue
            stem = re.sub(r"\.(ai|png|svg|webp)$", "", f, flags=re.I)
            stem = re.sub(r"\.(ai|png|svg|webp)$", "", stem, flags=re.I).strip()
            # Prefer the recognised animal so "CC_T-Rex_UPDATED" and "T_Rex.png"
            # land on the same subject; fall back to the filename for anything
            # outside the vocabulary (e.g. a hand-named "Axolotl.png").
            name = animal_from_filename(stem) or stem
            size = _png_size(os.path.join(ZOO_EXTRA_DIR, f)) or (9999, 9999)
            subjects.setdefault(name, []).append((2, size[0] * size[1], "art/zoo-extra/" + f))

    _ZOO_CACHE = {k: [u for _, _s, u in sorted(v, reverse=True)]
                  for k, v in sorted(subjects.items())}
    return _ZOO_CACHE


def voice_config():
    """Read voice.json for the ElevenLabs agent ID."""
    path = os.path.join(ROOT, "voice.json")
    try:
        with open(path) as f:
            cfg = json.load(f)
        agent = (cfg.get("agentId") or "").strip()
        # zoo.html can use its own animals agent; falls back to the main one
        return {"agentId": agent, "zooAgentId": (cfg.get("zooAgentId") or "").strip() or agent}
    except Exception:
        return {"agentId": "", "zooAgentId": ""}


def signed_url(agent_id):
    """Mint a short-lived signed URL for a private agent. Public agents don't need this."""
    if not ELEVEN_KEY:
        return {"error": "no_key"}
    url = ("https://api.elevenlabs.io/v1/convai/conversation/get-signed-url"
           "?agent_id=" + agent_id)
    req = urllib.request.Request(url, headers={"xi-api-key": ELEVEN_KEY})
    try:
        with urllib.request.urlopen(req, timeout=15) as r:
            return {"signedUrl": json.loads(r.read()).get("signed_url", "")}
    except Exception as e:
        return {"error": "signed url request failed: %s" % e}


def list_art():
    d = os.path.join(ROOT, "art")
    if not os.path.isdir(d):
        return []
    files = [f for f in os.listdir(d) if f.lower().endswith(IMG_EXT)]
    files.sort()
    return [{"name": f, "url": "art/" + f} for f in files]


def list_sequences():
    d = os.path.join(ROOT, "seq")
    out = []
    if not os.path.isdir(d):
        return out
    for name in sorted(os.listdir(d)):
        sub = os.path.join(d, name)
        if not os.path.isdir(sub):
            continue
        frames = [f for f in os.listdir(sub) if f.lower().endswith(IMG_EXT)]
        frames.sort()  # zero-padded frame_001.png sorts correctly
        if frames:
            out.append({
                "name": name,
                "count": len(frames),
                "frames": ["seq/" + name + "/" + f for f in frames],
            })
    return out


def list_motion():
    """Subject -> animated WebP loop, built offline by tools/build-motion.py.

    The manifest is the source of truth rather than a directory scan: it carries
    the pixel size the page needs to reserve space before the file arrives. Any
    entry whose file has gone missing is dropped, so a stale manifest degrades to
    the still art instead of a broken image on the tablet.
    """
    d = os.path.join(ROOT, "motion")
    try:
        with open(os.path.join(d, "index.json")) as f:
            index = json.load(f)
    except Exception:
        return {}
    out = {}
    for subject, meta in index.items():
        url = meta.get("url") or ""
        if url and os.path.isfile(os.path.join(ROOT, url)):
            out[subject] = meta
    return out


def _assets_key(*parts):
    bits = [ASSETS_PREFIX] if ASSETS_PREFIX else []
    bits.extend(p for p in parts if p)
    return "/".join(bits)


def s3_get_json(key):
    """Fetch a JSON object from ASSETS_BUCKET. Uses the task role / default creds."""
    try:
        import boto3
    except ImportError as e:
        raise RuntimeError(
            "boto3 is required when ASSETS_BUCKET is set (pip install boto3)"
        ) from e
    client = boto3.client("s3", region_name=AWS_REGION)
    body = client.get_object(Bucket=ASSETS_BUCKET, Key=key)["Body"].read()
    return json.loads(body)


def list_zoo_from_s3():
    """Precomputed subject map at {prefix}/zoo-index.json (see tools/build-zoo-index.py)."""
    data = s3_get_json(_assets_key("zoo-index.json"))
    if not isinstance(data, dict):
        raise ValueError("zoo-index.json must be a JSON object")
    return data


def list_motion_from_s3():
    """Motion manifest at {prefix}/motion/index.json — trust URLs (files live on CDN)."""
    data = s3_get_json(_assets_key("motion", "index.json"))
    if not isinstance(data, dict):
        raise ValueError("motion/index.json must be a JSON object")
    return data


def catalog_for_api():
    """Zoo + motion maps for /api/list — S3 catalogs or local disk."""
    if ASSETS_BUCKET:
        return list_zoo_from_s3(), list_motion_from_s3()
    return list_zoo(), list_motion()


def list_narration():
    """audio/page0.mp3 .. page3.mp3 -> {"0": "audio/page0.mp3", ...}"""
    d = os.path.join(ROOT, "audio")
    out = {}
    if not os.path.isdir(d):
        return out
    for f in sorted(os.listdir(d)):
        if not f.lower().endswith(AUD_EXT):
            continue
        stem = os.path.splitext(f)[0]
        if stem.startswith("page") and stem[4:].isdigit():
            out[stem[4:]] = "audio/" + f
    return out


class Handler(http.server.SimpleHTTPRequestHandler):
    def __init__(self, *a, **k):
        super().__init__(*a, directory=ROOT, **k)

    def _json(self, payload, code=200):
        blob = json.dumps(payload).encode()
        self.send_response(code)
        self.send_header("Content-Type", "application/json")
        self.send_header("Content-Length", str(len(blob)))
        self.send_header("Cache-Control", "no-store")
        self.end_headers()
        self.wfile.write(blob)

    def do_GET(self):
        parsed = urlparse(self.path)
        if parsed.path == "/api/list":
            cfg = voice_config()
            try:
                zoo, motion = catalog_for_api()
            except Exception as e:
                self._json({"error": "assets catalog failed: %s" % e}, 500)
                return
            self._json({
                "art": list_art(),
                "sequences": list_sequences(),
                "narration": list_narration(),
                "agentId": cfg["agentId"],
                "zooAgentId": cfg["zooAgentId"],
                "hasKey": bool(ELEVEN_KEY),
                "zoo": zoo,
                "motion": motion,
            })
            return
        if parsed.path == "/api/signed-url":
            cfg = voice_config()
            which = parse_qs(parsed.query).get("agent", ["main"])[0]
            aid = cfg["zooAgentId"] if which == "zoo" else cfg["agentId"]
            if not aid:
                self._json({"error": "no agentId in voice.json"}, 400)
                return
            self._json(signed_url(aid))
            return
        return super().do_GET()

    def do_POST(self):
        """Collect a diagnostic trace from the reading device.

        There are no devtools on the tablet, and the alignment stream is the one
        part of this system that cannot be observed from the laptop -- so the page
        posts a summary of what it actually received and it lands here as a file.
        """
        if urlparse(self.path).path != "/api/log":
            self._json({"error": "not found"}, 404)
            return
        try:
            n = int(self.headers.get("Content-Length") or 0)
            blob = self.rfile.read(min(n, 2_000_000))
            d = os.path.join(ROOT, "diag")
            os.makedirs(d, exist_ok=True)
            name = "align-%s.json" % time.strftime("%Y%m%d-%H%M%S")
            with open(os.path.join(d, name), "wb") as f:
                f.write(blob)
            self._json({"ok": True, "saved": name})
        except Exception as e:
            self._json({"error": str(e)}, 500)

    def end_headers(self):
        # Pages must never be cached or the device keeps running an old build.
        # Images stay cacheable on purpose — long sequences (100s of frames) must
        # replay from cache rather than re-download over wifi every loop.
        path = urlparse(self.path).path
        if path.endswith((".html", ".js", ".css", ".json")) or path.endswith("/"):
            self.send_header("Cache-Control", "no-store, must-revalidate")
        super().end_headers()

    def log_message(self, fmt, *args):
        pass  # keep the console quiet


def lan_ip():
    s = socket.socket(socket.AF_INET, socket.SOCK_DGRAM)
    try:
        s.connect(("8.8.8.8", 80))
        return s.getsockname()[0]
    except Exception:
        return "127.0.0.1"
    finally:
        s.close()


class Server(socketserver.ThreadingMixIn, socketserver.TCPServer):
    daemon_threads = True
    allow_reuse_address = True


if __name__ == "__main__":
    with Server(("0.0.0.0", PORT), Handler) as httpd:
        ip = lan_ip()
        print("\n  Bigme B6 bench rig is live.")
        print("  On the Bigme browser, open:\n")
        print(f"      http://{ip}:{PORT}/\n")
        print("  (Mac and Bigme must be on the same wifi.)")
        print("  Ctrl-C to stop.\n")
        try:
            httpd.serve_forever()
        except KeyboardInterrupt:
            print("\n  Stopped.")