#!/usr/bin/env python3
"""Read-only PixelRAG-style visual evidence capture worker.

This Python worker is the production-friendly counterpart of the Node worker:
it can dry-run without Playwright and only imports Playwright when actual capture
is requested.
"""

from __future__ import annotations

import argparse
import json
import math
from datetime import datetime, timezone
from pathlib import Path
from typing import Any
from urllib.parse import urlparse


CAPTURE_POLICY = "read_only_pixelrag_visual_capture_artifact_v1"
MANIFEST_POLICY = "read_only_pixelrag_visual_evidence_manifest_v1"
DEFAULT_OUTPUT_DIR = "runtime_artifacts/pixelrag_visual_evidence"
DEFAULT_TIMEOUT_MS = 30000
DEFAULT_SETTLE_MS = 500
DEFAULT_MAX_TILES = 80
FORBIDDEN_CONTEXT_HEADER_TOKENS = ("cookie", "authorization", "token", "secret", "key")
ALLOWED_PLATFORM_HOSTS = {
    "momo": {"m.momoshop.com.tw", "www.momoshop.com.tw"},
    "pchome": {"24h.pchome.com.tw", "ecshweb.pchome.com.tw", "ecapi-cdn.pchome.com.tw"},
    "shopee_tw": {"shopee.tw", "mall.shopee.tw"},
    "coupang_tw": {"www.tw.coupang.com", "tw.coupang.com"},
    "yahoo_shopping_tw": {"tw.buy.yahoo.com"},
    "etmall_tw": {"www.etmall.com.tw", "etmall.com.tw"},
    "friday_tw": {"shopping.friday.tw", "ec-m.shopping.friday.tw", "ec-w.shopping.friday.tw"},
    "rakuten_tw": {"www.rakuten.com.tw", "rakuten.com.tw"},
    "market_intel": set(),
    "external_market": set(),
}


def _positive_int(value: Any, fallback: int) -> int:
    try:
        parsed = int(value)
    except (TypeError, ValueError):
        return fallback
    return parsed if parsed > 0 else fallback


def _safe_name(value: Any) -> str:
    text = "".join(ch if ch.isalnum() else "_" for ch in str(value or "").lower())
    text = "_".join(part for part in text.split("_") if part)
    return text or "unknown"


def _load_manifest(args: argparse.Namespace) -> dict[str, Any]:
    if args.manifest_file:
        return json.loads(Path(args.manifest_file).read_text(encoding="utf-8"))
    return json.loads(args.manifest_json)


def _validate_manifest(manifest: dict[str, Any], *, require_allowed_host: bool) -> list[str]:
    errors: list[str] = []
    if manifest.get("policy") != MANIFEST_POLICY:
        errors.append(f"manifest policy must be {MANIFEST_POLICY}")
    target = manifest.get("capture_target") or {}
    raw_url = str(target.get("url") or "")
    parsed = urlparse(raw_url)
    if parsed.scheme not in {"http", "https"} or not parsed.netloc:
        errors.append("capture_target.url must be an absolute http(s) URL")
    if parsed.username or parsed.password:
        errors.append("URL credentials are not allowed")
    platform = str(target.get("platform") or "")
    allowed_hosts = ALLOWED_PLATFORM_HOSTS.get(platform)
    if (
        require_allowed_host
        and parsed.netloc
        and allowed_hosts
        and parsed.hostname not in allowed_hosts
    ):
        errors.append(f"host {parsed.hostname} is not allowed for platform {platform}")
    if not platform:
        errors.append("capture_target.platform is required")
    if not target.get("crawler"):
        errors.append("capture_target.crawler is required")
    if not target.get("trigger_reason"):
        errors.append("capture_target.trigger_reason is required")
    return errors


def _build_tile_plan(
    *,
    width: int,
    height: int,
    tile_width: int,
    tile_height: int,
    max_tiles: int,
) -> dict[str, Any]:
    tiles_x = max(1, math.ceil(width / tile_width))
    tiles_y = max(1, math.ceil(height / tile_height))
    tiles: list[dict[str, int]] = []
    for row in range(tiles_y):
        for col in range(tiles_x):
            if len(tiles) >= max_tiles:
                break
            x = col * tile_width
            y = row * tile_height
            tiles.append({
                "index": len(tiles),
                "row": row,
                "col": col,
                "x": x,
                "y": y,
                "width": min(tile_width, max(1, width - x)),
                "height": min(tile_height, max(1, height - y)),
            })
    planned = tiles_x * tiles_y
    return {
        "tiles_x": tiles_x,
        "tiles_y": tiles_y,
        "planned_tile_count": planned,
        "emitted_tile_count": len(tiles),
        "capped": len(tiles) < planned,
        "tiles": tiles,
    }


