#!/usr/bin/env -S uv run --script
import argparse
import datetime as dt
import json
import os
from pathlib import Path
import re
import sqlite3
import tempfile
from urllib.error import HTTPError, URLError
from urllib.parse import quote, urlencode, unquote, urlsplit
from urllib.request import Request, urlopen

POLL = 'tech.waow.pollz.poll'
VOTE = 'tech.waow.pollz.vote'
TABLE_COLUMNS = {
    'polls': ('uri', 'repo', 'rkey', 'text', 'options', 'created_at'),
    'votes': ('uri', 'subject', 'option', 'voter', 'created_at'),
    'profiles': ('did', 'handle', 'avatar_url', 'fetched_at'),
}
SCHEMA = '''
CREATE TABLE polls (uri TEXT PRIMARY KEY, repo TEXT NOT NULL, rkey TEXT NOT NULL,
 text TEXT NOT NULL, options TEXT NOT NULL, created_at TEXT NOT NULL);
CREATE TABLE votes (uri TEXT PRIMARY KEY, subject TEXT NOT NULL, option INTEGER NOT NULL,
 voter TEXT NOT NULL, created_at TEXT, UNIQUE(subject, voter));
CREATE INDEX idx_votes_subject ON votes(subject);
CREATE INDEX idx_votes_voter ON votes(subject, voter);
CREATE TABLE profiles (did TEXT PRIMARY KEY, handle TEXT, avatar_url TEXT, fetched_at INTEGER NOT NULL);
'''


class RecoveryError(Exception):
    pass


def now():
    return dt.datetime.now(dt.timezone.utc).isoformat()


def parse_uri(uri):
    if not isinstance(uri, str):
        raise RecoveryError('record URI is not a string')
    match = re.fullmatch(r'at://(did:[^/?#]+)/([^/?#]+)/([^/?#]+)', uri)
    if not match:
        raise RecoveryError('invalid record URI')
    return match.groups()


def discover_actors(seed):
    actors = set()
    warnings = []
    candidates = list(seed.get('actors', []))
    for row in seed['polls']:
        candidates.append(row['repo'])
        try:
            candidates.append(parse_uri(row['uri'])[0])
        except RecoveryError:
            warnings.append({'uri': row['uri'], 'reason': 'invalid historical poll URI'})
    for row in seed['votes']:
        candidates.append(row['voter'])
        for field in ('uri', 'subject'):
            try:
                candidates.append(parse_uri(row[field])[0])
            except RecoveryError:
                warnings.append({'uri': row['uri'], 'field': field, 'reason': 'historical URI does not identify a DID'})
    candidates.extend(row['did'] for row in seed['profiles'])
    for did in candidates:
        if isinstance(did, str) and re.fullmatch(r'did:(plc|web):[A-Za-z0-9._:%-]+', did):
            actors.add(did)
        else:
            warnings.append({'actor': did, 'reason': 'unsupported historical actor identifier'})
    return actors, warnings


def actor_set(seed):
    return discover_actors(seed)[0]


def write_json(path, value):
    path = Path(path)
    path.parent.mkdir(parents=True, exist_ok=True)
    fd, temporary = tempfile.mkstemp(prefix=path.name + '.', dir=path.parent)
    try:
        with os.fdopen(fd, 'w') as output:
            json.dump(value, output, ensure_ascii=False, sort_keys=True, indent=2)
            output.write('\n')
            output.flush()
            os.fsync(output.fileno())
        os.replace(temporary, path)
    finally:
        if os.path.exists(temporary):
            os.unlink(temporary)


def export_seed(source, output):
    if Path(output).exists():
        raise RecoveryError('seed output already exists; preserve the previous export')
    source = Path(source).resolve()
    conn = sqlite3.connect('file:' + quote(str(source), safe='/') + '?mode=ro', uri=True)
    conn.row_factory = sqlite3.Row
    try:
        conn.execute('BEGIN')
        seed = {'version': 1, 'exportedAt': now()}
        for table, columns in TABLE_COLUMNS.items():
            rows = conn.execute('SELECT ' + ','.join(columns) + ' FROM ' + table + ' ORDER BY ' + columns[0])
            seed[table] = [dict(row) for row in rows]
        actors, warnings = discover_actors(seed)
        seed['actors'] = sorted(actors)
        seed['discoveryWarnings'] = warnings
        seed['counts'] = {table: len(seed[table]) for table in TABLE_COLUMNS}
        conn.rollback()
    finally:
        conn.close()
    write_json(output, seed)
    return seed['counts']


def read_seed(path):
    with open(path) as file:
        seed = json.load(file)
    if seed.get('version') != 1:
        raise RecoveryError('unsupported seed version')
    for table, columns in TABLE_COLUMNS.items():
        if not isinstance(seed.get(table), list):
            raise RecoveryError('seed is missing public table ' + table)
        for row in seed[table]:
            if not isinstance(row, dict) or set(row) != set(columns):
                raise RecoveryError('unexpected public seed columns in ' + table)
    actor_set(seed)
    return seed


