"""
Scan TwitchDownloader chat JSONs for roster logins, download distinct profile images
(commenter.logo) as JPEGs under Story/pfp/, and print markdown snippets (or patch files).

Usage:
  python Assets/Scripts/pfp_from_chat.py [--chat-root PATH] [--story-root PATH] [--no-write-md]

Defaults:
  chat-root: MIXER_TWITCH_CHAT env, else Windows Synology path used in this project.
  Writes JPEGs under Story/pfp/<slug>/ and inserts a marked section into each roster .md.
  Also fills Story/Cameo-Creatures.md from CAMEO_ROSTER (see Assets/Scripts/pfp_from_chat.py).
"""
from __future__ import annotations

import argparse
import gzip
import json
import os
import re
import sys
from collections import defaultdict
from dataclasses import dataclass
from datetime import datetime
from pathlib import Path
from urllib.parse import urlparse
from urllib.request import Request, urlopen

# slug -> list of twitch logins (lowercase)
ROSTER: dict[str, list[str]] = {
    "adrian": ["axe7adrian"],
    "noname": ["noname106668"],
    "agate": ["loonyagate"],
    "notorious-rooster": ["notorious_rooster"],
    "ubear": ["verify52w", "sky_city_2013", "ancientmalgru"],
    "azure": ["guidingflyer530", "actuallynotazure"],
    "starboy": ["starboy_journeys"],
    "jenni": ["jennimilano"],
    "heart-and-mind": ["heart_cccc", "brush_colourful"],
    "beanie": ["beaniee__"],
    "raincloud": ["raincloudthedragon"],
}

MD_FILES: dict[str, str] = {
    "adrian": "Adrian.md",
    "noname": "Noname.md",
    "agate": "Agate.md",
    "notorious-rooster": "NotoriousRooster.md",
    "ubear": "Ubear.md",
    "azure": "Azure.md",
    "starboy": "Starboy.md",
    "jenni": "Jenni.md",
    "heart-and-mind": "Heart-and-Mind.md",
    "beanie": "Beanie.md",
    "raincloud": "RaincloudTheDragon.md",
}

# Grouped cameo handles (see Story/Cameo-Creatures.md). Slug = folder under Story/pfp/<slug>/.
CAMEO_ROSTER: dict[str, list[str]] = {
    "cameo-experimenta1ic3": ["branndongames", "experimenta1ic3"],
    "cameo-noncritical": ["noncriticalmother", "noncriticalgamingttv"],
    "cameo-pirate": ["pirate_protogen"],
    "cameo-foxy": ["foxy_fnaf5_ucn"],
    "cameo-gymrat": ["basedgymrat"],
    "cameo-miclbero": ["miclbero"],
    "cameo-queen": ["cameoqueen86"],
    "cameo-rayne": ["rayne8856"],
    "cameo-bd": ["bd_cum_lube"],
}

CAMEO_MD = "Cameo-Creatures.md"

DEFAULT_CHAT_ROOT = os.environ.get(
    "MIXER_TWITCH_CHAT",
    r"C:\Users\Nathan\SynologyDrive\YouTube\Streams\MixerTwitch",
)


def _avatar_fingerprint(logo_url: str) -> str | None:
    """Stable id for the same uploaded image (ignores 70x70 vs 300x300, etc.)."""
    if not logo_url:
        return None
    path = urlparse(logo_url).path
    base = re.sub(r"-\d+x\d+\.(png|jpe?g|webp)$", "", path, flags=re.IGNORECASE)
    return base or path


def _best_logo_url(logo_url: str) -> str:
    """Prefer a larger Twitch CDN size for download."""
    if not logo_url:
        return logo_url
    u = logo_url
    for small, large in (("70x70", "300x300"), ("36x36", "300x300"), ("50x50", "300x300")):
        if small in u:
            u = u.replace(small, large)
            break
    return u


def _parse_created_at(raw) -> datetime | None:
    if raw is None:
        return None
    if isinstance(raw, (int, float)):
        return datetime.utcfromtimestamp(raw)
    s = str(raw).replace("Z", "+00:00")
    try:
        return datetime.fromisoformat(s)
    except ValueError:
        return None


