#!/usr/bin/env python3
"""Capture retained Bubble Visual URLs and prepare a private, hash-bound review bundle."""

import argparse
import copy
import hashlib
import json
import os
import re
import stat
import sys
import urllib.request
from datetime import UTC, datetime
from pathlib import Path
from urllib.parse import urlsplit

CDN = '13a55be71f3a6a4c9ea356833df823c4.cdn.bubble.io'
MAX_ASSET = 16 * 1024 * 1024
MAX_TOTAL = 256 * 1024 * 1024
ROOT = Path(__file__).resolve().parents[1]


class Refused(Exception):
    pass


def sha(data):
    return hashlib.sha256(data).hexdigest()


def canonical_url(value):
    if not isinstance(value, str) or any(ord(c) <= 32 for c in value):
        raise Refused('invalid_visual_url')
    value = 'https:' + value if value.startswith('//') else value
    try:
        url = urlsplit(value)
        valid = (url.scheme == 'https' and url.hostname == CDN and url.port in (None, 443)
                 and url.username is None and url.password is None and not url.query
                 and not url.fragment and url.path.startswith('/') and '\\' not in value)
    except ValueError:
        valid = False
    if not valid:
        raise Refused('outside_reviewed_visual_origin')
    return value


class NoRedirect(urllib.request.HTTPRedirectHandler):
    def redirect_request(self, req, fp, code, msg, headers, newurl):
        raise Refused('visual_redirect_refused')


def image_extension(raw, media_type):
    if raw.startswith(b'\x89PNG\r\n\x1a\n') and media_type == 'image/png':
        return 'png'
    if raw.startswith(b'\xff\xd8\xff') and media_type == 'image/jpeg':
        return 'jpg'
    if (len(raw) >= 12 and raw[:4] == b'RIFF' and raw[8:12] == b'WEBP'
            and media_type == 'image/webp'):
        return 'webp'
    raise Refused('unsupported_visual_bytes_or_type')


def fetch(url):
    url = canonical_url(url)
    opener = urllib.request.build_opener(urllib.request.ProxyHandler({}), NoRedirect())
    request = urllib.request.Request(url, headers={'User-Agent': 'coulomb-visual-capture/1.0'})
    with opener.open(request, timeout=20) as response:
        if response.status != 200 or response.geturl() != url:
            raise Refused('unexpected_visual_response')
        length = response.headers.get('Content-Length')
        if length and (not length.isdigit() or int(length) > MAX_ASSET):
            raise Refused('visual_size_limit')
        raw = response.read(MAX_ASSET + 1)
        if len(raw) > MAX_ASSET:
            raise Refused('visual_size_limit')
        if length and int(length) != len(raw):
            raise Refused('incomplete_visual_response')
        extension = image_extension(raw, response.headers.get_content_type())
        return raw, extension


def read_private(root, name, expected=None):
    relative = Path(name)
    if relative.is_absolute() or '..' in relative.parts or not relative.parts:
        raise Refused('unsafe_bundle_path')
    path = root / relative
    if any(p.is_symlink() for p in (path, *path.parents)):
        raise Refused('bundle_symlink_refused')
    info = path.stat()
    if (not stat.S_ISREG(info.st_mode) or info.st_uid != os.getuid()
            or info.st_mode & 0o077 or info.st_nlink != 1 or info.st_size > MAX_ASSET):
        raise Refused('private_bounded_source_required')
    raw = path.read_bytes()
    if expected is not None and sha(raw) != expected:
        raise Refused('source_file_changed')
    return raw


def publish(root, name, raw):
    relative = Path(name)
    if relative.is_absolute() or '..' in relative.parts or not relative.parts:
        raise Refused('unsafe_output_path')
    path = root / relative
    if any(p.is_symlink() for p in (path, *path.parents)):
        raise Refused('output_symlink_refused')
    parent = root
    for part in relative.parts[:-1]:
        parent = parent / part
        parent.mkdir(mode=0o700, exist_ok=True)
        info = parent.stat()
        if info.st_uid != os.getuid() or info.st_mode & 0o077:
            raise Refused('private_output_directory_required')
    fd = os.open(path, os.O_WRONLY | os.O_CREAT | os.O_EXCL | os.O_NOFOLLOW, 0o600)
    with os.fdopen(fd, 'wb') as stream:
        stream.write(raw)
        stream.flush()
        os.fsync(stream.fileno())
    if read_private(root, name) != raw:
        raise Refused('capture_readback_failed')


def with_visual(raw, relative, *, index=False):
    # Replace only the compiler's known empty field; never parse/reformat the body.
    head, body = raw.split(b'\n---\n', 1)
    if not head.startswith(b'---\n'):
        raise Refused('draft_frontmatter_required')
    new = b'visual: ' + json.dumps(relative).encode()
    if index:
        if re.search(rb'^visual:', head, re.M):
            raise Refused('unexpected_index_visual')
        head += b'\n' + new
    else:
        if len(re.findall(rb'^visual: ""$', head, re.M)) != 1:
            raise Refused('empty_draft_visual_required')
        head = re.sub(rb'^visual: ""$', lambda _: new, head, flags=re.M)
    return head + b'\n---\n' + body


