#!/usr/bin/env python3
"""Mirror full raw OSV JSON files per ecosystem with conditional HTTP requests.
No third-party Python dependencies. All-or-nothing validation per ecosystem.
"""
import argparse
import io
import json
import os
import re
import shutil
import sys
import tempfile
import time
import urllib.error
import urllib.parse
import urllib.request
import zipfile
from datetime import datetime, timezone
from pathlib import Path

ROOT = Path(__file__).resolve().parents[1]
BASE = 'https://storage.googleapis.com/osv-vulnerabilities'
DEFAULT = 'PyPI,npm,Go,Maven,NuGet,crates.io,Packagist,RubyGems,GitHub Actions,OSS-Fuzz'
VALID_ECOSYSTEM = re.compile(r'^[A-Za-z0-9._ -]{1,60}$')
VALID_ID = re.compile(r'^[A-Za-z0-9._-]{2,150}$')

def atomic_write(path, data):
    path.parent.mkdir(parents=True, exist_ok=True)
    if path.exists() and path.read_bytes() == data:
        return False
    tmp = path.with_name(path.name + '.tmp')
    try:
        tmp.write_bytes(data)
        os.replace(tmp, path)
    finally:
        tmp.unlink(missing_ok=True)
    return True

def fetch(url, etag=None, last_modified=None):
    headers = {'User-Agent': 'vuln-osv-gitea-mirror/1.0'}
    if etag:
        headers['If-None-Match'] = etag
    elif last_modified:
        headers['If-Modified-Since'] = last_modified
    for attempt in range(5):
        try:
            with urllib.request.urlopen(urllib.request.Request(url, headers=headers), timeout=180) as response:
                # Download into a temp file; don't hold huge archives in RAM.
                with tempfile.TemporaryFile() as temp:
                    shutil.copyfileobj(response, temp)
                    temp.seek(0)
                    body = temp.read()
                return body, response.headers.get('ETag'), response.headers.get('Last-Modified')
        except urllib.error.HTTPError as exc:
            if exc.code == 304:
                return None, etag, last_modified
            if exc.code not in (408, 429, 500, 502, 503, 504) or attempt == 4:
                raise
        except (urllib.error.URLError, TimeoutError):
            if attempt == 4:
                raise
        time.sleep(min(120, 3 * 2**attempt))
    raise RuntimeError('Unreachable')

def process(ecosystem, state, allow_deletions):
    if not VALID_ECOSYSTEM.fullmatch(ecosystem):
        raise ValueError(f'Invalid ecosystem name: {ecosystem!r}')
    url = f'{BASE}/{urllib.parse.quote(ecosystem)}/all.zip'
    previous = state.get('ecosystems', {}).get(ecosystem, {})
    body, etag, modified = fetch(url, previous.get('etag'), previous.get('last_modified'))
    if body is None:
        print(f'{ecosystem}: unchanged (304)', flush=True)
        return False
    # Validate all records before changing the working tree.
    records = {}
    with zipfile.ZipFile(io.BytesIO(body)) as archive:
        for info in archive.infolist():
            if info.is_dir() or not info.filename.endswith('.json'):
                continue
            if info.file_size > 32*1024*1024:
                raise ValueError(f'OSV JSON too large: {info.filename}')
            identifier = Path(info.filename).stem
            if not VALID_ID.fullmatch(identifier):
                raise ValueError(f'Unsafe ID: {identifier!r}')
            content = archive.read(info)
            entry = json.loads(content)
            if not isinstance(entry, dict) or entry.get('id') != identifier:
                raise ValueError(f'Invalid OSV record: {info.filename}')
            if identifier in records and records[identifier] != content:
                raise ValueError(f'Duplicate OSV ID: {identifier}')
            records[identifier] = content
    if not records:
        raise ValueError('Empty OSV archive: refusing to delete data')
    directory = ROOT / 'data' / ecosystem
    existing = {p.stem: p for p in directory.glob('*.json')}
    written = sum(atomic_write(directory / f'{ident}.json', raw) for ident, raw in records.items())
    deleted = 0
    if allow_deletions:
        for ident, path in existing.items():
            if ident not in records:
                path.unlink()
                deleted += 1
    state.setdefault('ecosystems', {})[ecosystem] = {
        'etag': etag, 'last_modified': modified,
        'records': len(records), 'last_success': datetime.now(timezone.utc).isoformat()
    }
    print(f'{ecosystem}: {len(records)} records, {written} written, {deleted} deleted', flush=True)
    return True

def main():
    parser = argparse.ArgumentParser()
    parser.add_argument('--ecosystems', default=os.getenv('OSV_ECOSYSTEMS', DEFAULT))
    parser.add_argument('--allow-deletions', action='store_true', help='Remove records no longer in verified full snapshots')
    args = parser.parse_args()
    ecosystems = [s.strip() for s in args.ecosystems.split(',') if s.strip()]
    if len(set(ecosystems)) != len(ecosystems):
        raise ValueError('Duplicate ecosystem')
    state_path = ROOT / 'state' / 'sync.json'
    state = json.loads(state_path.read_text()) if state_path.exists() else {'version': 1, 'ecosystems': {}}
    for ecosystem in ecosystems:
        # Persist each successfully downloaded ecosystem. A retry resumes the rest.
        if process(ecosystem, state, args.allow_deletions):
            atomic_write(state_path, (json.dumps(state, ensure_ascii=False, sort_keys=True, indent=2) + '\n').encode())

if __name__ == '__main__':
    main()