#!/usr/bin/env python3
"""Fetch one policy-selected catalog artifact from its live source.

The Go worker passes a small JSON manifest (argv[1]) describing the intent.
This process never receives database credentials or secret values; it locates
credentials only via the pre-mounted /run/source-secrets paths, exactly like
live_catalog_refresh.py. Behavior per source type:

  alist  -> direct read-only API: raw_url then download.
  pp     -> manifest.json local_path first (already on disk), else public URL.
  dingtalk -> prefer the local export tree (converted docx/xlsx/pdf); if it is
              older than the source-side modification time, export fresh bytes
              through the web internal API captured headers.

Before any network call the fetcher compares against local candidates by
size and, when a hash is cheap to obtain, by content: a match means "本地已有
- no pull". Verified bytes are moved into the ingestion spool and a single
JSON result line is printed on stdout. Credentials, tokens, signed URLs and
authorization material are never printed.
"""
import csv
import hashlib
import json
import os
import pathlib
import re
import shutil
import subprocess
import sys
import tempfile
import time
import urllib.error
import urllib.parse
import urllib.request


SUPPORTED = {
    "pdf": "pdf", "docx": "docx", "xlsx": "xlsx", "md": "md", "txt": "txt",
    "csv": "csv", "json": "json",
    "png": "png", "jpg": "jpg", "jpeg": "jpeg", "gif": "gif", "bmp": "bmp", "webp": "webp",
}
# Office types we refuse to ingest directly: PPT/PPTX must arrive as a PDF twin
# (embedded images travel inside those files already).
SLIDES = {"ppt", "pptx"}
LEGACY = {"doc": "doc", "xls": "xls"}
MAX_BYTES_DEFAULT = 200 * 1024 * 1024


class FetchError(Exception):
    def __init__(self, message, retryable=True, code="fetch_failed"):
        super().__init__(message)
        self.retryable = retryable
        self.code = code


class MissingSource(Exception):
    pass


def emit(result):
    print(json.dumps(result, ensure_ascii=False))


def sha256_file(path):
    digest = hashlib.sha256()
    with open(path, "rb") as handle:
        for block in iter(lambda: handle.read(1 << 20), b""):
            digest.update(block)
    return digest.hexdigest()


def digest_of(value):
    return hashlib.sha256(value.encode()).hexdigest()


def validate_magic(path, suffix):
    with open(path, "rb") as handle:
        head = handle.read(12)
    checks = {
        "pdf": head.startswith(b"%PDF-"),
        "docx": head.startswith(b"PK\x03\x04"),
        "xlsx": head.startswith(b"PK\x03\x04"),
        "doc": head.startswith(b"\xd0\xcf\x11\xe0"),
        "xls": head.startswith(b"\xd0\xcf\x11\xe0"),
        "png": head.startswith(b"\x89PNG\r\n\x1a\n"),
        "jpg": head[:2] == b"\xff\xd8" and head[2:3] == b"\xff",
        "jpeg": head[:2] == b"\xff\xd8" and head[2:3] == b"\xff",
        "gif": head[:6] in (b"GIF87a", b"GIF89a"),
        "bmp": head[:2] == b"BM",
        "webp": head[:4] == b"RIFF" and head[8:12] == b"WEBP",
    }
    if suffix in checks and not checks[suffix]:
        raise FetchError("downloaded bytes are not a valid %s file" % suffix,
                         retryable=False, code="invalid_format")
    if suffix in ("md", "txt", "csv", "json"):
        try:
            with open(path, "rb") as handle:
                handle.read(4096).decode("utf-8", errors="strict")
        except UnicodeDecodeError:
            raise FetchError("text artifact is not valid UTF-8", retryable=False, code="invalid_format")


def pick_suffix(name, declared_type):
    suffix = pathlib.PurePosixPath(name).suffix.lower().lstrip(".")
    if declared_type in ("adoc",) and suffix in ("", "adoc"):
        return "docx"
    if declared_type in ("able", "axls") and suffix in ("", "able", "axls"):
        return "xlsx"
    if suffix in SUPPORTED:
        return SUPPORTED[suffix]
    if suffix == "adoc":
        return "docx"
    if suffix in ("able", "axls"):
        return "xlsx"
    return None