def open_chat_json(path: Path) -> dict | None:
    try:
        with path.open("rb") as f:
            magic = f.read(4)
        if magic[:2] == b"\x1f\x8b":
            with gzip.open(path, "rt", encoding="utf-8") as f:
                return json.load(f)
        with path.open("r", encoding="utf-8") as f:
            return json.load(f)
    except (OSError, json.JSONDecodeError) as e:
        print(f"skip {path}: {e}", file=sys.stderr)
        return None


@dataclass
class SeenAvatar:
    fingerprint: str
    first_at: datetime
    url: str


def collect_avatars(
    chat_root: Path, roster: dict[str, list[str]]
) -> dict[tuple[str, str], dict[str, SeenAvatar]]:
    """
    Returns map (slug, login) -> fingerprint -> SeenAvatar (earliest message wins per fp).
    """
    login_to_slug: dict[str, str] = {}
    for slug, logins in roster.items():
        for lg in logins:
            login_to_slug[lg.lower()] = slug

    # (slug, login) -> fp -> SeenAvatar
    acc: dict[tuple[str, str], dict[str, SeenAvatar]] = defaultdict(dict)

    files = sorted(chat_root.glob("*/chat/*.json"))
    if not files:
        print(f"No JSON under {chat_root}/*/chat/", file=sys.stderr)
        return acc

    for jp in files:
        data = open_chat_json(jp)
        if not data:
            continue
        for c in data.get("comments") or []:
            com = c.get("commenter") or {}
            login = (com.get("name") or com.get("login") or "").strip().lower()
            if not login or login not in login_to_slug:
                continue
            logo = com.get("logo")
            if not logo or not str(logo).startswith("http"):
                continue
            fp = _avatar_fingerprint(str(logo))
            if not fp:
                continue
            slug = login_to_slug[login]
            key = (slug, login)
            created = _parse_created_at(c.get("created_at"))
            if created is None:
                created = datetime.min
            cur = acc[key].get(fp)
            url = _best_logo_url(str(logo))
            if cur is None or created < cur.first_at:
                acc[key][fp] = SeenAvatar(fp, created, url)

    return acc


def download_as_jpeg(url: str, dest: Path) -> bool:
    dest.parent.mkdir(parents=True, exist_ok=True)
    req = Request(url, headers={"User-Agent": "BattleRoyale-pfp-fetch/1.0"})
    try:
        with urlopen(req, timeout=60) as resp:
            data = resp.read()
    except OSError as e:
        print(f"download failed {url}: {e}", file=sys.stderr)
        return False

    try:
        from io import BytesIO

        from PIL import Image

        im = Image.open(BytesIO(data)).convert("RGB")
        im.save(dest, "JPEG", quality=92)
    except ImportError:
        # No Pillow: write raw if already jpeg
        if data[:2] == b"\xff\xd8":
            dest.write_bytes(data)
        else:
            print("Install Pillow for PNG/WebP -> JPEG: pip install pillow", file=sys.stderr)
            return False
    except OSError as e:
        print(f"jpeg encode {url}: {e}", file=sys.stderr)
        return False
    return True


SECTION_START = "<!-- pfp:chat-exports -->"
SECTION_END = "<!-- /pfp:chat-exports -->"


def build_section_md(images: list[tuple[str, str]]) -> str:
    """images: list of (caption, relative_path from Story/)"""
    lines = [SECTION_START, "", "## Profile pictures", ""]
    for caption, rp in images:
        lines.append(f"**{caption}**")
        lines.append("")
        lines.append(f"![]({rp})")
        lines.append("")
    lines.append(SECTION_END)
    return "\n".join(lines) + "\n"


def inject_section(content: str, section_md: str) -> str:
    if SECTION_START in content and SECTION_END in content:
        pre, _, rest = content.partition(SECTION_START)
        _, _, post = rest.partition(SECTION_END)
        return pre.rstrip() + "\n\n" + section_md + "\n" + post.lstrip()
    # After first heading line
    lines = content.splitlines()
    if not lines:
        return section_md + content
    out = [lines[0], "", section_md]
    if len(lines) > 1 and lines[1].strip():
        out.append("")
    out.extend(lines[1:])
    return "\n".join(out)


