"""Random access to 32-book continuous timelines; model consumes first twelve."""
from collections import OrderedDict
import hashlib,io,json,sqlite3
import numpy as np
from common import *

class Data:
 def __init__(self,pilot=False,path=None):
  if pilot:raise RuntimeError('No partial-corpus production mode')
  self.root=RUN/'large';self.ready=json.loads((self.root/'DATA_READY.json').read_text())
  self.db=sqlite3.connect('file:'+str(path or self.root/'mix.sqlite')+'?mode=ro&immutable=1',uri=True)
  self.db.row_factory=sqlite3.Row;self.files=OrderedDict();self.cache=OrderedDict()
  self.count=self.db.execute('SELECT count(*) FROM presentations').fetchone()[0]
 def episode(self,uid):
  if uid not in self.cache:
   r=self.db.execute('SELECT * FROM episodes WHERE uid=?',(uid,)).fetchone()
   if not r:raise KeyError(uid)
   if r['tar'] not in self.files:
    self.files[r['tar']]=open(r['tar'],'rb')
    if len(self.files)>16:self.files.popitem(last=False)[1].close()
   f=self.files[r['tar']];f.seek(r['code_offset']);blob=f.read(r['code_size'])
   if hashlib.sha256(blob).hexdigest()!=r['code_sha']:raise RuntimeError('Continuous codec timeline hash changed')
   codes=np.load(io.BytesIO(blob),allow_pickle=False)
   if codes.ndim!=2 or codes.shape[1]!=32 or codes.dtype.kind not in 'iu' or codes.min()<0 or codes.max()>1023:raise ValueError('Invalid podcast codec array')
   self.cache[uid]=codes
   if len(self.cache)>8:self.cache.popitem(last=False)
  self.cache.move_to_end(uid);return self.cache[uid]
 def row(self,key,condition='reference'):
  r=self.db.execute('SELECT * FROM samples WHERE key=?',(int(key),)).fetchone()
  if not r:raise KeyError(key)
  ex=json.loads(r['example_json']);codes=self.episode(r['episode'])
  a,b=ex['target_frame_start'],ex['target_frame_end'];ca,cb=ex['context_frame_start'],ex['context_frame_end']
  if not 0<=ca<=cb<=a<b<=len(codes) or cb-ca>750:raise ValueError('Future context or invalid boundaries')
  refs={}
  for speaker,q in ex['speaker_references'].items():
   ra,rb=q['start_frame'],q['end_frame']
   if not 0<=ra<rb<=a or not 63<=rb-ra<=187:raise ValueError('Reference must be past-only and 5–15s')
   refs[speaker]=codes[ra:rb,:12].copy()
  if set(refs)!={'S1','S2'}:raise ValueError('Two past identity references required')
  description='\n'.join(f"{q['speaker_id']}: {max(0,q['start_s']-ca/12.5):.2f}–{q['end_s']-ca/12.5:.2f} seconds in previous audio" for q in ex['context_speaker_intervals'])
  return dict(uid=r['uid'],episode=r['episode'],source=r['source'],language=r['language'],plain=r['plain'],script=r['inline'],
   instruction='GENERAL: '+r['general']+'\nSCRIPT:\n'+r['inline'],frames=r['frames'],target_codes=codes[a:b,:12].copy(),
   speaker_references=refs,reference_codes=None,context_codes=codes[ca:cb,:12].copy(),context_description=description,
   conditioning='two_speaker_references_and_past_context',target_turns=ex['target_turns'])
 def presentation(self,pos):
  r=self.db.execute('SELECT * FROM presentations WHERE pos=?',(int(pos),)).fetchone()
  if not r:raise KeyError(pos)
  key=r['sample_key']
  if key is None:key=self.ready['canary_keys'][0]
  return self.row(key),r['sample_key'] is not None
 def valid_count(self,start,width):
  return self.db.execute('SELECT count(*) FROM presentations WHERE pos>=? AND pos<? AND sample_key IS NOT NULL',(start,start+width)).fetchone()[0]
 def canaries(self):return self.ready['canary_keys']
 def validation_keys(self):return self.ready['validation_keys']
 def close(self):
  for f in self.files.values():f.close()
  self.db.close()