def http_get(url, headers, timeout, max_bytes, cookie_opener=None):
    request = urllib.request.Request(url, headers=headers or {})

    class NoRedirect(urllib.request.HTTPRedirectHandler):
        location = ""

        def redirect_request(self, req, fp, code, msg, headers, newurl):
            NoRedirect.location = newurl
            return None

    handler = NoRedirect()
    opener = cookie_opener or urllib.request.build_opener(handler)
    try:
        response = opener.open(request, timeout=timeout)
    except urllib.error.HTTPError as error:
        if error.code in (301, 302, 303, 307, 308):
            target = handler.location or error.headers.get("Location") or ""
            host = urllib.parse.urlparse(target).netloc.lower()
            if any(marker in host for marker in ("login.", "sso", "oauth", "passport")):
                raise FetchError("文件所在的来源存储需要网页登录授权（存储后端未配置服务凭据），请联系管理员", retryable=False, code="source_auth_required")
            raise FetchError("redirect refused", retryable=True, code="redirect")
        if error.code == 404:
            raise MissingSource("source returned 404")
        raise FetchError("source returned HTTP %d" % error.code, retryable=error.code >= 500, code="http_%d" % error.code)
    with response:
        declared = response.headers.get("Content-Length")
        if declared and int(declared) > max_bytes:
            raise FetchError("source file exceeds ingest size limit", retryable=False, code="over_limit")
        destination = tempfile.NamedTemporaryFile(delete=False, suffix=".fetch")
        total = 0
        digest = hashlib.sha256()
        try:
            while True:
                block = response.read(1 << 20)
                if not block:
                    break
                total += len(block)
                if total > max_bytes:
                    raise FetchError("source stream exceeds ingest size limit", retryable=False, code="over_limit")
                digest.update(block)
                destination.write(block)
        finally:
            destination.close()
        if total == 0:
            os.unlink(destination.name)
            raise FetchError("source returned empty body", retryable=True, code="empty")
        return destination.name, total, digest.hexdigest()


# ---------------------------------------------------------------- AList ------

def alist_download(manifest, max_bytes):
    sys.path.insert(0, "/sources/alist")
    from alist_readonly_inventory import AListReadOnlyClient, read_connections  # noqa

    token = pathlib.Path("/run/source-secrets/alist_api_token").read_text(encoding="utf-8").strip()
    if not token:
        raise FetchError("AList token empty", retryable=False, code="credentials")
    rows = [row for row in read_connections(pathlib.Path("/run/source-secrets/alist_connections.xlsx"), require_token=False)
            if row["category"] == manifest["source_name"]]
    if len(rows) != 1:
        raise FetchError("AList connection configuration missing", retryable=False, code="credentials")
    base_url = rows[0]["base_url"].rstrip("/")
    client = AListReadOnlyClient(base_url, token)
    path = "/" + manifest["relative_path"].lstrip("/")
    payload = json.dumps({"path": path, "password": ""}).encode("utf-8")
    request = urllib.request.Request(base_url + "/api/fs/get", data=payload, method="POST", headers={
        "x-api-token": token, "Content-Type": "application/json", "Accept": "application/json",
        "User-Agent": "governance-artifact-fetch/1.0",
    })
    try:
        with client.opener.open(request, timeout=30) as response:
            body = json.loads(response.read().decode("utf-8"))
    except Exception as error:
        raise FetchError("AList file lookup failed (%s)" % type(error).__name__, code="source_unavailable")
    if int(body.get("code") or 0) == 404:
        raise MissingSource("AList file not present at expected path")
    if int(body.get("code") or 0) != 200:
        raise FetchError("AList file lookup refused (code %s)" % body.get("code"), code="source_unavailable")
    data = body.get("data") or {}
    raw_url = str(data.get("raw_url") or "")
    if not raw_url:
        raise MissingSource("AList returned no download locator for this file")
    if manifest.get("size_bytes") in (None, 0) and data.get("size"):
        manifest["size_bytes"] = int(data["size"])
    download_url = urllib.parse.urljoin(base_url + "/", raw_url.lstrip("/"))
    # This AList deployment gates the download route with x-api-token (the
    # Authorization header is NOT recognized there): wrong auth yields a 302
    # to the DingTalk login page instead of bytes.
    try:
        return http_get(download_url, {"x-api-token": token}, 180, max_bytes)
    except FetchError as error:
        if error.code == "source_auth_required":
            try:
                return http_get(download_url, {"Authorization": token}, 180, max_bytes)
            except FetchError:
                raise error
        if error.code == "redirect":
            # Same-origin /d/ route 302s to a signed external URL: follow once
            # WITHOUT the source token (never forward credentials cross-origin).
            class LocationCapture(urllib.request.HTTPRedirectHandler):
                captured = ""

                def redirect_request(self, req, fp, code, msg, headers, newurl):
                    LocationCapture.captured = newurl
                    return None

            probe = urllib.request.build_opener(LocationCapture)
            try:
                probe.open(urllib.request.Request(download_url, headers={"Authorization": token}), timeout=30)
            except urllib.error.HTTPError:
                pass
            except Exception as failure:
                raise FetchError("AList redirect probe failed (%s)" % type(failure).__name__, code="source_unavailable")
            if not LocationCapture.captured:
                raise
            return http_get(LocationCapture.captured, {}, 180, max_bytes)
        raise


