"""Rehearse one expand/backfill/cutover/contract recipe in an owned SQLite database."""
from __future__ import annotations
import argparse,hashlib,json,sqlite3,tempfile,time
from pathlib import Path
class MigrationError(Exception):pass
class InterruptedBackfill(Exception):pass

def connect(path,timeout=.15):
 db=sqlite3.connect(path,timeout=timeout);db.row_factory=sqlite3.Row;db.execute('PRAGMA foreign_keys=ON');db.execute('PRAGMA journal_mode=WAL');return db

def create_fixture(path,rows=5000,invalid=None):
 db=connect(path)
 db.executescript('CREATE TABLE assets(id INTEGER PRIMARY KEY,tag TEXT,description TEXT NOT NULL); CREATE TABLE rollout(key TEXT PRIMARY KEY,value TEXT NOT NULL); INSERT INTO rollout VALUES (\'stage\',\'legacy\'); INSERT INTO rollout VALUES (\'old_app_active\',\'yes\'); INSERT INTO rollout VALUES (\'checkpoint\',\'0\');')
 db.executemany('INSERT INTO assets VALUES (?,?,?)',((i,f'  AHU-{i:06d}  ' if i%3==0 else f'ahu-{i:06d}',f'Synthetic air handler {i}') for i in range(1,rows+1)))
 if invalid=='duplicate':db.execute('UPDATE assets SET tag=(SELECT tag FROM assets WHERE id=1) WHERE id=2')
 if invalid=='null':db.execute('UPDATE assets SET tag=NULL WHERE id=2')
 if invalid=='empty':db.execute("UPDATE assets SET tag='   ' WHERE id=2")
 db.commit();db.close()

def metadata(db,key):return db.execute('SELECT value FROM rollout WHERE key=?',(key,)).fetchone()[0]
def set_meta(db,key,value):db.execute('UPDATE rollout SET value=? WHERE key=?',(str(value),key))
def schema(db):return [dict(x) for x in db.execute("SELECT type,name,sql FROM sqlite_master WHERE name NOT LIKE 'sqlite_%' ORDER BY type,name")]
def content_digest(db):
 rows=db.execute('SELECT id,description FROM assets ORDER BY id').fetchall()
 return hashlib.sha256(json.dumps([list(r) for r in rows],separators=(',',':')).encode()).hexdigest()
def profile(db):
 columns=[r['name'] for r in db.execute('PRAGMA table_info(assets)')]
 key='asset_key' if 'asset_key' in columns else 'lower(trim(tag))'
 return {'rows':db.execute('SELECT COUNT(*) FROM assets').fetchone()[0],
         'null_or_empty_keys':db.execute(f"SELECT COUNT(*) FROM assets WHERE {key} IS NULL OR {key}='' ").fetchone()[0],
         'duplicate_keys':[dict(r) for r in db.execute(f'SELECT {key} AS value,COUNT(*) AS count FROM assets GROUP BY {key} HAVING COUNT(*)>1 LIMIT 20')],
         'checkpoint':int(metadata(db,'checkpoint')),'stage':metadata(db,'stage')}

def atomic(db,fn):
 db.execute('BEGIN IMMEDIATE')
 try:value=fn();db.commit();return value
 except BaseException:db.rollback();raise

def expand(db):
 if metadata(db,'stage')!='legacy':raise MigrationError('Expansion requires the legacy stage')
 def change():
  db.execute('ALTER TABLE assets ADD COLUMN asset_key TEXT')
  db.execute("CREATE TRIGGER assets_legacy_insert AFTER INSERT ON assets BEGIN UPDATE assets SET asset_key=lower(trim(NEW.tag)) WHERE id=NEW.id; END")
  db.execute("CREATE TRIGGER assets_legacy_update AFTER UPDATE OF tag ON assets BEGIN UPDATE assets SET asset_key=lower(trim(NEW.tag)) WHERE id=NEW.id; END")
  set_meta(db,'stage','expanded')
 atomic(db,change)

