#!/usr/bin/env python3
"""Conservative SQLite index from CVEProject, NVD and OSV raw JSON.
Do not mistake this index for a complete vulnerability matching engine.
"""
import argparse
import hashlib
import json
import os
import re
import shutil
import sqlite3
import subprocess
import tempfile
import sys
from datetime import datetime, timezone
from pathlib import Path

# Make the top-level importers package importable when invoked as scripts/build.py.
sys.path.insert(0, str(Path(__file__).resolve().parents[1]))
from importers import IMPORTERS

CVE = re.compile(r'^CVE-\d{4}-\d{4,}$', re.I)
FILES = ('vulnerabilities','aliases','affected_packages','cpe_matches')


def canonical(obj):
    return json.dumps(obj, ensure_ascii=False, sort_keys=True, separators=(',', ':'))


def get_git_revision(path):
    return subprocess.check_output(['git', '-C', str(path), 'rev-parse', 'HEAD'], text=True).strip()


def input_files(repo):
    d = repo / 'data'
    if not d.exists():
        return

    # Nur tatsächliche CVE-Datensätze einlesen
    if repo.name == 'vuln-cve':
        yield from sorted(d.rglob('CVE-*.json'))
        return

    # Andere Quellen unverändert verarbeiten
    yield from sorted(d.rglob('*.json'))


def stream_source(db, name, path):
    importer = IMPORTERS[name]
    count = 0
    bad = 0
    scanned = 0
    for file in input_files(path):
        scanned += 1
        try:
            with file.open('r', encoding='utf-8') as stream:
                obj = json.load(stream)
            count += int(bool(importer(db, obj)))
        except (ValueError, UnicodeError, TypeError, KeyError, sqlite3.Error) as exc:
            bad += 1
            if bad <= 5:
                print(f'INVALID {name} {file}: {exc}', flush=True)
    if bad:
        raise RuntimeError(f'{name}: {bad} invalid records; database not published')
    if name == 'vuln-nvd' and count == 0:
        raise RuntimeError('vuln-nvd: no usable documents; check source data and JSON shape')
    if name in ('vuln-cve', 'vuln-osv') and count == 0:
        raise RuntimeError(f'{name}: no usable documents; refusing to publish')
    print(f'{name}: scanned={scanned} accepted={count}', flush=True)
    return count


def dump_table(db,table):
    # memory efficient cursor; a canonical row object for deltas
    cols=[r[1] for r in db.execute(f'PRAGMA table_info({table})')]
    for row in db.execute(f'SELECT * FROM {table} ORDER BY 1'):
        yield row[0],dict(zip(cols,row))


def make_delta(old_path,new_db,path):
    old=sqlite3.connect(old_path)
    with path.open('w',encoding='utf-8') as f:
        for table in FILES:
            prev=dump_table(old,table); curr=dump_table(new_db,table)
            a=next(prev,None); b=next(curr,None)
            while a is not None or b is not None:
                if a is not None and (b is None or a[0]<b[0]):
                    f.write(canonical({'table':table,'op':'delete','key':a[0]})+'\n'); a=next(prev,None)
                elif b is not None and (a is None or b[0]<a[0]):
                    f.write(canonical({'table':table,'op':'upsert','row':b[1]})+'\n'); b=next(curr,None)
                else:
                    if a[1]!=b[1]:
                        f.write(canonical({'table':table,'op':'upsert','row':b[1]})+'\n')
                    a=next(prev,None);b=next(curr,None)
    old.close()


def compress(src,out):
    subprocess.run(['zstd','-q','-f','-T1','-19',str(src),'-o',str(out)],check=True)


def build_signature(source_revs, schema_path):
    """Hash source revisions plus ALL executable importer logic and schema."""
    root = Path(__file__).resolve().parents[1]
    h = hashlib.sha256()
    inputs = sorted((root / 'scripts').glob('*.py'))
    inputs += sorted((root / 'importers').glob('*.py'))
    inputs += [schema_path.resolve()]
    for path in inputs:
        if not path.is_file():
            raise RuntimeError(f'Missing build input: {path}')
        logical_name = path.relative_to(root).as_posix() if path.is_relative_to(root) else 'sql/schema.sql'
        h.update(logical_name.encode('utf-8') + b'\0')
        h.update(hashlib.sha256(path.read_bytes()).digest())
    h.update(canonical(source_revs).encode('utf-8'))
    return h.hexdigest()


def existing_release_ok(output, manifest):
    """Only skip if current release artifacts exist and match recorded hash."""
    dbfile, packed = output / 'master.db', output / 'master.db.zst'
    if not dbfile.is_file() or not packed.is_file() or not manifest.get('master_sha256'):
        return False
    h = hashlib.sha256()
    with packed.open('rb') as stream:
        for block in iter(lambda: stream.read(1024 * 1024), b''):
            h.update(block)
    return h.hexdigest() == manifest['master_sha256']