# ------------------------------------------------------------------ PP -------

def pp_find_entry(manifest):
    locator = manifest.get("locator_hash", "")
    if not locator:
        return None
    manifest_path = pathlib.Path("/sources/pp/data/public-documents/manifest.json")
    ledger = json.loads(manifest_path.read_text(encoding="utf-8"))
    for result in ledger.get("results", []):
        if digest_of(str(result.get("url") or "")) == locator:
            return result
    return None


def pp_download(manifest, max_bytes):
    entry = pp_find_entry(manifest)
    if entry is None:
        raise MissingSource("PP ledger entry for this locator no longer exists")
    # Ledger titles may omit the extension; recover the true file name from the
    # local path or source URL before type checks run.
    manifest_name = manifest.setdefault("name", "")
    if not pathlib.PurePosixPath(manifest_name).suffix:
        from urllib.parse import unquote, urlparse
        for probe in (entry.get("local_path", ""), unquote(urlparse(str(entry.get("url", ""))).path)):
            suffix = pathlib.PurePosixPath(probe).suffix.lower().lstrip(".")
            if suffix:
                manifest["name"] = manifest_name + "." + suffix
                break
    local_path = entry.get("local_path", "")
    if local_path:
        candidate = pathlib.Path(local_path.replace("/home/zheromiku/dobot-portal-sync", "/sources/pp", 1))
        if not candidate.is_file():
            candidate = pathlib.Path(local_path)
        if candidate.is_file():
            size = candidate.stat().st_size
            expected_hash = str(entry.get("sha256", "")).lower()
            actual = sha256_file(candidate)
            if expected_hash and len(expected_hash) == 64 and actual != expected_hash:
                raise FetchError("PP local copy hash mismatch (ledger changed?)", retryable=True, code="local_mismatch")
            return str(candidate), size, actual, True
    url = entry.get("url", "")
    if not url:
        raise MissingSource("PP ledger entry has no local file and no download URL")
    path, size, actual = http_get(url, {}, 180, max_bytes)
    expected_hash = str(entry.get("sha256", "")).lower()
    if expected_hash and len(expected_hash) == 64 and actual != expected_hash:
        os.unlink(path)
        raise FetchError("PP download hash mismatch", retryable=True, code="hash_mismatch")
    return path, size, actual, False


# ------------------------------------------------------------ DingTalk -------

def load_capture_headers():
    capture = json.loads(pathlib.Path("/run/source-secrets/dingtalk_capture.json").read_text(encoding="utf-8"))
    headers = capture.get("headers") or {}
    def header(name):
        for key, value in headers.items():
            if key.lower() == name:
                return value
        return ""
    a_token = header("a-token")
    xsrf = header("x-xsrf-token")
    cookie = header("cookie")
    uid = ""
    match = re.search(r'"?uid"?[=:]"?(\d{4,})', cookie or "")
    if match:
        uid = match.group(1)
    if not a_token:
        raise FetchError("DingTalk capture has no a-token; re-capture in 来源与凭据", retryable=False, code="credentials")
    return {"a-token": a_token, "x-xsrf-token": xsrf, "cookie": cookie, "uid": uid,
            "referer": "https://alidocs.dingtalk.com/"}