def read_json(url, timeout=15):
    try:
        with urlopen(Request(url, headers={'Accept': 'application/json', 'User-Agent': 'pollz-public-bootstrap/1'}), timeout=timeout) as response:
            body = response.read(16 * 1024 * 1024 + 1)
            if len(body) > 16 * 1024 * 1024:
                raise RecoveryError('response exceeds bootstrap page limit')
        value = json.loads(body)
        if not isinstance(value, dict):
            raise RecoveryError('response is not a JSON object')
        return value
    except HTTPError as error:
        error.close()
        raise RecoveryError('HTTP ' + str(error.code) + ' (not evidence of record deletion)') from error
    except (URLError, TimeoutError, OSError, ValueError) as error:
        raise RecoveryError(type(error).__name__ + ' fetching public records (not evidence of deletion)') from error


def validate_pds(endpoint, allow_http_loopback=False):
    if not isinstance(endpoint, str):
        raise RecoveryError('PDS endpoint is not a string')
    try:
        parsed = urlsplit(endpoint)
    except ValueError as error:
        raise RecoveryError('invalid PDS origin') from error
    loopback = parsed.hostname in ('127.0.0.1', 'localhost', '::1')
    if parsed.scheme != 'https' and not (allow_http_loopback and parsed.scheme == 'http' and loopback):
        raise RecoveryError('PDS endpoint requires HTTPS')
    if not parsed.hostname or parsed.username or parsed.password or parsed.query or parsed.fragment or parsed.path not in ('', '/'):
        raise RecoveryError('PDS endpoint must be an origin')
    return endpoint.rstrip('/')


def resolve_pds(did, overrides, allow_http_loopback=False):
    if did in overrides:
        return validate_pds(overrides[did], allow_http_loopback)
    if did.startswith('did:plc:'):
        url = 'https://plc.directory/' + quote(did, safe=':')
    elif did.startswith('did:web:'):
        domain = unquote(did[len('did:web:'):])
        if '/' in domain or '@' in domain or '?' in domain or '#' in domain:
            raise RecoveryError('invalid did:web host')
        url = 'https://' + domain + '/.well-known/did.json'
    else:
        raise RecoveryError('unsupported DID method')
    document = read_json(url)
    if document.get('id') != did:
        raise RecoveryError('resolved DID document identity mismatch')
    for service in document.get('service', []):
        if isinstance(service, dict) and service.get('id') in ('#atproto_pds', did + '#atproto_pds') and service.get('type') == 'AtprotoPersonalDataServer':
            return validate_pds(service.get('serviceEndpoint'), allow_http_loopback)
    raise RecoveryError('DID document has no PDS service')


def list_records(pds, did, collection):
    records = []
    cursors = set()
    uris = set()
    cursor = None
    while True:
        params = {'repo': did, 'collection': collection, 'limit': 100}
        if cursor is not None:
            params['cursor'] = cursor
        page = read_json(pds + '/xrpc/com.atproto.repo.listRecords?' + urlencode(params))
        if not isinstance(page.get('records'), list):
            raise RecoveryError('listRecords response has no records array')
        for record in page['records']:
            if not isinstance(record, dict) or 'value' not in record or not isinstance(record.get('cid'), str):
                raise RecoveryError('invalid listRecords entry')
            owner, nsid, _ = parse_uri(record.get('uri'))
            if owner != did or nsid != collection or record['uri'] in uris:
                raise RecoveryError('record identity mismatch or repeated page contents')
            uris.add(record['uri'])
            records.append(record)
        cursor = page.get('cursor')
        if cursor is None or cursor == '':
            return records
        if not isinstance(cursor, str) or cursor in cursors:
            raise RecoveryError('listRecords pagination did not advance')
        cursors.add(cursor)


def insert_record(conn, did, collection, record):
    value = record['value']
    if not isinstance(value, dict):
        return False
    uri = record['uri']
    if collection == POLL:
        if not isinstance(value.get('text'), str) or not isinstance(value.get('options'), list) or not isinstance(value.get('createdAt'), str):
            return False
        conn.execute('INSERT OR IGNORE INTO polls VALUES (?,?,?,?,?,?)',
                     (uri, did, parse_uri(uri)[2], json.dumps(value['text'], ensure_ascii=False), json.dumps(value['options'], ensure_ascii=False), value['createdAt']))
    else:
        option = value.get('option')
        if not isinstance(value.get('subject'), str) or type(option) is not int or option < -(2**31) or option >= 2**31:
            return False
        try:
            _, nsid, _ = parse_uri(value['subject'])
        except RecoveryError:
            return False
        if nsid != POLL:
            return False
        created_at = value.get('createdAt') if isinstance(value.get('createdAt'), str) else None
        conn.execute('''INSERT INTO votes VALUES (?,?,?,?,?)
         ON CONFLICT(subject,voter) DO UPDATE SET uri=excluded.uri,option=excluded.option,created_at=excluded.created_at
         WHERE excluded.created_at > votes.created_at OR votes.created_at IS NULL''',
                     (uri, value['subject'], option, did, created_at))
    return True