def contracts(db,old_required=True):
 stage=metadata(db,'stage');columns={r['name'] for r in db.execute('PRAGMA table_info(assets)')}
 result={'stage':stage,'old_required':old_required}
 if old_required:
  if 'tag' not in columns:raise MigrationError('Old query contract no longer exists')
  old=db.execute('SELECT id,lower(trim(tag)) AS key,description FROM assets ORDER BY id').fetchall()
  result['old_rows']=len(old)
 if 'asset_key' in columns:
  expr='asset_key' if stage=='contracted' else 'coalesce(asset_key,lower(trim(tag)))'
  new=db.execute(f'SELECT id,{expr} AS key,description FROM assets ORDER BY id').fetchall()
  if old_required and [tuple(x) for x in new]!=[tuple(x) for x in old]:raise MigrationError('Old and new reads disagree')
  result['new_rows']=len(new);result['consistent']=True
 elif not old_required:raise MigrationError('New query contract is not available')
 return result

def backfill(db,batch_size=250,interrupt_after=None):
 if metadata(db,'stage') not in ('expanded','backfilled'):raise MigrationError('Backfill requires expansion')
 if batch_size<1:raise ValueError('batch_size must be positive')
 batches=0;updated=0
 while True:
  def batch():
   # Scan missing keys, rather than trusting an ID checkpoint for old-application writes.
   ids=[r[0] for r in db.execute('SELECT id FROM assets WHERE asset_key IS NULL ORDER BY id LIMIT ?',(batch_size,))]
   if not ids:return 0
   placeholders=','.join('?' for _ in ids)
   bad=db.execute(f"SELECT id FROM assets WHERE id IN ({placeholders}) AND (tag IS NULL OR trim(tag)='')",ids).fetchall()
   if bad:raise MigrationError('Null or empty legacy tag at IDs '+','.join(str(x[0]) for x in bad))
   db.execute(f'UPDATE assets SET asset_key=lower(trim(tag)) WHERE id IN ({placeholders})',ids)
   set_meta(db,'checkpoint',max(int(metadata(db,'checkpoint')),max(ids)))
   return len(ids)
  n=atomic(db,batch)
  if not n:break
  updated+=n;batches+=1
  if interrupt_after and batches==interrupt_after:raise InterruptedBackfill(f'Interrupted after {batches} committed batches')
 def validate():
  bad=db.execute("SELECT COUNT(*) FROM assets WHERE asset_key IS NULL OR asset_key='' OR asset_key<>lower(trim(tag))").fetchone()[0]
  if bad:raise MigrationError('Incomplete or inconsistent backfill')
  db.execute('CREATE UNIQUE INDEX IF NOT EXISTS assets_key_unique ON assets(asset_key)')
  # Verify the actual index definition, not just its name.
  ix=db.execute("SELECT sql FROM sqlite_master WHERE type='index' AND name='assets_key_unique'").fetchone()
  cols=[r['name'] for r in db.execute("PRAGMA index_info('assets_key_unique')")]
  props=[r for r in db.execute("PRAGMA index_list('assets')") if r['name']=='assets_key_unique']
  if not ix or cols!=['asset_key'] or not props or props[0]['unique']!=1:raise MigrationError('Unexpected assets_key_unique index definition')
  set_meta(db,'stage','backfilled')
 atomic(db,validate)
 return {'updated_rows':updated,'committed_batches':batches,'checkpoint':int(metadata(db,'checkpoint'))}

def cutover(db):
 if metadata(db,'stage')!='backfilled':raise MigrationError('Cutover requires a validated backfill')
 def change():set_meta(db,'old_app_active','no');set_meta(db,'stage','cutover')
 atomic(db,change)

def contract(db):
 if metadata(db,'stage')!='cutover' or metadata(db,'old_app_active')!='no':raise MigrationError('Retire old application before contraction')
 before=content_digest(db)
 def change():
  db.execute('DROP TRIGGER assets_legacy_insert');db.execute('DROP TRIGGER assets_legacy_update')
  db.execute('CREATE TABLE assets_new(id INTEGER PRIMARY KEY,asset_key TEXT NOT NULL UNIQUE CHECK(length(asset_key)>0),description TEXT NOT NULL)')
  db.execute('INSERT INTO assets_new SELECT id,asset_key,description FROM assets')
  db.execute('DROP TABLE assets');db.execute('ALTER TABLE assets_new RENAME TO assets');set_meta(db,'stage','contracted')
  if content_digest(db)!=before:raise MigrationError('Non-key content changed during contraction')
 atomic(db,change)