def dingtalk_local_export(manifest):
    """Locate an already-exported copy; returns (path, fresher_than_source) or None."""
    name = pathlib.PurePosixPath(manifest["relative_path"]).name
    stem = pathlib.PurePosixPath(name).stem
    raw_suffix = pathlib.PurePosixPath(name).suffix.lower().lstrip(".")
    suffix = pick_suffix(name, manifest.get("declared_type", ""))
    root = pathlib.Path("/sources/dingtalk/exported")
    folder = str(pathlib.PurePosixPath(manifest["relative_path"]).parent).strip(".")
    search_suffixes = [suffix] if raw_suffix not in SLIDES else ["pdf", "pptx", "ppt"]
    candidates = []
    for probe in search_suffixes:
        if probe is None:
            continue
        exact = root / manifest["source_name"] / folder / (stem + "." + probe)
        if exact.is_file():
            candidates.append(exact)
        else:
            for base in (root / manifest["source_name"], pathlib.Path("/sources/dingtalk/converted_ppt_pdf"),
                         pathlib.Path("/sources/dingtalk/converted_ppt_pdf_full")):
                if not base.is_dir():
                    continue
                for hit in base.rglob(stem + "." + probe):
                    candidates.append(hit)
        if candidates:
            break
    if not candidates:
        if raw_suffix not in SLIDES and suffix is None:
            raise MissingSource("dingtalk file type is not ingestable")
        return None
    modified = str(manifest.get("modified_at") or "")
    candidate = candidates[0]
    local_modified = time.strftime("%Y-%m-%dT%H:%M:%S+08:00", time.localtime(candidate.stat().st_mtime))
    fresh = (not modified) or (local_modified >= modified[:19] + local_modified[19:])
    return str(candidate), fresh


def dingtalk_raw_keys(manifest):
    """Recover the dentry uuid behind a node id via the scanner inventory."""
    state_path = pathlib.Path("/sources/dingtalk/data/dingtalk-internal/inventory-state.json")
    state = json.loads(state_path.read_text(encoding="utf-8"))
    roots = state.get("roots") or {}
    space_key = _space_key(roots, manifest["source_name"])
    match = [value for key, value in sorted(roots.items()) if key == space_key]
    if len(match) != 1:
        raise FetchError("DingTalk space not present in the scanner inventory; refresh 来源目录 first",
                         retryable=False, code="space_missing")
    sys.path.insert(0, "/app/scripts")
    from source_catalog_import import node_id as catalog_node_id  # noqa
    source_id = "dingtalk-" + space_key
    nodes = match[0].get("nodes") or {}
    for raw, item in nodes.items():
        if catalog_node_id(source_id, raw) == manifest["node_id"]:
            return raw, str(item.get("dentry_type") or item.get("type") or "")
    raise MissingSource("DingTalk node not found in current inventory (moved or deleted?)")


def _space_key(roots, name):
    for key, value in roots.items():
        if value.get("live_space_name") == name:
            return key
    return ""


def dingtalk_download(manifest, max_bytes):
    export = dingtalk_local_export(manifest)
    modified = str(manifest.get("modified_at") or "")
    if export:
        path, fresh = export
        size = os.path.getsize(path)
        if size > 0 and (fresh or not modified):
            return path, size, sha256_file(path), True
        # stale local export -> fall through to a live export when available
    raw_key, declared_type = dingtalk_raw_keys(manifest)
    headers = load_capture_headers()
    live_path = live_export_dingtalk_document(raw_key, headers, manifest, max_bytes)
    if export and not live_path:
        path, fresh = export
        return path, os.path.getsize(path), sha256_file(path), True  # keep old bytes; staleness shown in UI
    if not live_path:
        raise MissingSource("DingTalk file is not exported and cannot be exported automatically")
    return live_path


def live_export_dingtalk_document(raw_key, headers, manifest, max_bytes):
    opener = urllib.request.build_opener()

    def api(url, payload=None):
        request = urllib.request.Request(url, headers={
            "a-token": headers["a-token"], "x-xsrf-token": headers.get("x-xsrf-token", ""),
            "cookie": headers.get("cookie", ""), "content-type": "application/json",
            "referer": headers["referer"], "user-agent": "Mozilla/5.0",
        }, data=json.dumps(payload).encode() if payload is not None else None)
        try:
            with opener.open(request, timeout=60) as response:
                return json.loads(response.read().decode("utf-8"))
        except (urllib.error.URLError, urllib.error.HTTPError, TimeoutError) as error:
            raise FetchError("DingTalk web API unavailable (%s)" % type(error).__name__, code="source_unavailable")

    info = api("https://alidocs.dingtalk.com/box/api/v2/dentry/info?dentryUuid=" + urllib.parse.quote(raw_key))
    data = info.get("data") or {}
    space_id, dentry_id = str(data.get("spaceId") or ""), str(data.get("dentryId") or "")
    if not space_id or not dentry_id:
        raise MissingSource("DingTalk dentry not resolvable")
    download = api("https://alidocs.dingtalk.com/box/api/v1/file/download?spaceId=%s&dentryId=%s" % (space_id, dentry_id))
    payload = download.get("data") or {}
    if payload.get("downloadType") == "URL_PRE_SIGNATURE":
        urls = (payload.get("ossUrlPreSignatureInfo") or {}).get("preSignUrls") or []
        if urls:
            return http_get(urls[0], {}, 180, max_bytes) + (False,)
    return None