def main() -> int:
    ap = argparse.ArgumentParser()
    ap.add_argument("--chat-root", type=Path, default=Path(DEFAULT_CHAT_ROOT))
    ap.add_argument("--story-root", type=Path, default=Path(__file__).resolve().parents[1] / "Story")
    ap.add_argument(
        "--no-write-md",
        action="store_true",
        help="Only download images; do not edit Story/*.md",
    )
    args = ap.parse_args()

    if not args.chat_root.is_dir():
        print(f"Chat root not found: {args.chat_root}", file=sys.stderr)
        print("Set MIXER_TWITCH_CHAT or pass --chat-root.", file=sys.stderr)
        return 1

    full_roster: dict[str, list[str]] = {**ROSTER, **CAMEO_ROSTER}
    acc = collect_avatars(args.chat_root, full_roster)
    pfp_root = args.story_root / "pfp"

    # Flatten per slug for multi-login: still per-login files
    for slug, md_name in MD_FILES.items():
        story_rel = f"pfp/{slug}"
        section_images: list[tuple[str, str]] = []

        display = Path(MD_FILES[slug]).stem.replace("-", " ")

        if slug == "heart-and-mind":
            order = [("Heart", "heart_cccc"), ("Mind", "brush_colourful")]
            for label, login in order:
                key = (slug, login)
                avs = acc.get(key, {})
                ordered = sorted(avs.values(), key=lambda x: x.first_at)
                for i, av in enumerate(ordered):
                    fn = f"{login}_{i}.jpg"
                    rel = f"{story_rel}/{fn}"
                    dest = pfp_root / slug / fn
                    if download_as_jpeg(av.url, dest):
                        cap = f"{label} (`{login}`)"
                        if len(ordered) > 1:
                            cap = f"{cap} — {i + 1}"
                        section_images.append((cap, rel))
        else:
            logins = ROSTER[slug]
            for login in logins:
                key = (slug, login)
                avs = acc.get(key, {})
                ordered = sorted(avs.values(), key=lambda x: x.first_at)
                if not ordered:
                    continue
                for i, av in enumerate(ordered):
                    fn = f"{login}_{i}.jpg"
                    rel = f"{story_rel}/{fn}"
                    dest = pfp_root / slug / fn
                    if not download_as_jpeg(av.url, dest):
                        continue
                    if len(logins) > 1:
                        cap = f"`{login}`"
                    else:
                        cap = display
                    if len(ordered) > 1:
                        cap = f"{cap} — {i + 1}"
                    section_images.append((cap, rel))

        if not section_images:
            print(f"No avatars found for {slug}", file=sys.stderr)
            continue

        section_md = build_section_md(section_images)
        md_path = args.story_root / md_name
        if not args.no_write_md and md_path.is_file():
            text = md_path.read_text(encoding="utf-8")
            md_path.write_text(inject_section(text, section_md), encoding="utf-8")
            print(f"updated {md_path}")

    # Cameo creatures: one combined section in Cameo-Creatures.md
    cameo_images: list[tuple[str, str]] = []
    for slug in CAMEO_ROSTER:
        logins = CAMEO_ROSTER[slug]
        story_rel = f"pfp/{slug}"
        for login in logins:
            key = (slug, login)
            avs = acc.get(key, {})
            ordered = sorted(avs.values(), key=lambda x: x.first_at)
            if not ordered:
                print(f"No avatars in exports for cameo `{login}` ({slug})", file=sys.stderr)
                continue
            for i, av in enumerate(ordered):
                fn = f"{login}_{i}.jpg"
                rel = f"{story_rel}/{fn}"
                dest = pfp_root / slug / fn
                if not download_as_jpeg(av.url, dest):
                    continue
                cap = f"`{login}`"
                if len(ordered) > 1:
                    cap = f"{cap} — {i + 1}"
                cameo_images.append((cap, rel))

    if cameo_images:
        cameo_md = build_section_md(cameo_images)
        cameo_path = args.story_root / CAMEO_MD
        if not args.no_write_md and cameo_path.is_file():
            cameo_path.write_text(
                inject_section(cameo_path.read_text(encoding="utf-8"), cameo_md),
                encoding="utf-8",
            )
            print(f"updated {cameo_path}")
    elif not args.no_write_md:
        print("No cameo avatars downloaded (check exports and logins).", file=sys.stderr)

    return 0


if __name__ == "__main__":
    raise SystemExit(main())