#!/usr/bin/env python3
"""Refresh stale low-resolution cached FetLife images from saved picture pages.

Dry-run by default. Use ``--apply`` to download replacements. The script only
rewrites existing local asset files, so DB asset URLs keep pointing at the same
paths.
"""
from __future__ import annotations

import argparse
from dataclasses import dataclass
from pathlib import Path
import sys

import httpx
from bs4 import BeautifulSoup

ROOT = Path(__file__).resolve().parents[1]
if str(ROOT) not in sys.path:
    sys.path.insert(0, str(ROOT))

from fetlife_collect import (  # noqa: E402
    ASSET_ROOT,
    BASE_URL,
    DEFAULT_HEADERS,
    RAW_ROOT,
    _fetlife_cdn_width,
    picture_cards,
)

IMAGE_SUFFIXES = {".jpg", ".jpeg", ".png", ".webp"}


@dataclass(frozen=True)
class AssetCandidate:
    fetish_id: str
    path: Path
    attachment_id: str
    width: int
    height: int
    source_url: str
    source_width: int


def jpeg_size(path: Path) -> tuple[int, int] | None:
    try:
        with path.open("rb") as fh:
            if fh.read(2) != b"\xff\xd8":
                return None
            while True:
                byte = fh.read(1)
                while byte and byte != b"\xff":
                    byte = fh.read(1)
                if not byte:
                    return None
                marker = fh.read(1)
                while marker == b"\xff":
                    marker = fh.read(1)
                if not marker:
                    return None
                marker_value = marker[0]
                if marker_value in (0xD8, 0xD9):
                    continue
                if marker_value == 0xDA:
                    return None
                raw_length = fh.read(2)
                if len(raw_length) != 2:
                    return None
                length = int.from_bytes(raw_length, "big")
                if length < 2:
                    return None
                if 0xC0 <= marker_value <= 0xCF and marker_value not in (0xC4, 0xC8, 0xCC):
                    data = fh.read(5)
                    if len(data) != 5:
                        return None
                    height = int.from_bytes(data[1:3], "big")
                    width = int.from_bytes(data[3:5], "big")
                    return width, height
                fh.seek(length - 2, 1)
    except OSError:
        return None


def png_size(path: Path) -> tuple[int, int] | None:
    try:
        with path.open("rb") as fh:
            if fh.read(8) != b"\x89PNG\r\n\x1a\n":
                return None
            length = int.from_bytes(fh.read(4), "big")
            chunk_type = fh.read(4)
            if length != 13 or chunk_type != b"IHDR":
                return None
            data = fh.read(8)
            if len(data) != 8:
                return None
            return int.from_bytes(data[0:4], "big"), int.from_bytes(data[4:8], "big")
    except OSError:
        return None


def image_size(path: Path) -> tuple[int, int] | None:
    suffix = path.suffix.lower()
    if suffix in {".jpg", ".jpeg"}:
        return jpeg_size(path)
    if suffix == ".png":
        return png_size(path)
    return None


def fetish_ids_from_assets(asset_root: Path) -> list[str]:
    if not asset_root.is_dir():
        return []
    return sorted(path.name for path in asset_root.iterdir() if path.is_dir())


def asset_files(asset_root: Path, fetish_id: str) -> list[Path]:
    root = asset_root / fetish_id
    if not root.is_dir():
        return []
    return sorted(path for path in root.iterdir() if path.is_file() and path.suffix.lower() in IMAGE_SUFFIXES)


def raw_picture_file(raw_root: Path, fetish_id: str) -> Path | None:
    path = raw_root / f"fetishes__{fetish_id}__pictures.html"
    return path if path.is_file() else None


def picture_cards_by_attachment(raw_root: Path, fetish_id: str) -> dict[str, dict[str, str]]:
    path = raw_picture_file(raw_root, fetish_id)
    if path is None:
        return {}
    soup = BeautifulSoup(path.read_text(errors="ignore"), "html.parser")
    return {card["attachment_id"]: card for card in picture_cards(soup)}


