"""Merge disjoint compact-index partitions and validate their inventory."""
import argparse
from datetime import datetime, timezone
import hashlib
import json
from pathlib import Path
import sqlite3

from compact_index import DDL, finalize_indexes
from remote_catalog import BUCKET


def main():
    p=argparse.ArgumentParser(description=__doc__)
    p.add_argument('--inventory',type=Path,required=True)
    p.add_argument('--parts',type=Path,nargs='+',required=True)
    p.add_argument('--output',type=Path,required=True)
    args=p.parse_args()
    work=args.output.with_suffix('.merging.sqlite')
    work.unlink(missing_ok=True)
    c=sqlite3.connect(work)
    c.execute('PRAGMA journal_mode=OFF'); c.execute('PRAGMA synchronous=OFF')
    c.execute('PRAGMA cache_size=-262144'); c.executescript(DDL)
    file_offset=row_offset=0
    for path in args.parts:
        c.execute('ATTACH DATABASE ? AS part',(str(path),))
        manifest=json.loads(c.execute("SELECT value FROM part.metadata WHERE key='manifest'").fetchone()[0])
        if manifest.get('schema_version')!=3: raise ValueError(f'Unexpected schema in {path}')
        c.execute('INSERT OR IGNORE INTO contexts(assembly_accession,assembly_base,metadata_json) SELECT assembly_accession,assembly_base,metadata_json FROM part.contexts')
        c.execute('CREATE TEMP TABLE context_map(old INTEGER PRIMARY KEY,new INTEGER)')
        c.execute('INSERT INTO context_map SELECT p.id,c.id FROM part.contexts p JOIN contexts c ON c.metadata_json=p.metadata_json')
        c.execute('INSERT INTO files SELECT id+?,path,hash,size,rows,row_groups,max_group_bytes,bytes_read,range_reads FROM part.files',(file_offset,))
        c.execute('''INSERT INTO segment_data SELECT s.id+?,s.record_name,s.record_base,s.segment_start_bp,s.segment_end_bp,
                     s.aligned_bp_length,s.segment_index,s.segment_count,s.file_id+?,s.row_group,s.row_in_group,m.new,s.source_key
                     FROM part.segment_data s JOIN context_map m ON m.old=s.context_id''',(row_offset,file_offset))
        file_offset=c.execute('SELECT max(id) FROM files').fetchone()[0] or 0
        row_offset=c.execute('SELECT max(id) FROM segment_data').fetchone()[0] or 0
        c.execute('DROP TABLE context_map'); c.commit(); c.execute('DETACH DATABASE part')
        print('Merged',path,'last_row',row_offset,flush=True)
    expected={r['path']:(r['hash'],r['size']) for r in map(json.loads,args.inventory.open())}
    actual={r[0]:(r[1],r[2]) for r in c.execute('SELECT path,hash,size FROM files')}
    if actual!=expected: raise ValueError('Merged object inventory does not match listing')
    print('Building search indexes',flush=True)
    finalize_indexes(c); c.commit()
    rows,bases=c.execute('SELECT count(*),sum(segment_end_bp-segment_start_bp) FROM segment_data').fetchone()
    if rows!=c.execute('SELECT sum(rows) FROM files').fetchone()[0]: raise ValueError('Row totals differ')
    assemblies=[r[0] for r in c.execute('SELECT DISTINCT assembly_accession FROM contexts ORDER BY assembly_accession')]
    examples={}
    for cid,encoded in c.execute('SELECT id,metadata_json FROM contexts'):
        division=json.loads(encoded)['division']
        if division not in examples:
            row=c.execute('SELECT record_name FROM segment_data WHERE context_id=? ORDER BY id LIMIT 1',(cid,)).fetchone()
            if row: examples[division]=row[0]
    manifest={'schema_version':3,'bucket_id':BUCKET,'created_at':datetime.now(timezone.utc).isoformat(),
              'scope':'All published annotation Parquet objects in the dated inventory; assemblies may be partial.',
              'full_snapshot':True,'source_count':len(expected),'sources':[],
              'source_ledger':'SQLite files table contains every object path, hash, size and row-group summary.',
              'rows':rows,'bases':bases,'assemblies':assemblies,'divisions':sorted(examples),
              'example_record_names':list(examples.values()),
              'inventory_sha256':hashlib.sha256(args.inventory.read_bytes()).hexdigest(),
              'inventory_at':datetime.fromtimestamp(args.inventory.stat().st_mtime,timezone.utc).isoformat()}
    c.execute("INSERT OR REPLACE INTO metadata VALUES ('manifest',?)",(json.dumps(manifest),)); c.commit()
    if c.execute('PRAGMA quick_check').fetchone()[0]!='ok': raise ValueError('SQLite check failed')
    c.execute('PRAGMA journal_mode=DELETE'); c.close(); work.replace(args.output)
    report={k:v for k,v in manifest.items() if k!='assemblies'}
    report.update(assembly_count=len(assemblies),index_bytes=args.output.stat().st_size)
    args.output.with_suffix('.report.json').write_text(json.dumps(report,indent=2)+'\n')
    print(json.dumps(report),flush=True)


if __name__=='__main__': main()