"""Resumable, unpublished COPY followed by complete ordered content verification."""
import argparse
from contextlib import closing
import fcntl
import hashlib
import json
from pathlib import Path
import re
import shutil
import sqlite3
import uuid
import zlib

from paper_atlas.db.submissions import timestamp
from paper_atlas.postgres.client import connect, SEARCH_VECTOR
from paper_atlas.postgres.schema import TABLES, inspect, fingerprint, quote, select_sql, table_sql
from paper_atlas.postgres import text


def digest_file(path):
    digest = hashlib.sha256()
    with path.open('rb') as stream:
        while block := stream.read(8 * 1024 ** 2):
            digest.update(block)
    return digest.hexdigest()


def update_hash(digest, row):
    # Raw JSON columns remain text: whitespace, escapes and source evidence survive.
    data = json.dumps(tuple(row), ensure_ascii=False, separators=(',', ':')).encode()
    digest.update(len(data).to_bytes(8, 'big')); digest.update(data)


def save(path, result):
    path.parent.mkdir(parents=True, exist_ok=True)
    part = path.with_suffix(path.suffix + '.part')
    part.write_text(json.dumps(result, ensure_ascii=False, indent=2) + '\n')
    part.replace(path)


def free_space(path, minimum):
    if path and shutil.disk_usage(path).free < minimum * 1024 ** 3:
        raise RuntimeError('Migration storage headroom fell below the configured minimum')


def qualify(schema, table):
    return quote(schema) + '.' + quote(table)


def add_constraints(pg, stage, plan):
    def constraint(table, name, clause):
        exists = pg.execute('''SELECT 1 FROM pg_constraint co JOIN pg_class cl ON cl.oid=co.conrelid
            JOIN pg_namespace ns ON ns.oid=cl.relnamespace WHERE ns.nspname=%s AND cl.relname=%s AND co.conname=%s''',
                            (stage, table, name)).fetchone()
        if not exists:
            pg.execute('ALTER TABLE ' + qualify(stage, table) + ' ADD CONSTRAINT ' + quote(name) + ' ' + clause)

    for table, spec in plan.items():
        constraint(table, table + '_pkey', 'PRIMARY KEY (' + ','.join(map(quote, spec['pk'])) + ')')
        for number, index in enumerate(spec['indexes']):
            if index['origin'] == 'u':
                constraint(table, table + '_unique_' + str(number), 'UNIQUE (' + ','.join(quote(n) for n, desc in index['terms']) + ')')
            else:
                terms = ','.join(quote(n) + (' DESC' if desc else '') for n, desc in index['terms'])
                pg.execute('CREATE ' + ('UNIQUE ' if index['unique'] else '') + 'INDEX IF NOT EXISTS ' +
                           quote(index['name']) + ' ON ' + qualify(stage, table) + '(' + terms + ')')
    for table, spec in plan.items():
        for number, fk in enumerate(spec['foreign_keys']):
            clause = ('FOREIGN KEY (' + ','.join(map(quote, fk['columns'])) + ') REFERENCES ' +
                      qualify(stage, fk['table']) + '(' + ','.join(map(quote, fk['references'])) + ') NOT VALID')
            constraint(table, table + '_fk_' + str(number), clause)
        for number, check in enumerate(spec['checks']):
            constraint(table, table + '_check_' + str(number), 'CHECK (' + check + ') NOT VALID')
    constraints = pg.execute('''SELECT cl.relname,co.conname FROM pg_constraint co
        JOIN pg_class cl ON cl.oid=co.conrelid JOIN pg_namespace ns ON ns.oid=cl.relnamespace
        WHERE ns.nspname=%s AND NOT co.convalidated ORDER BY cl.relname,co.conname''', (stage,)).fetchall()
    for table, name in constraints:
        pg.execute('ALTER TABLE ' + qualify(stage, table) + ' VALIDATE CONSTRAINT ' + quote(name))