def rebuild(seed_path, database, report_path, overrides=None, allow_http_loopback=False):
    if Path(seed_path).resolve() == Path(report_path).resolve():
        raise RecoveryError('report must not overwrite the historical seed')
    seed = read_seed(seed_path)
    database = Path(database)
    if database.exists():
        raise RecoveryError('target database already exists; rebuild requires a fresh path')
    if database.resolve() in (Path(seed_path).resolve(), Path(report_path).resolve()):
        raise RecoveryError('seed, database and report must have distinct paths')
    database.parent.mkdir(parents=True, exist_ok=True)
    conn = sqlite3.connect(database)
    report = {'version': 1, 'seedExportedAt': seed['exportedAt'], 'startedAt': now(),
              'complete': False, 'collections': [], 'failures': [], 'confirmedAbsent': [], 'invalidRecords': [],
              'discoveryWarnings': seed.get('discoveryWarnings', []) + discover_actors(seed)[1],
              'cutover': 'Offline snapshot only. Replay a durable stream overlap from seedExportedAt before serving.'}
    try:
        conn.executescript(SCHEMA)
        for table, columns in TABLE_COLUMNS.items():
            conn.executemany('INSERT INTO ' + table + ' VALUES (' + ','.join('?' for _ in columns) + ')',
                             [tuple(row[column] for column in columns) for row in seed[table]])
        conn.commit()
        pending = actor_set(seed)
        visited = set()
        while pending:
            did = min(pending)
            pending.remove(did)
            if did in visited:
                continue
            visited.add(did)
            try:
                pds = resolve_pds(did, overrides or {}, allow_http_loopback)
            except RecoveryError as error:
                report['failures'].append({'did': did, 'stage': 'resolve', 'error': str(error)})
                continue
            for collection in (POLL, VOTE):
                try:
                    records = list_records(pds, did, collection)
                    table, owner = ('polls', 'repo') if collection == POLL else ('votes', 'voter')
                    existing = {row[0] for row in conn.execute('SELECT uri FROM ' + table + ' WHERE ' + owner + '=?', (did,))}
                    present = {record['uri'] for record in records}
                    absent = sorted(existing - present)
                    with conn:
                        for uri in absent:
                            conn.execute('DELETE FROM ' + table + ' WHERE uri=?', (uri,))
                        for record in records:
                            if not insert_record(conn, did, collection, record):
                                report['invalidRecords'].append(record['uri'])
                            if collection == VOTE and isinstance(record['value'], dict):
                                try:
                                    author, nsid, _ = parse_uri(record['value'].get('subject'))
                                except RecoveryError:
                                    continue
                                if nsid == POLL and author not in visited:
                                    pending.add(author)
                    report['confirmedAbsent'].extend(absent)
                    report['collections'].append({'did': did, 'collection': collection, 'records': len(records)})
                except (RecoveryError, sqlite3.Error) as error:
                    report['failures'].append({'did': did, 'collection': collection, 'error': str(error)})
        report['actors'] = sorted(visited)
        report['counts'] = {table: conn.execute('SELECT COUNT(*) FROM ' + table).fetchone()[0] for table in TABLE_COLUMNS}
        report['complete'] = not report['failures'] and not report['invalidRecords'] and not report['discoveryWarnings']
        report['finishedAt'] = now()
        write_json(report_path, report)
    finally:
        conn.close()
    return report


def main():
    parser = argparse.ArgumentParser(description='Preserve public Pollz history and rebuild a fresh read-only AppView from PDS records.')
    commands = parser.add_subparsers(dest='command', required=True)
    export = commands.add_parser('export')
    export.add_argument('--source', required=True)
    export.add_argument('--output', required=True)
    build = commands.add_parser('rebuild')
    build.add_argument('--seed', required=True)
    build.add_argument('--database', required=True)
    build.add_argument('--report', required=True)
    build.add_argument('--pds-map')
    build.add_argument('--allow-http-loopback', action='store_true')
    args = parser.parse_args()
    try:
        if args.command == 'export':
            print(json.dumps(export_seed(args.source, args.output)))
        else:
            overrides = json.loads(Path(args.pds_map).read_text()) if args.pds_map else {}
            report = rebuild(args.seed, args.database, args.report, overrides, args.allow_http_loopback)
            print(json.dumps({'complete': report['complete'], 'counts': report['counts'], 'failedCollections': len(report['failures']), 'confirmedAbsent': len(report['confirmedAbsent'])}))
            return 0 if report['complete'] else 2
    except (RecoveryError, sqlite3.Error, OSError, ValueError) as error:
        parser.exit(1, 'bootstrap: ' + str(error) + '\n')
    return 0


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