def lock_probe(path):
 holder=connect(path);contender=connect(path,timeout=.075);holder.execute('BEGIN IMMEDIATE');start=time.perf_counter()
 try:
  try:contender.execute("UPDATE rollout SET value=value WHERE key='stage'");contender.commit();raise MigrationError('Competing write unexpectedly succeeded')
  except sqlite3.OperationalError as e:
   if 'locked' not in str(e).lower():raise
   return {'observed':'database is locked','wait_seconds':round(time.perf_counter()-start,6),'configured_timeout_seconds':.075}
 finally:holder.rollback();holder.close();contender.close()

def rehearse(out=Path('results'),rows=5000,batch_size=250,interrupt_after=3,invalid=None):
 out=Path(out);out.mkdir(parents=True,exist_ok=True)
 report={'engine':'SQLite','sqlite_version':sqlite3.sqlite_version,'passed':False,'fixture_rows':rows,'invalid_fixture':invalid,'stages':[],'limits':['SQLite lock behavior only; not PostgreSQL lock or concurrent-index behavior.','Synthetic fixture timing is not a production rollout estimate.']}
 start=time.perf_counter()
 with tempfile.TemporaryDirectory(prefix='migration-rehearsal-') as temp:
  path=Path(temp)/'rehearsal.sqlite';db=None
  def record(name,fn):
   t=time.perf_counter();v=fn();report['stages'].append({'stage':name,'seconds':round(time.perf_counter()-t,6),'result':v});return v
  try:
   record('fixture',lambda:create_fixture(path,rows,invalid));db=connect(path);before=content_digest(db);report['schema_before']=schema(db);report['profile_before']=profile(db)
   record('expand',lambda:expand(db));record('overlap_contracts',lambda:contracts(db))
   if interrupt_after:
    try:backfill(db,batch_size,interrupt_after)
    except InterruptedBackfill as e:
     report['interruption']={'message':str(e),'committed_rows':db.execute('SELECT COUNT(*) FROM assets WHERE asset_key IS NOT NULL').fetchone()[0],'checkpoint':int(metadata(db,'checkpoint'))}
     db.close();db=connect(path) # Resume through a fresh connection from committed on-disk state.
   record('resume_backfill',lambda:backfill(db,batch_size));record('backfilled_contracts',lambda:contracts(db));record('lock_probe',lambda:lock_probe(path))
   record('cutover',lambda:cutover(db));record('contract',lambda:contract(db));record('new_contract',lambda:contracts(db,False))
   report['profile_after']=profile(db);report['schema_after']=schema(db);report['content_preserved']=content_digest(db)==before
   if not report['content_preserved']:raise MigrationError('ID/description checksum changed')
   report['passed']=True
  except Exception as e:
   report['error']=f'{type(e).__name__}: {e}'
   if db:
    report['failure_profile']=profile(db);report['schema_at_failure']=schema(db)
  finally:
   if db:db.close()
  report['temporary_directory_removed_after_run']=True
 report['total_seconds']=round(time.perf_counter()-start,6)
 (out/'run.json').write_text(json.dumps(report,indent=2)+'\n')
 return report

def main():
 p=argparse.ArgumentParser(description=__doc__);p.add_argument('--out',type=Path,default=Path('results'));p.add_argument('--rows',type=int,default=5000);p.add_argument('--batch-size',type=int,default=250);p.add_argument('--interrupt-after',type=int,default=3);p.add_argument('--invalid',choices=['duplicate','null','empty']);a=p.parse_args()
 if not 2<=a.rows<=1000000:p.error('rows must be 2..1000000')
 if a.batch_size<1:p.error('batch-size must be positive')
 r=rehearse(a.out,a.rows,a.batch_size,a.interrupt_after,a.invalid);print(('PASS' if r['passed'] else 'FAIL')+': '+str(a.out/'run.json'))
 if 'error' in r:print(r['error'])
 return 0 if r['passed'] else 1
if __name__=='__main__':raise SystemExit(main())