#!/usr/bin/env python3
# -*- coding: utf-8 -*-
"""Coleta UMA fonte (subprocesso isolado — estado limpo da biblioteca datasets,
RAM devolvida ao SO ao terminar). Recebe o plano em JSON, escreve registros
convertidos (multimodal em base64) num JSONL parcial."""
import argparse
import base64
import json
import sys

sys.path.insert(0, "/home/z/my-project/aurora/src")

from aurora.dados.streaming import StreamingCorpus, via_especial


def _serializa(r: dict) -> dict:
    out = {}
    for k, v in r.items():
        if v is None:
            continue
        if isinstance(v, bytes):
            out[k] = base64.b64encode(v).decode("ascii")
        else:
            out[k] = v
    return out


def main():
    ap = argparse.ArgumentParser()
    ap.add_argument("--indice", required=True, type=int,
                    help="índice da fonte em StreamingCorpus.PLANO")
    ap.add_argument("--saida", required=True)
    ap.add_argument("--estado", required=True, help="JSON de evento (status/n)")
    args = ap.parse_args()

    corpus = StreamingCorpus()
    plano = corpus.PLANO[args.indice]
    registros = []
    status, detalhe = "OK", plano.get("via", "")
    try:
        for ex in corpus._iterar_fonte(plano):
            if len(registros) >= plano.get("max", 0):
                break
            reg = plano["ad"](ex) if plano.get("ad") else \
                (ex if via_especial(plano["via"]) else None)
            if reg is None:
                continue
            reg["fonte"] = plano["id"]
            reg.setdefault("meta", {})
            registros.append(reg)
    except Exception as e:
        status = "ERRO"
        detalhe = f"{type(e).__name__}: {str(e)[:180]}"
    if corpus.eventos:                      # erro/PULADO interno da fonte
        ev = corpus.eventos[0]
        status, detalhe = ev.status, ev.detalhe
    with open(args.saida, "w", encoding="utf-8") as f:
        for r in registros:
            f.write(json.dumps(_serializa(r), ensure_ascii=False) + "\n")
    with open(args.estado, "w", encoding="utf-8") as f:
        json.dump({"fonte": plano["id"], "status": status,
                   "detalhe": detalhe, "n": len(registros)}, f,
                  ensure_ascii=False)
    print(f"{status}: {plano['id']} n={len(registros)}")


if __name__ == "__main__":
    main()