llm/scripts/download_1c_its_media.py

from __future__ import annotations

import argparse
import hashlib
import json
import mimetypes
import os
import re
import time
import urllib.error
import urllib.parse
import urllib.request
from pathlib import Path
from typing import Any


ROOT = Path(__file__).resolve().parents[1]
DEFAULT_MANIFEST = ROOT / "plugins" / "1c" / "rag" / "official-docs" / "normalized" / "manifest.json"
DEFAULT_OUTPUT_DIR = ROOT / "plugins" / "1c" / "rag" / "official-docs" / "media"
DEFAULT_OUTPUT_MANIFEST = DEFAULT_OUTPUT_DIR / "manifest.json"


def load_json(path: Path) -> dict[str, Any]:
    if not path.exists():
        return {}
    try:
        data = json.loads(path.read_text(encoding="utf-8-sig"))
    except json.JSONDecodeError:
        return {}
    return data if isinstance(data, dict) else {}


def read_cookie(cookie_file: Path | None) -> str:
    if cookie_file:
        return cookie_file.read_text(encoding="utf-8").strip()
    return os.environ.get("ONEC_ITS_COOKIE", "").strip()


def extension_from_url(url: str, content_type: str = "") -> str:
    suffix = Path(urllib.parse.urlparse(url).path).suffix.lower()
    if suffix:
        return suffix
    guessed = mimetypes.guess_extension(content_type.split(";", 1)[0].strip())
    return guessed or ".bin"


def safe_media_name(url: str, content_type: str = "") -> str:
    parsed = urllib.parse.urlparse(url)
    stem = re.sub(r"[^A-Za-zА-Яа-яЁё0-9_.-]+", "_", Path(parsed.path).stem, flags=re.UNICODE).strip("_") or "media"
    digest = hashlib.sha256(url.encode("utf-8")).hexdigest()[:12]
    return f"{stem[:80]}__{digest}{extension_from_url(url, content_type)}"


def iter_media_items(manifest: dict[str, Any]) -> list[dict[str, Any]]:
    items = []
    seen = set()
    for page in manifest.get("pages") or []:
        for image in ((page.get("media") or {}).get("images") or []):
            url = str(image.get("url") or "")
            if not url or url in seen:
                continue
            seen.add(url)
            items.append(
                {
                    "url": url,
                    "page_title": page.get("title"),
                    "page_url": page.get("url"),
                    "source_id": page.get("source_id"),
                    "source_type": page.get("source_type"),
                    "alt": image.get("alt"),
                    "title": image.get("title"),
                    "width": image.get("width"),
                    "height": image.get("height"),
                }
            )
    return items


def fetch_media(url: str, *, cookie: str, timeout: int, user_agent: str) -> tuple[bytes, dict[str, str], int]:
    headers = {"User-Agent": user_agent}
    if cookie:
        headers["Cookie"] = cookie
    request = urllib.request.Request(url, headers=headers)
    with urllib.request.urlopen(request, timeout=timeout) as response:
        return response.read(), {key: value for key, value in response.headers.items()}, int(response.status)


def download_media(
    *,
    manifest_path: Path,
    output_dir: Path,
    output_manifest: Path,
    cookie: str,
    timeout: int,
    user_agent: str,
    limit: int,
) -> dict[str, Any]:
    manifest = load_json(manifest_path)
    output_dir.mkdir(parents=True, exist_ok=True)
    items = iter_media_items(manifest)
    if limit > 0:
        items = items[:limit]

    downloaded = []
    errors = []
    for item in items:
        url = item["url"]
        try:
            probe_name = safe_media_name(url)
            existing = next(output_dir.glob(f"{Path(probe_name).stem}.*"), None)
            if existing and existing.is_file():
                body = existing.read_bytes()
                item.update(
                    {
                        "status": "cached",
                        "file": existing.name,
                        "bytes": len(body),
                        "sha256": hashlib.sha256(body).hexdigest(),
                    }
                )
                downloaded.append(item)
                continue
            body, headers, status = fetch_media(url, cookie=cookie, timeout=timeout, user_agent=user_agent)
            filename = safe_media_name(url, headers.get("Content-Type", ""))
            target = output_dir / filename
            target.write_bytes(body)
            item.update(
                {
                    "status": "downloaded",
                    "file": filename,
                    "content_type": headers.get("Content-Type", ""),
                    "http_status": status,
                    "bytes": len(body),
                    "sha256": hashlib.sha256(body).hexdigest(),
                }
            )
            downloaded.append(item)
        except (urllib.error.URLError, TimeoutError, OSError) as exc:
            error = dict(item)
            error.update({"status": "error", "error": str(exc)})
            errors.append(error)

    result = {
        "schema": "onec_its_media_manifest.v1",
        "created_at_unix": int(time.time()),
        "source_manifest": str(manifest_path),
        "output_dir": str(output_dir),
        "counts": {
            "media_items": len(iter_media_items(manifest)),
            "attempted": len(items),
            "saved": len(downloaded),
            "errors": len(errors),
        },
        "items": downloaded,
        "errors": errors,
    }
    output_manifest.parent.mkdir(parents=True, exist_ok=True)
    output_manifest.write_text(json.dumps(result, ensure_ascii=False, indent=2) + "\n", encoding="utf-8")
    return result


def main() -> int:
    parser = argparse.ArgumentParser(description="Download media assets referenced by normalized private 1C:ITS pages.")
    parser.add_argument("--manifest", type=Path, default=DEFAULT_MANIFEST)
    parser.add_argument("--output-dir", type=Path, default=DEFAULT_OUTPUT_DIR)
    parser.add_argument("--output-manifest", type=Path, default=DEFAULT_OUTPUT_MANIFEST)
    parser.add_argument("--cookie-file", type=Path)
    parser.add_argument("--timeout", type=int, default=30)
    parser.add_argument("--limit", type=int, default=0, help="0 means all media items.")
    parser.add_argument("--user-agent", default="Codex-1C-RAG/1.0 (+licensed private 1C:ITS access)")
    args = parser.parse_args()

    result = download_media(
        manifest_path=args.manifest,
        output_dir=args.output_dir,
        output_manifest=args.output_manifest,
        cookie=read_cookie(args.cookie_file),
        timeout=args.timeout,
        user_agent=args.user_agent,
        limit=args.limit,
    )
    print(json.dumps({"counts": result["counts"], "output": str(args.output_manifest)}, ensure_ascii=False, indent=2))
    return 0 if result["counts"]["errors"] == 0 else 2


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