def _safe_public_browser_context(manifest: dict[str, Any]) -> dict[str, Any]:
    raw_context = manifest.get("public_browser_context") or {}
    raw_headers = raw_context.get("extra_http_headers") or {}
    headers: dict[str, str] = {}
    if isinstance(raw_headers, dict):
        for key, value in raw_headers.items():
            clean_key = str(key or "").strip()
            lower_key = clean_key.lower()
            if not clean_key:
                continue
            if any(token in lower_key for token in FORBIDDEN_CONTEXT_HEADER_TOKENS):
                continue
            headers[clean_key] = str(value or "")
    return {
        "context_policy": "public_empty_browser_context_no_login",
        "locale": str(raw_context.get("locale") or "zh-TW"),
        "timezone_id": str(
            raw_context.get("timezone_id") or raw_context.get("timezoneId") or "Asia/Taipei"
        ),
        "extra_http_headers": headers,
        "credentialed_session_allowed": False,
        "storage_state_allowed": False,
        "raw_cookie_or_session_read_allowed": False,
        "login_allowed": False,
    }


def _output_paths(manifest: dict[str, Any], output_dir: str) -> dict[str, Path]:
    target = manifest.get("capture_target") or {}
    manifest_id = manifest.get("manifest_id") or _safe_name(
        f"{target.get('platform')}-{target.get('crawler')}-{target.get('url')}"
    )[:20]
    root = Path(output_dir).resolve() / _safe_name(target.get("platform")) / manifest_id
    return {
        "root": root,
        "receipt": root / "capture_receipt.json",
        "screenshot": root / "fullpage.png",
        "tiles_dir": root / "tiles",
    }


def _base_artifact(
    manifest: dict[str, Any],
    args: argparse.Namespace,
    errors: list[str],
) -> dict[str, Any]:
    viewport = manifest.get("viewport") or {}
    page_size = manifest.get("page_size") or {}
    tile_size = manifest.get("tile_size") or {}
    width = _positive_int(page_size.get("width"), _positive_int(viewport.get("width"), 1440))
    height = _positive_int(page_size.get("height"), _positive_int(viewport.get("height"), 950))
    tile_width = _positive_int(tile_size.get("width"), 512)
    tile_height = _positive_int(tile_size.get("height"), 512)
    paths = _output_paths(manifest, args.output_dir)
    return {
        "success": not errors,
        "policy": CAPTURE_POLICY,
        "status": "rejected" if errors else ("dry_run_ready" if args.dry_run else "capture_pending"),
        "generated_at": datetime.now(timezone.utc).isoformat(),
        "errors": errors,
        "manifest_id": manifest.get("manifest_id") or "",
        "manifest_policy": manifest.get("policy") or "",
        "capture_target": manifest.get("capture_target") or {},
        "controlled_apply": {
            "db_write": False,
            "secret_read": False,
            "github_dependency": False,
            "production_price_write": False,
            "artifact_write": bool(not args.dry_run and not errors),
        },
        "planned_output": {
            "root": str(paths["root"]),
            "receipt": str(paths["receipt"]),
            "screenshot": str(paths["screenshot"]),
            "tiles_dir": str(paths["tiles_dir"]),
        },
        "viewport": {
            "name": viewport.get("name") or "desktop-1440",
            "width": _positive_int(viewport.get("width"), 1440),
            "height": _positive_int(viewport.get("height"), 950),
        },
        "page_size": {"width": width, "height": height},
        "tile_size": {"width": tile_width, "height": tile_height},
        "tile_plan": _build_tile_plan(
            width=width,
            height=height,
            tile_width=tile_width,
            tile_height=tile_height,
            max_tiles=_positive_int(args.max_tiles, DEFAULT_MAX_TILES),
        ),
        "files": [],
        "public_browser_context": _safe_public_browser_context(manifest),
    }