def sequences(pg, source, stage, plan):
    old = dict(source.execute('SELECT name,seq FROM sqlite_sequence'))
    history = {'papers': ('paper_merges', 'former_paper_id'),
               'submissions': ('submission_merges', 'former_submission_id'),
               'decisions': ('decision_merges', 'former_decision_id')}
    result = {}
    for table, spec in plan.items():
        if spec['pk'] != ['id']:
            continue
        value = max(source.execute('SELECT coalesce(max(id),0) FROM ' + quote(table)).fetchone()[0], old.get(table, 0))
        if table in history:
            audit, column = history[table]
            value = max(value, source.execute('SELECT coalesce(max(' + quote(column) + '),0) FROM ' + quote(audit)).fetchone()[0])
        sequence = pg.execute('SELECT pg_get_serial_sequence(%s,%s)', (qualify(stage, table), 'id')).fetchone()[0]
        pg.execute('SELECT setval(%s,%s,%s)', (sequence, max(value, 1), value > 0))
        result[table] = value
    return result


def upgrade_text_codec(pg, control, stage, plan):
    """Upgrade committed legacy staging rows once; never recopy their source."""
    changed = {}
    with pg.transaction():
        copied = {r[0] for r in pg.execute('SELECT name FROM ' + qualify(control, 'tables'))}
        if not copied <= plan.keys():
            raise ValueError('Unknown committed staging tables')
        # An older run may have reached constraint validation. Replacing SOH
        # can temporarily equal another legacy key even though the final mapping
        # is bijective. Rebuild unpublished constraints/indexes after the upgrade.
        constraints = pg.execute('''SELECT cl.relname,co.conname FROM pg_constraint co
            JOIN pg_class cl ON cl.oid=co.conrelid JOIN pg_namespace ns ON ns.oid=cl.relnamespace
            WHERE ns.nspname=%s AND co.contype IN ('f','p','u')
            ORDER BY CASE WHEN co.contype='f' THEN 0 ELSE 1 END''', (stage,)).fetchall()
        for table, name in constraints:
            pg.execute('ALTER TABLE ' + qualify(stage, table) + ' DROP CONSTRAINT ' + quote(name))
        indexes = pg.execute('''SELECT cl.relname FROM pg_index i JOIN pg_class cl ON cl.oid=i.indexrelid
            JOIN pg_namespace ns ON ns.oid=cl.relnamespace WHERE ns.nspname=%s AND i.indisunique''', (stage,)).fetchall()
        for name, in indexes:
            pg.execute('DROP INDEX ' + qualify(stage, name))
        for table in plan:
            if table not in copied:
                continue
            columns = [quote(c['name']) for c in plan[table]['columns'] if c['type'] == 'TEXT']
            if columns:
                assignments = ','.join(c + '=replace(' + c + ',%s,%s)' for c in columns)
                predicate = ' OR '.join('strpos(' + c + ',%s)>0' for c in columns)
                params = [v for c in columns for v in (text.ESCAPE, text.encode(text.ESCAPE))]
                params += [text.ESCAPE] * len(columns)
                changed[table] = pg.execute('UPDATE ' + qualify(stage, table) + ' SET ' + assignments +
                                            ' WHERE ' + predicate, params).rowcount
        pg.execute('UPDATE ' + qualify(control, 'run') + ' SET text_codec=%s WHERE id=1', (text.CODEC,))
    return changed