def main():
    ap=argparse.ArgumentParser()
    ap.add_argument('--sources',type=Path,required=True)
    ap.add_argument('--output',type=Path,required=True)
    ap.add_argument('--schema',type=Path,default=Path(__file__).resolve().parents[1]/'sql/schema.sql')
    ap.add_argument('--force', action='store_true', help='Force rebuild even when input signature is unchanged')
    args=ap.parse_args()
    args.output.mkdir(parents=True,exist_ok=True)
    active=[(p.name,p) for p in sorted(args.sources.iterdir()) if p.is_dir() and (p/'.git').exists()]
    if not active:
        raise SystemExit('No source repositories checked out; nothing published')
    # Fail closed if expected core sources are missing.
    for name in IMPORTERS:
        if name not in [n for n,_ in active]:
            raise SystemExit(f'Missing core source: {name}')
    source_revs={name:get_git_revision(path) for name,path in active}
    signature = build_signature(source_revs, args.schema)
    manifest_path = args.output / 'manifest.json'
    previous_manifest = None
    if manifest_path.is_file():
        try:
            previous_manifest = json.loads(manifest_path.read_text(encoding='utf-8'))
        except (ValueError, UnicodeError) as exc:
            raise RuntimeError(f'Invalid existing manifest, refusing to overwrite: {exc}') from exc
    if (not args.force and previous_manifest
            and previous_manifest.get('build_signature') == signature
            and existing_release_ok(args.output, previous_manifest)):
        print(f"NO_CHANGES: version={previous_manifest['version']} inputs unchanged; master and delta retained", flush=True)
        return
    if not args.force and previous_manifest and previous_manifest.get('build_signature') == signature:
        print('WARNING: unchanged inputs but missing/invalid release artifact; rebuilding', flush=True)
    old_db=args.output/'master.db' 
    with tempfile.TemporaryDirectory(prefix='vuln-build-',dir=args.output) as td:
        tmp=Path(td); new=tmp/'master.db'; db=sqlite3.connect(new)
        db.executescript(args.schema.read_text(encoding='utf-8'))
        stats={}
        for name,path in active:
            if name in IMPORTERS:
                stats[name]=stream_source(db,name,path)
            else:
                stats[name]=0
            db.execute('INSERT INTO sources VALUES(?,?,?,?)',(name,source_revs[name],int(name in IMPORTERS),stats[name]))
            print(f'{name}: indexed={stats[name]} revision={source_revs[name][:12]}',flush=True)
        if not any(stats.get(n) for n in IMPORTERS):
            raise RuntimeError('Core source indexes are empty; refusing to publish')
        old_version=None
        if old_db.is_file():
            with sqlite3.connect(old_db) as o:
                row=o.execute("SELECT value FROM meta WHERE key='version'").fetchone()
                old_version=row[0] if row else None
        version=datetime.now(timezone.utc).strftime('%Y%m%dT%H%M%S%fZ')
        for k,v in (('version',version),('schema_version','1')):
            db.execute('INSERT INTO meta VALUES(?,?)',(k,v))
        db.commit(); db.execute('PRAGMA optimize'); db.execute('VACUUM'); db.commit()
        assert db.execute('PRAGMA integrity_check').fetchone()[0]=='ok'
        db.close()
        manifest={'version':version,'schema_version':1,'previous':old_version,'source_commits':source_revs,
                  'build_signature':signature,
                  'indexed_documents':stats,'format':'sqlite+zstd','delta_format':'jsonl+zstd',
                  'note':'CVE/NVD/OSV v1 index; AND/NOT CPE conditions are not evaluated as independent matches',
                  'active_importers':sorted(IMPORTERS)}
        compress(new,tmp/'master.db.zst')
        if old_version:
            with sqlite3.connect(new) as nd:
                make_delta(old_db,nd,tmp/'delta.jsonl')
            compress(tmp/'delta.jsonl',tmp/f'delta-{old_version}-to-{version}.jsonl.zst')
        # Atomically replace individual artifacts; manifest written last as commit marker.
        os.replace(new,old_db)
        os.replace(tmp/'master.db.zst',args.output/'master.db.zst')
        if old_version:
            os.replace(tmp/f'delta-{old_version}-to-{version}.jsonl.zst',args.output/f'delta-{old_version}-to-{version}.jsonl.zst')
        with (args.output/'master.db.zst').open('rb') as hf:
            h=hashlib.sha256()
            for chunk in iter(lambda: hf.read(1024*1024),b''):
                h.update(chunk)
            manifest['master_sha256']=h.hexdigest()
        (tmp/'manifest.json').write_text(json.dumps(manifest,indent=2,ensure_ascii=False)+'\n',encoding='utf-8')
        os.replace(tmp/'manifest.json',args.output/'manifest.json')
        print(f'Build complete: {version}',flush=True)

if __name__=='__main__':
    main()