# ---------------------------------------------------------------- main -------

def suffix_from_magic(path):
    with open(path, "rb") as handle:
        head = handle.read(12)
    if head.startswith(b"%PDF-"):
        return "pdf"
    if head.startswith(b"PK\x03\x04"):
        # zip container: peek one entry name to pick office family
        try:
            import zipfile
            with zipfile.ZipFile(path) as archive:
                names = set(archive.namelist()[:50])
            if "word/document.xml" in names:
                return "docx"
            if "xl/workbook.xml" in names:
                return "xlsx"
        except Exception:
            return None
        return None
    if head.startswith(b"\xd0\xcf\x11\xe0"):
        return None
    for suffix in ("png", "jpg", "gif", "bmp", "webp"):
        checks = validate_magic_map(head)
        if checks.get(suffix):
            return suffix
    try:
        with open(path, "rb") as handle:
            handle.read(4096).decode("utf-8", errors="strict")
        return "md"
    except (UnicodeDecodeError, OSError):
        return None


def validate_magic_map(head):
    return {
        "png": head.startswith(b"\x89PNG\r\n\x1a\n"),
        "jpg": head[:2] == b"\xff\xd8" and head[2:3] == b"\xff",
        "gif": head[:6] in (b"GIF87a", b"GIF89a"),
        "bmp": head[:2] == b"BM",
        "webp": head[:4] == b"RIFF" and head[8:12] == b"WEBP",
    }


def main():
    manifest = json.loads(sys.argv[1])
    max_bytes = int(os.environ.get("ARTIFACT_FETCH_MAX_BYTES") or MAX_BYTES_DEFAULT)
    spool_root = pathlib.Path(os.environ.get("INGESTION_SPOOL_ROOT") or "/var/lib/governance/ingestion")
    name = manifest["name"]
    declared_type = manifest.get("declared_type", "")
    matched_local = False
    if manifest["source_type"] == "alist":
        path, size, digest = alist_download(manifest, max_bytes)
    elif manifest["source_type"] == "pp":
        path, size, digest, matched_local = pp_download(manifest, max_bytes)
    elif manifest["source_type"] == "dingtalk":
        path, size, digest, matched_local = dingtalk_download(manifest, max_bytes)
    else:
        raise FetchError("unknown source type", retryable=False, code="bad_manifest")
    # Type is decided on the RESOLVED bytes, not the catalog title: PP ledger
    # titles may omit extensions and DingTalk exports always change them.
    suffix = pick_suffix(pathlib.PurePosixPath(path).name, declared_type) or pick_suffix(name, declared_type)
    if suffix is None:
        suffix = suffix_from_magic(path)
    if suffix is None:
        raw_suffix = pathlib.PurePosixPath(name).suffix.lower().lstrip(".")
        hint = " (convert PPT to PDF first)" if raw_suffix in SLIDES else ""
        raise FetchError("file type .%s is not supported by the parser%s" % (raw_suffix or "?", hint),
                         retryable=False, code="unsupported_format")
    try:
        validate_magic(path, suffix)
        if size > max_bytes:
            raise FetchError("artifact exceeds ingest size limit", retryable=False, code="over_limit")
        relative_ref = "artifact/%s/%s/%s.%s" % (digest[:2], digest[2:4], manifest["node_id"], suffix)
        destination = spool_root / relative_ref
        destination.parent.mkdir(parents=True, exist_ok=True)
        staged = str(destination) + ".part"
        shutil.copyfile(path, staged)
        if sha256_file(staged) != digest:
            raise FetchError("spooled bytes failed verification", code="spool_mismatch")
        os.replace(staged, destination)
        result = {"ok": True, "spec_id": manifest["spec_id"], "node_id": manifest["node_id"],
                  "connection_id": manifest["connection_id"], "sha256": digest, "size_bytes": size,
                  "spool_ref": os.environ.get("SPOOL_REF_PREFIX", "") + relative_ref, "validated_format": suffix, "matched_local": matched_local}
    finally:
        if path.startswith(tempfile.gettempdir()) and os.path.exists(path):
            os.unlink(path)
    emit(result)
    return 0


if __name__ == "__main__":
    try:
        sys.exit(main())
    except MissingSource as error:
        emit({"ok": False, "missing": True, "error": str(error)})
        sys.exit(3)
    except FetchError as error:
        emit({"ok": False, "retryable": error.retryable, "code": error.code, "error": str(error)})
        sys.exit(4 if error.retryable else 2)