def migrate(db, service, out, *, schema='research', storage_path=None, min_free_gib=8):
    if not re.fullmatch(r'[a-z][a-z0-9_]{0,35}', schema):
        raise ValueError('Use a short lowercase target schema name')
    db, out = Path(db).resolve(), Path(out).resolve()
    if not db.is_file() or out == db:
        raise ValueError('Existing source and a distinct report are required')
    if out.exists() and json.loads(out.read_text()).get('status') == 'complete':
        raise ValueError('Completed report must be preserved')
    free_space(storage_path, min_free_gib)
    control = schema + '_migration'
    report = {'status': 'preflight', 'source': str(db), 'service': service, 'schema': schema,
              'started_at': timestamp(), 'tables': {}, 'scope': 'All 21 business tables; FTS5 internals rebuilt as native PostgreSQL GIN.'}
    with Path(str(db) + '.maintenance.lock').open('a') as lock:
        fcntl.flock(lock, fcntl.LOCK_EX | fcntl.LOCK_NB)
        if Path(str(db) + '-wal').exists() and Path(str(db) + '-wal').stat().st_size:
            raise ValueError('Freeze writers and checkpoint SQLite WAL before migrating')
        with closing(sqlite3.connect(db.as_uri() + '?mode=ro', uri=True)) as source, closing(connect(service)) as pg:
            source.execute('PRAGMA query_only=ON'); source.execute('PRAGMA cache_size=-16384'); source.execute('BEGIN')
            plan = inspect(source)
            stat = db.stat()
            identity = {'file_sha256': digest_file(db), 'bytes': stat.st_size, 'mtime_ns': stat.st_mtime_ns,
                        'schema_sha256': fingerprint(plan)}
            report['source_identity'] = identity
            advisory = zlib.crc32(schema.encode()) & 0x7fffffff
            if not pg.execute('SELECT pg_try_advisory_lock(%s,%s)', (615173, advisory)).fetchone()[0]:
                raise RuntimeError('Another migration holds the target lock')
            if pg.execute('SELECT 1 FROM pg_namespace WHERE nspname=%s', (schema,)).fetchone():
                raise ValueError('Target schema already exists; published data will never be overwritten')
            pg.execute('CREATE SCHEMA IF NOT EXISTS ' + quote(control))
            pg.execute('CREATE TABLE IF NOT EXISTS ' + qualify(control, 'run') +
                       '(id INTEGER PRIMARY KEY CHECK(id=1), source_json TEXT NOT NULL, stage TEXT NOT NULL, status TEXT NOT NULL)')
            pg.execute('CREATE TABLE IF NOT EXISTS ' + qualify(control, 'tables') +
                       '(name TEXT PRIMARY KEY, rows BIGINT NOT NULL, sha256 TEXT NOT NULL, verified BOOLEAN NOT NULL DEFAULT false)')
            pg.execute('ALTER TABLE ' + qualify(control, 'run') + " ADD COLUMN IF NOT EXISTS text_codec TEXT NOT NULL DEFAULT 'legacy'")
            run = pg.execute('SELECT source_json,stage,status,text_codec FROM ' + qualify(control, 'run') + ' WHERE id=1').fetchone()
            if run:
                if json.loads(run[0]) != identity or run[2] == 'complete':
                    raise ValueError('Resume requires the exact same unpublished SQLite snapshot')
                stage = run[1]
                if run[3] == 'legacy':
                    report['legacy_text_upgrade_rows'] = upgrade_text_codec(pg, control, stage, plan)
                elif run[3] != text.CODEC:
                    raise ValueError('Unsupported staging text storage codec')
            else:
                stage = schema + '_stage_' + uuid.uuid4().hex[:12]
                with pg.transaction():
                    pg.execute('CREATE SCHEMA ' + quote(stage))
                    for table, spec in plan.items():
                        pg.execute(table_sql(stage, table, spec))
                    pg.execute('INSERT INTO ' + qualify(control, 'run') +
                               '(id,source_json,stage,status,text_codec) VALUES(1,%s,%s,%s,%s)',
                               (json.dumps(identity, sort_keys=True), stage, 'copying', text.CODEC))
            report.update(status='copying', stage=stage, text_codec=text.CODEC,
                          content_hash_scope='Decoded original SQLite values; storage escapes are reversible.')
            save(out, report)
            try:
                for table, spec in plan.items():
                    free_space(storage_path, min_free_gib)
                    previous = pg.execute('SELECT rows,sha256 FROM ' + qualify(control, 'tables') + ' WHERE name=%s', (table,)).fetchone()
                    if previous:
                        rows, digest = previous
                    else:
                        digest = hashlib.sha256(); rows = 0
                        with pg.transaction():
                            # Only this run's unpublished, incomplete table can be reset.
                            pg.execute('TRUNCATE ' + qualify(stage, table))
                            columns = ','.join(quote(c['name']) for c in spec['columns'])
                            with pg.cursor().copy('COPY ' + qualify(stage, table) + '(' + columns + ') FROM STDIN') as copy:
                                for row in source.execute(select_sql(None, table, spec)):
                                    copy.write_row(tuple(map(text.encode, row))); update_hash(digest, row); rows += 1
                                    if rows % 10000 == 0:
                                        free_space(storage_path, min_free_gib)
                                    if rows % 100000 == 0:
                                        print(json.dumps({'stage': 'copying', 'table': table, 'rows': rows}), flush=True)
                            digest = digest.hexdigest()
                            pg.execute('INSERT INTO ' + qualify(control, 'tables') + '(name,rows,sha256) VALUES(%s,%s,%s)', (table, rows, digest))
                    report['tables'][table] = {'source_rows': rows, 'source_sha256': digest}
                    save(out, report)
                    print(json.dumps({'stage': 'copied', 'table': table, 'rows': rows}), flush=True)
                report['status'] = 'validating_constraints'; save(out, report)
                add_constraints(pg, stage, plan)
                report['sequences'] = sequences(pg, source, stage, plan)
                report['status'] = 'verifying_contents'; save(out, report)
                for table, spec in plan.items():
                    digest = hashlib.sha256(); storage_digest = hashlib.sha256(); rows = 0
                    with pg.transaction(), pg.cursor(name='verify_' + table) as cursor:
                        cursor.itersize = 512
                        cursor.execute(select_sql(stage, table, spec))
                        for row in cursor:
                            update_hash(storage_digest, row)
                            update_hash(digest, tuple(map(text.decode, row))); rows += 1
                    entry = report['tables'][table]
                    entry.update(target_rows=rows, target_sha256=digest.hexdigest(),
                                 target_storage_sha256=storage_digest.hexdigest())
                    if rows != entry['source_rows'] or digest.hexdigest() != entry['source_sha256']:
                        raise ValueError('Complete content verification failed for ' + table)
                    pg.execute('UPDATE ' + qualify(control, 'tables') + ' SET verified=true WHERE name=%s', (table,))
                    save(out, report)
                    print(json.dumps({'stage': 'verified', 'table': table, 'rows': rows}), flush=True)
                report['status'] = 'building_search_index'; save(out, report)
                free_space(storage_path, min_free_gib)
                pg.execute('CREATE INDEX IF NOT EXISTS papers_search_gin ON ' + qualify(stage, 'papers') + ' USING gin (' + SEARCH_VECTOR + ')')
                search = pg.execute('''SELECT am.amname,i.indisvalid,i.indisready,t.relname
                    FROM pg_index i JOIN pg_class ix ON ix.oid=i.indexrelid
                    JOIN pg_class t ON t.oid=i.indrelid JOIN pg_namespace ns ON ns.oid=ix.relnamespace
                    JOIN pg_am am ON am.oid=ix.relam WHERE ns.nspname=%s AND ix.relname='papers_search_gin' ''', (stage,)).fetchone()
                if search != ('gin', True, True, 'papers'):
                    raise ValueError('Native full-text index is not valid and ready')
                for table in TABLES:
                    pg.execute('ANALYZE ' + qualify(stage, table))
                invalid = pg.execute('''SELECT count(*) FROM pg_constraint co JOIN pg_class cl ON cl.oid=co.conrelid
                    JOIN pg_namespace ns ON ns.oid=cl.relnamespace WHERE ns.nspname=%s AND NOT co.convalidated''', (stage,)).fetchone()[0]
                if invalid:
                    raise ValueError('Unvalidated constraints remain')
                if (db.stat().st_size, db.stat().st_mtime_ns) != (identity['bytes'], identity['mtime_ns']):
                    raise ValueError('Source changed during migration')
                if Path(str(db) + '-wal').exists() and Path(str(db) + '-wal').stat().st_size:
                    raise ValueError('A source writer committed during migration')
                report.update(status='prepared', foreign_keys='all_validated', checks='all_validated', native_search_index=True)
                save(out, report)
                with pg.transaction():
                    pg.execute('COMMENT ON SCHEMA ' + quote(stage) + " IS 'paper-atlas:text-codec=" + text.CODEC + "'")
                    pg.execute('ALTER SCHEMA ' + quote(stage) + ' RENAME TO ' + quote(schema))
                    pg.execute('UPDATE ' + qualify(control, 'run') + " SET status='complete' WHERE id=1")
                report.update(status='complete', finished_at=timestamp())
                save(out, report)
                return report
            except BaseException as error:
                report.update(status='error', error_type=type(error).__name__, error=str(error), failed_at=timestamp())
                save(out, report)
                raise


def main(argv=None, *, prog=None):
    parser = argparse.ArgumentParser(prog=prog)
    parser.add_argument('--db', required=True, help='Frozen SQLite source')
    parser.add_argument('--pg-service', required=True, help='Protected libpq service name')
    parser.add_argument('--schema', default='research')
    parser.add_argument('--out', required=True)
    parser.add_argument('--storage-path', help='Target filesystem path for free-space checks')
    parser.add_argument('--min-free-gib', type=float, default=8)
    args = parser.parse_args(argv)
    result = migrate(args.db, args.pg_service, args.out, schema=args.schema,
                     storage_path=args.storage_path, min_free_gib=args.min_free_gib)
    print(json.dumps(result, ensure_ascii=False), flush=True)