#!/usr/bin/env python3
"""Coleta retomável dos logs públicos. Nenhum dado da tabela de terceiros é lido."""
import argparse, collections, concurrent.futures, datetime as dt, hashlib, io, json, pathlib, sqlite3, subprocess, threading, time, zipfile, zlib
BASE='https://resultados.tse.jus.br/oficial/ele2026/arquivo-urna/3221'
UFS='ac al ap am ba ce df es go ma mt ms mg pa pb pr pe pi rj rn rs ro rr sc sp se to'.split()
ROOT=pathlib.Path(__file__).resolve().parent
DATE='25/10/2026'
lock=threading.Lock(); next_request=0; RPS=10

def get(url):
 global next_request
 for attempt in range(4):
  with lock:
   now=time.monotonic();delay=max(0,next_request-now);next_request=max(now,next_request)+1/RPS
  if delay:time.sleep(delay)
  r=subprocess.run(['curl','--fail','--location','--silent','--show-error','--max-time','45',url],capture_output=True)
  if r.returncode==0:return r.stdout
  if attempt<3:time.sleep(2**attempt)
 raise RuntimeError(r.stderr.decode(errors='replace')[:250])

def parse_log(blob):
 events=set();openings=set();closings=set();devices=set();backwards=False
 with zipfile.ZipFile(io.BytesIO(blob)) as z:
  files=[n for n in z.namelist() if pathlib.PurePosixPath(n).name.lower()=='logd.dat']
  if not files:raise ValueError('logd.dat ausente')
  for name in files:
   previous={}
   for line in z.read(name).decode('latin1').splitlines():
    cols=line.split('\t')
    if len(cols)<5 or not cols[0].startswith(DATE+' '):continue
    stamp=cols[0][11:19]
    try:h,m,s=map(int,stamp.split(':')); sec=h*3600+m*60+s
    except ValueError:continue
    dev=cols[2];event=cols[4].strip();devices.add(dev)
    if dev in previous and sec<previous[dev]:backwards=True
    previous[dev]=sec
    if event=='O voto do eleitor foi computado':events.add(line)
    elif event=='Urna pronta para receber votos':openings.add((dev,sec))
    elif event=='Procedimento de encerramento confirmado':closings.add((dev,sec))
 return dict(events=events,openings=openings,closings=closings,devices=devices,clock_backwards=backwards)

def download_section(row):
 key,uf,mun,zone,section,name=row
 url=f'{BASE}/dados/{uf}/{mun}/{zone}/{section}'
 auxurl=f'{url}/p003221-{uf}-m{mun}-z{zone}-s{section}-aux.json'
 raw=get(auxurl);aux=json.loads(raw);urls=set()
 for h in aux.get('hashes',[]):
  for a in h.get('arq',[]):
   if a.get('tp')=='log' or a.get('nm','').endswith('-log.jez'):
    urls.add(f"{url}/{h['hash']}/{a['nm']}")
 if not urls:raise ValueError('Auxiliar sem arquivo de log')
 all_events=set();opens=set();closes=set();devices=set();backwards=False;files=[]
 for logurl in sorted(urls):
  blob=get(logurl);v=parse_log(blob)
  all_events.update(v['events']);opens.update(v['openings']);closes.update(v['closings']);devices.update(v['devices']);backwards|=v['clock_backwards']
  files.append(dict(url=logurl,sha256=hashlib.sha256(blob).hexdigest(),bytes=len(blob)))
 hist=collections.Counter()
 for line in all_events:
  h,m,s=map(int,line[11:19].split(':'));hist[str(h*3600+m*60+s)]+=1
 return dict(key=key,aux_url=auxurl,aux_sha256=hashlib.sha256(raw).hexdigest(),aux=aux,files=files,histogram=hist,openings=sorted(opens),closings=sorted(closes),devices=sorted(devices),clock_backwards=backwards,events=len(all_events),downloaded_at=dt.datetime.now(dt.timezone.utc).isoformat())

def connect():
 db=sqlite3.connect(ROOT/'coleta.sqlite');db.execute('PRAGMA journal_mode=WAL')
 db.execute('CREATE TABLE IF NOT EXISTS sections (id TEXT PRIMARY KEY, uf TEXT, mun TEXT, zone TEXT, section TEXT, name TEXT, status TEXT DEFAULT "pending", payload BLOB, error TEXT)');return db

def seed(db,ufs):
 (ROOT/'indices').mkdir(exist_ok=True)
 for uf in ufs:
  path=ROOT/'indices'/f'{uf}.json'
  if not path.exists():
   blob=get(f'{BASE}/config/{uf}/{uf}-p003221-cs.json');json.loads(blob);path.write_bytes(blob)
  j=json.loads(path.read_bytes())
  for abr in j['abr']:
   for m in abr['mu']:
    for z in m['zon']:
     for s in z['sec']:
      row=(f"{uf}/{m['cd']}/{z['cd']}/{s['ns']}",uf,m['cd'],z['cd'],s['ns'],m['nm'])
      db.execute('INSERT OR IGNORE INTO sections(id,uf,mun,zone,section,name) VALUES(?,?,?,?,?,?)',row)
  db.commit()

def status(db):
 counts=dict(db.execute('SELECT status,count(*) FROM sections GROUP BY status'))
 out=dict(updated_at=dt.datetime.now(dt.timezone.utc).isoformat(),sections=counts,total=sum(counts.values()),source='TSE — coleta direta',requests_per_second=RPS,complete=bool(counts) and counts.get('ok',0)==sum(counts.values()))
 tmp=ROOT/'status.tmp';tmp.write_text(json.dumps(out,ensure_ascii=False,indent=2));tmp.replace(ROOT/'status.json');return out

def main():
 global RPS
 ap=argparse.ArgumentParser();ap.add_argument('--ufs',nargs='+',default=UFS,choices=UFS);ap.add_argument('--limit',type=int);ap.add_argument('--rps',type=float,default=10);ap.add_argument('--workers',type=int,default=8);ap.add_argument('--retry-errors',action='store_true');args=ap.parse_args()
 RPS=max(.1,min(args.rps,10));db=connect();seed(db,args.ufs)
 if args.retry_errors:db.execute('UPDATE sections SET status="pending" WHERE status="error"');db.commit()
 placeholders=','.join('?' for _ in args.ufs)
 rows=db.execute(f'SELECT id,uf,mun,zone,section,name FROM sections WHERE status="pending" AND uf IN ({placeholders}) ORDER BY uf,mun,zone,section',args.ufs).fetchall()
 if args.limit is not None:rows=rows[:args.limit]
 print(json.dumps(status(db)),flush=True);iterator=iter(rows);done=0
 with concurrent.futures.ThreadPoolExecutor(max_workers=args.workers) as pool:
  pending={}
  def submit():
   try:row=next(iterator)
   except StopIteration:return
   pending[pool.submit(download_section,row)]=row
  for _ in range(args.workers):submit()
  while pending:
   ready,_=concurrent.futures.wait(pending,timeout=30,return_when=concurrent.futures.FIRST_COMPLETED)
   for f in ready:
    row=pending.pop(f)
    try:
     result=f.result();payload=zlib.compress(json.dumps(result,ensure_ascii=False).encode());db.execute('UPDATE sections SET status="ok",payload=?,error=NULL WHERE id=?',(payload,row[0]))
    except Exception as e:db.execute("UPDATE sections SET status='error',error=? WHERE id=?",(str(e),row[0]))
    done+=1;submit()
   db.commit()
   if done%100<args.workers or not ready:print(json.dumps(status(db)),flush=True)
 print(json.dumps(status(db)),flush=True)
if __name__=='__main__':main()