def capture(bundle, output, fetcher=fetch):
    bundle = Path(os.path.abspath(bundle.expanduser()))
    output = Path(os.path.abspath(output.expanduser()))
    if (output.is_relative_to(ROOT) or output.is_relative_to(bundle)
            or bundle.is_relative_to(output)
            or any(p.is_symlink() for p in (output, *output.parents))
            or any((p / '.git').exists() for p in (output, *output.parents))):
        raise Refused('private_separate_output_required')
    parent = output.parent.stat()
    if parent.st_uid != os.getuid() or parent.st_mode & 0o077:
        raise Refused('private_owned_output_parent_required')
    original = read_private(bundle, 'reconciliation.json')
    report = json.loads(original)
    if (report.get('schema') != 'coulomb.bubble-review-bundle.v1'
            or report.get('status') != 'completed_offline_preparation'):
        raise Refused('completed_original_bundle_required')
    # Verify every declared file and the captured source bindings before any network read.
    files = {name: read_private(bundle, name, digest)
             for name, digest in report['file_sha256'].items()}
    if any(not name.startswith(('source-bodies/', 'review-trees/')) for name in files):
        raise Refused('unexpected_original_bundle_namespace')
    for space in report['spaces']:
        tree = space['tree']
        if space.get('import_ready') is not False:
            raise Refused('unapproved_original_readiness')
        if tree is None:
            continue
        if not re.fullmatch(r'review-trees/bubble-[0-9a-f]{24}', tree):
            raise Refused('unexpected_review_tree_path')
        for page in space['articles']:
            if not re.fullmatch(re.escape(tree) + r'/pages/article-[0-9a-f]{24}\.md',
                                page['file']):
                raise Refused('unexpected_review_page_path')
    for name, digest in report['source_sha256'].items():
        files['source-capture/' + name + '.json'] = read_private(
            bundle, 'source-capture/' + name + '.json', digest)
    visuals = {}
    for article in report['articles']:
        refs = [ref['source'] for ref in article['references'] if ref['kind'] == 'visual']
        if len(refs) > 1:
            raise Refused('ambiguous_visual_reference')
        if refs:
            visuals[article['id']] = canonical_url(refs[0])
    output.mkdir(mode=0o700)  # Existing output always refuses; never replace a capture.
    captured, total = {}, 0
    for url in sorted(set(visuals.values())):
        raw, extension = fetcher(url)
        if len(raw) > MAX_ASSET or total + len(raw) > MAX_TOTAL:
            raise Refused('visual_capture_budget_exceeded')
        total += len(raw)
        name = 'captured-visuals/' + sha(raw) + '.' + extension
        if name not in files:
            files[name] = raw
        captured[url] = {'file': name, 'sha256': sha(raw), 'bytes': len(raw)}
        # Checkpoint captured bytes immediately; no final report means incomplete output.
        if not (output / name).exists():
            publish(output, name, raw)
    updated = copy.deepcopy(report)
    by_id = {a['id']: a for a in updated['articles']}
    materialized = 0
    for space in updated['spaces']:
        tree = space['tree']
        if tree is None:
            continue
        for page in space['articles']:
            url = visuals.get(page['id'])
            if not url:
                continue
            visual = captured[url]
            relative = 'assets/' + Path(visual['file']).name
            files[tree + '/' + relative] = files[visual['file']]
            files[page['file']] = with_visual(files[page['file']], relative)
            page['file_sha256'] = sha(files[page['file']])
            materialized += 1
            if page['id'] == space['root_article_id']:
                index = tree + '/pages/index.md'
                files[index] = with_visual(files[index], relative, index=True)
        # Keep all unresolved gates: retained bytes do not certify source/UI fidelity or grants.
        space['import_ready'] = False
    for identity, url in visuals.items():
        article = by_id[identity]
        article['captured_visual'] = captured[url]
        for ref in article['references']:
            if ref['kind'] == 'visual':
                ref['capture_status'] = 'bytes_retained_and_hashed'
                ref['resolved'] = False
    # Recheck the immutable inputs after fetching; source drift cannot publish completion.
    if read_private(bundle, 'reconciliation.json') != original:
        raise Refused('source_report_changed_during_capture')
    for name, digest in report['file_sha256'].items():
        read_private(bundle, name, digest)
    for name, digest in report['source_sha256'].items():
        read_private(bundle, 'source-capture/' + name + '.json', digest)
    files['source-capture/reconciliation-before-visuals.json'] = original
    for name, raw in sorted(files.items()):
        if not name.startswith('captured-visuals/'):
            publish(output, name, raw)
    updated.update(status='completed_private_visual_preparation',
                   original_report_sha256=sha(original),
                   captured_at=datetime.now(UTC).isoformat(),
                   visual_capture={'references': len(visuals), 'unique_urls': len(captured),
                                   'downloaded_bytes': total, 'materialized_pages': materialized,
                                   'body_rewriting_performed': False, 'grants_changed': False,
                                   'browser_render_verified': False, 'production_imported': False},
                   network_requests=len(captured),
                   file_sha256={name: sha(raw) for name, raw in sorted(files.items())})
    for directory in sorted([output, *(p for p in output.rglob('*') if p.is_dir())],
                            key=lambda p: len(p.parts), reverse=True):
        fd = os.open(directory, os.O_RDONLY | os.O_DIRECTORY)
        try:
            os.fsync(fd)
        finally:
            os.close(fd)
    raw_report = (json.dumps(updated, indent=2, sort_keys=True) + '\n').encode()
    publish(output, 'reconciliation.json', raw_report)
    fd = os.open(output, os.O_RDONLY | os.O_DIRECTORY)
    try:
        os.fsync(fd)
    finally:
        os.close(fd)
    return updated['visual_capture']


def main():
    parser = argparse.ArgumentParser(description=__doc__)
    parser.add_argument('bundle', type=Path)
    parser.add_argument('--output', type=Path, required=True)
    args = parser.parse_args()
    try:
        print(json.dumps(capture(args.bundle, args.output)))
    except Exception:  # Private URLs, titles and provider errors never enter stdout/stderr.
        print('Visual capture refused or incomplete; retain any partial output for review. '
              'No final reconciliation.json means it is not a completed bundle.', file=sys.stderr)
        return 1
    return 0


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