#!/usr/bin/env python3
# -*- coding: utf-8 -*-
"""Teste FINAL do item (b) — modelo KHTST TREINADO (checkpoint fase-som).

Verifica, com o modelo pós-treino real:
  1. 4 requisições submetidas em série ANTES da 1ª resposta;
  2. recebimento (estado RECEBIDA) confirmado para todas antes do processamento;
  3. processamento (PROCESSANDO → CONCLUIDA) com tokens gerados;
  4. stopping (cancelamento) honrado em ≤ 1 passo durante geração ativa;
  5. continue (retomada de pedido pausado).
Saída: telemetria_out/teste_item_b_v10.json
"""
import json
import sys
import threading
import time

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

import torch

from khtst.config import Config
from khtst.dados.tokenizador import TokenizadorKHTST
from khtst.memoria.orquestrador import OrquestradorSOM
from khtst.nucleo.modelo import KHTSTModel
from khtst.servico import SessaoKHTST

CFG_PATH = "/home/z/my-project/khtst/configs/base.json"
TK_PATH = "/home/z/my-project/khtst/cache_dados/tokenizador.json"
SAIDA = "/home/z/my-project/khtst/telemetria_out/teste_item_b_v10.json"

cfg = Config().de_arquivo(CFG_PATH)
torch.manual_seed(cfg.dados.semente)
modelo = KHTSTModel(cfg, usar_multimodal=True)
tk = TokenizadorKHTST(TK_PATH)

# retoma o estado final (fase-som)
from khtst.dados.checkpoints import GestorCheckpoints
ck = GestorCheckpoints("/home/z/my-project/khtst/estados_local", "estados/fase-v10")
tag = ck.ultima_tag()
if tag:
    from khtst.treino.estado_integral import EstadoIntegral
    from khtst.memoria.orquestrador import OrquestradorSOM
    # v9 — estado integral: roteador S-SOM + orquestrador restaurados
    orq = OrquestradorSOM(cfg.modelo.d_modelo, cfg.som)
    ck.estado = EstadoIntegral(modelo, {"roteador": getattr(modelo, "roteador_ssom", None),
                                        "orquestrador": orq})
    ok = ck.carregar(modelo, tag)
    print(f"estado retomado: {tag} (ok={ok}, estado integral)")
modelo.eval()

sessao = SessaoKHTST(modelo, decodificar=tk)
prompts = [
    "O KHTST é um modelo multimodal que",
    "A previsão do tempo para amanhã indica",
    "Para cozinhar arroz, primeiro",
    "O melhor time do Brasil é",
]
# ---- 1) 4 requisições EM SÉRIE antes da 1ª resposta ----
pids = [sessao.submeter(p, max_novos=10, tarefa="lm") for p in prompts]
estados_iniciais = [sessao.pedidos[i].estado.value for i in pids]
check_recebimento = all(e == "recebida" for e in estados_iniciais)
print("1) 4 pedidos ANTES da 1ª resposta:", pids, "| estados:", estados_iniciais)

# ---- 2/3) processa todos (a bomba adhoc roda ENTRE tokens) ----
t0 = time.time()
saidas = sessao.executar_todos(max_passos=400)
dur = time.time() - t0
rel = sessao.relatorio()
check_processamento = all(rel[i]["estado"] == "concluida" and rel[i]["n_tokens"] == 10
                          for i in pids)
print("2) processamento:", [(i, rel[i]["estado"], rel[i]["n_tokens"]) for i in pids],
      f"| {dur:.1f}s")
for i in pids:
    texto = tk.decodificar(saidas[i])
    print(f"   pedido {i}: «{prompts[i]}» → {texto[:70]!r}")

# ---- 4) stopping durante processamento ativo (thread externa) ----
sessao2 = SessaoKHTST(modelo, decodificar=tk)
p_st = sessao2.submeter("Escreva uma história longa sobre o mar e as estrelas",
                        max_novos=60, tarefa="lm")
def para_no_meio():
    time.sleep(0.4)
    sessao2.stop(p_st, cancelar=True)
threading.Thread(target=para_no_meio, daemon=True).start()
sessao2.executar_todos(max_passos=400)
rel2 = sessao2.relatorio()[p_st]
check_stop = rel2["estado"] == "cancelada" and 0 < rel2["n_tokens"] < 60
print("4) stopping:", rel2["estado"], "| tokens antes do stop:", rel2["n_tokens"],
      "| eventos:", rel2["eventos"])

# ---- 5) pause → continue ----
sessao3 = SessaoKHTST(modelo, decodificar=tk)
p_pc = sessao3.submeter("O livro da floresta narra", max_novos=12, tarefa="lm")
sessao3.canal.stop(p_pc, cancelar=False)          # pausa no 1º dreno
def retoma():
    time.sleep(0.3)
    sessao3.continuar(p_pc)
threading.Thread(target=retoma, daemon=True).start()
sessao3.executar_todos(max_passos=400)
rel3 = sessao3.relatorio()[p_pc]
check_continue = rel3["estado"] == "concluida" and "pausada" in rel3["eventos"] \
    and "retomada" in rel3["eventos"]
print("5) pause→continue:", rel3["estado"], "| eventos:", rel3["eventos"])

resultado = {
    "modelo": "KHTST v10 (checkpoint: " + str(tag) + ")",
    "item_b_4_requisicoes_antes_da_1a": bool(check_recebimento),
    "item_b_processamento_e_conclusao": bool(check_processamento),
    "item_b_stopping_durante_ativa": bool(check_stop),
    "item_b_pause_continue": bool(check_continue),
    "estados_iniciais": estados_iniciais,
    "relatorio_completo": {str(i): rel[i] for i in rel},
    "relatorio_stop": rel2,
    "relatorio_pause_continue": rel3,
    "duracao_s": round(dur, 2),
}
ok_all = all([check_recebimento, check_processamento, check_stop, check_continue])
resultado["item_b_aprovado"] = bool(ok_all)
with open(SAIDA, "w", encoding="utf-8") as f:
    json.dump(resultado, f, ensure_ascii=False, indent=2, default=str)
print(f"\nITEM B: {'APROVADO' if ok_all else 'REPROVADO'} → {SAIDA}")
sys.exit(0 if ok_all else 1)