#!/usr/bin/env python3
"""Chat interativo usando exclusivamente pesos treinados do zero."""

from __future__ import annotations

import argparse
from pathlib import Path

import torch

from usar_gpt_zero import carregar
from treinar_gpt_zero import gerar


SYSTEM_DEFAULT = "Você é um assistente local útil e educado."


def main() -> None:
    parser = argparse.ArgumentParser(description="Conversa com o GPT treinado do zero.")
    parser.add_argument("--checkpoint", type=Path, default=Path("modelo_dialogos_multi_v2.pt"))
    parser.add_argument("--system", default=SYSTEM_DEFAULT)
    parser.add_argument("--tokens", type=int, default=80)
    parser.add_argument("--temperature", type=float, default=0.05)
    parser.add_argument("--top-k", type=int, default=20)
    parser.add_argument("--top-p", type=float, default=0.9)
    parser.add_argument("--repetition-penalty", type=float, default=1.15)
    parser.add_argument("--device", choices=["auto", "cpu", "cuda"], default="auto")
    args = parser.parse_args()

    if args.device == "auto":
        device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
    else:
        device = torch.device(args.device)
    modelo, tokenizer, checkpoint, bos_id, eos_id = carregar(args.checkpoint, device)
    print("Modelo treinado do zero pronto. Digite /sair para terminar.")
    historico: list[tuple[str, str]] = []
    while True:
        pergunta = input("Você: ").strip()
        if pergunta == "/sair":
            break
        historico.append(("user", pergunta))
        partes = [f"<|system|>{args.system}<|end|>"]
        for indice, (papel, conteudo) in enumerate(historico[-4:]):
            partes.append(f"<|{papel}|>{conteudo}<|end|>")
            if indice == len(historico[-4:]) - 1 and papel == "user":
                partes.append("<|assistant|>")
        prompt = "".join(partes)
        resposta = gerar(modelo, tokenizer, prompt, device, args.tokens, args.temperature, args.top_k, args.top_p, bos_id, eos_id, args.repetition_penalty)
        historico.append(("assistant", resposta))
        print(f"Modelo: {resposta}\n")


if __name__ == "__main__":
    main()