def collect_candidates(
    asset_root: Path,
    raw_root: Path,
    fetish_ids: list[str],
    min_long_edge: int,
    include_all: bool = False,
    allow_unknown_source: bool = False,
    candidate_limit: int = 0,
) -> tuple[list[AssetCandidate], dict[str, str]]:
    candidates: list[AssetCandidate] = []
    skipped: dict[str, str] = {}
    for fetish_id in fetish_ids:
        files = asset_files(asset_root, fetish_id)
        if not files:
            skipped[fetish_id] = "no cached image files"
            continue
        cards = picture_cards_by_attachment(raw_root, fetish_id)
        if not cards:
            skipped[fetish_id] = "no raw picture capture"
            continue
        matched = 0
        for path in files:
            size = image_size(path)
            if size is None:
                if not include_all:
                    continue
                width = height = 0
            else:
                width, height = size
                if not include_all and max(width, height) > min_long_edge:
                    continue
            attachment_id = path.stem
            card = cards.get(attachment_id)
            if not card:
                continue
            source_url = card["src"]
            source_width = _fetlife_cdn_width(source_url)
            if not allow_unknown_source and source_width and source_width <= max(width, height):
                continue
            if not allow_unknown_source and not source_width and not include_all:
                continue
            candidates.append(
                AssetCandidate(
                    fetish_id=fetish_id,
                    path=path,
                    attachment_id=attachment_id,
                    width=width,
                    height=height,
                    source_url=source_url,
                    source_width=source_width,
                )
            )
            matched += 1
            if candidate_limit > 0 and len(candidates) >= candidate_limit:
                return candidates, skipped
        if matched == 0 and fetish_id not in skipped:
            skipped[fetish_id] = "no replaceable low-resolution assets"
    return candidates, skipped


def refresh_candidate(client: httpx.Client, candidate: AssetCandidate) -> bool:
    headers = {**DEFAULT_HEADERS, "referer": f"{BASE_URL}/fetishes/{candidate.fetish_id}/pictures"}
    response = client.get(candidate.source_url, headers=headers)
    if response.status_code != 200 or not response.headers.get("content-type", "").startswith("image/"):
        return False
    candidate.path.write_bytes(response.content)
    return True


def parse_args() -> argparse.Namespace:
    parser = argparse.ArgumentParser(description=__doc__)
    parser.add_argument("--asset-root", type=Path, default=ASSET_ROOT)
    parser.add_argument("--raw-root", type=Path, default=RAW_ROOT)
    parser.add_argument("--fetish-id", action="append", default=[], help="limit to a fetish id; may be repeated")
    parser.add_argument("--min-long-edge", type=int, default=160, help="refresh assets at or below this long edge")
    parser.add_argument("--limit", type=int, default=0, help="max candidates/downloads; 0 means no limit")
    parser.add_argument("--all", action="store_true", help="consider every cached asset, not only low-resolution files")
    parser.add_argument(
        "--allow-unknown-source",
        action="store_true",
        help="allow replacement when the raw source URL has no cNNN width marker",
    )
    mode = parser.add_mutually_exclusive_group()
    mode.add_argument("--apply", action="store_true", help="download and overwrite matching local assets")
    mode.add_argument("--dry-run", action="store_true", help="report only (default)")
    return parser.parse_args()


def main() -> int:
    args = parse_args()
    if args.fetish_id:
        fetish_ids = sorted(set(str(fid) for fid in args.fetish_id))
    else:
        fetish_ids = fetish_ids_from_assets(args.asset_root)

    candidates, skipped = collect_candidates(
        asset_root=args.asset_root,
        raw_root=args.raw_root,
        fetish_ids=fetish_ids,
        min_long_edge=args.min_long_edge,
        include_all=args.all,
        allow_unknown_source=args.allow_unknown_source,
        candidate_limit=args.limit,
    )

    print(f"asset root: {args.asset_root}")
    print(f"raw root: {args.raw_root}")
    print(f"fetishes scanned: {len(fetish_ids)}")
    print(f"refresh candidates: {len(candidates)}")
    for candidate in candidates[:30]:
        src = f"c{candidate.source_width}" if candidate.source_width else "unknown source size"
        print(
            f"  {candidate.fetish_id}/{candidate.path.name}: "
            f"{candidate.width}x{candidate.height} -> {src} {candidate.source_url}"
        )
    if len(candidates) > 30:
        print(f"  ... {len(candidates) - 30} more")
    for fetish_id, reason in list(skipped.items())[:20]:
        print(f"  skipped {fetish_id}: {reason}")
    if len(skipped) > 20:
        print(f"  ... {len(skipped) - 20} more skipped")

    if not args.apply:
        print("dry-run only; pass --apply to refresh assets")
        return 0

    refreshed = 0
    with httpx.Client(timeout=30.0, follow_redirects=True) as client:
        for candidate in candidates:
            if refresh_candidate(client, candidate):
                refreshed += 1
    print(f"refreshed assets: {refreshed}")
    return 0


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