def _capture(
    manifest: dict[str, Any],
    args: argparse.Namespace,
    artifact: dict[str, Any],
) -> dict[str, Any]:
    from playwright.sync_api import sync_playwright

    paths = _output_paths(manifest, args.output_dir)
    paths["tiles_dir"].mkdir(parents=True, exist_ok=True)
    target = manifest.get("capture_target") or {}
    viewport = artifact["viewport"]
    timeout_ms = _positive_int(args.timeout * 1000, DEFAULT_TIMEOUT_MS)

    with sync_playwright() as playwright:
        browser = playwright.chromium.launch(headless=True)
        public_context = _safe_public_browser_context(manifest)
        context_kwargs = {
            "ignore_https_errors": True,
            "viewport": {
                "width": viewport["width"],
                "height": viewport["height"],
            },
            "locale": public_context["locale"],
            "timezone_id": public_context["timezone_id"],
        }
        if public_context["extra_http_headers"]:
            context_kwargs["extra_http_headers"] = public_context["extra_http_headers"]
        context = browser.new_context(**context_kwargs)
        page = context.new_page()
        try:
            response = page.goto(target["url"], wait_until=args.wait_until, timeout=timeout_ms)
            try:
                page.wait_for_selector("body", timeout=min(timeout_ms, 5000))
            except Exception:
                pass
            if args.settle_ms > 0:
                page.wait_for_timeout(args.settle_ms)
            metrics = page.evaluate(
                """() => {
                    const root = document.documentElement;
                    const body = document.body;
                    return {
                        final_url: location.href,
                        title: document.title,
                        scroll_width: Math.max(root.scrollWidth, body.scrollWidth, window.innerWidth),
                        scroll_height: Math.max(root.scrollHeight, body.scrollHeight, window.innerHeight)
                    };
                }"""
            )
            page.screenshot(path=str(paths["screenshot"]), full_page=True)
            tile_plan = _build_tile_plan(
                width=_positive_int(metrics.get("scroll_width"), artifact["page_size"]["width"]),
                height=_positive_int(metrics.get("scroll_height"), artifact["page_size"]["height"]),
                tile_width=artifact["tile_size"]["width"],
                tile_height=artifact["tile_size"]["height"],
                max_tiles=_positive_int(args.max_tiles, DEFAULT_MAX_TILES),
            )
            files: list[dict[str, Any]] = [
                {"kind": "fullpage_screenshot", "path": str(paths["screenshot"])},
            ]
            for tile in tile_plan["tiles"]:
                file_path = paths["tiles_dir"] / f"tile_{tile['index']:03d}.png"
                page.screenshot(path=str(file_path), clip={
                    "x": tile["x"],
                    "y": tile["y"],
                    "width": tile["width"],
                    "height": tile["height"],
                })
                files.append({"kind": "tile", "path": str(file_path), "tile": tile})
            captured = {
                **artifact,
                "success": True,
                "status": "captured",
                "http_status": response.status if response else 0,
                "page_metrics": metrics,
                "page_size": {
                    "width": _positive_int(metrics.get("scroll_width"), artifact["page_size"]["width"]),
                    "height": _positive_int(metrics.get("scroll_height"), artifact["page_size"]["height"]),
                },
                "tile_plan": tile_plan,
                "files": files,
            }
            paths["receipt"].write_text(json.dumps(captured, ensure_ascii=False, indent=2) + "\n", encoding="utf-8")
            captured["files"].append({"kind": "receipt", "path": str(paths["receipt"])})
            return captured
        finally:
            context.close()
            browser.close()


def main() -> int:
    parser = argparse.ArgumentParser(description="Capture PixelRAG visual evidence from a manifest.")
    parser.add_argument("--manifest-file")
    parser.add_argument("--manifest-json")
    parser.add_argument("--output-dir", default=DEFAULT_OUTPUT_DIR)
    parser.add_argument("--dry-run", action="store_true")
    parser.add_argument("--timeout", type=int, default=30)
    parser.add_argument("--settle-ms", type=int, default=DEFAULT_SETTLE_MS)
    parser.add_argument("--max-tiles", type=int, default=DEFAULT_MAX_TILES)
    parser.add_argument("--wait-until", default="domcontentloaded")
    parser.add_argument("--allow-any-host", action="store_true")
    args = parser.parse_args()
    if not args.manifest_file and not args.manifest_json:
        parser.error("Pass --manifest-file or --manifest-json.")

    manifest = _load_manifest(args)
    errors = _validate_manifest(manifest, require_allowed_host=not args.allow_any_host)
    artifact = _base_artifact(manifest, args, errors)
    if errors or args.dry_run:
        print(json.dumps(artifact, ensure_ascii=False, indent=2))
        return 1 if errors else 0
    captured = _capture(manifest, args, artifact)
    print(json.dumps(captured, ensure_ascii=False, indent=2))
    return 0


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