import argparse
import json
import sys
from pathlib import Path
import re

import torch
import torch.nn.functional as F
from transformers import AutoTokenizer

try:
    from tokenizers import Tokenizer as HFTokenizer
except ImportError:
    HFTokenizer = None

DEFAULT_CHECKPOINT = r"D:\Downloads\hssm_fineweb_edu_final.pt"
DEFAULT_TOKENIZER = r"D:\Downloads\simple_tokenizer_20k.json"
RUBINET_HSSM_PATH = r"C:\Users\ASUS\.anaconda"

sys.path.append(RUBINET_HSSM_PATH)
from RubiNet_HSSM import HierarchicalSSM
from hssm_v2_gpu_pretrain import HSSMV2Config, HSSMV2LM

if hasattr(sys.stdout, "reconfigure"):
    sys.stdout.reconfigure(encoding="utf-8", errors="replace")
if hasattr(sys.stderr, "reconfigure"):
    sys.stderr.reconfigure(encoding="utf-8", errors="replace")


class CompatibleTokenizer:
    def __init__(self, tokenizer_path: str):
        path = Path(tokenizer_path)
        with path.open("r", encoding="utf-8") as f:
            data = json.load(f)

        self.backend = None
        if HFTokenizer is not None:
            try:
                self.backend = HFTokenizer.from_file(str(path))
            except Exception:
                self.backend = None

        if "model" in data and isinstance(data["model"], dict) and "vocab" in data["model"]:
            vocab = data["model"]["vocab"]
        elif "vocab" in data:
            vocab = data["vocab"]
        else:
            raise ValueError(f"Unsupported tokenizer format: {path}")

        self.vocab = {str(token): int(idx) for token, idx in vocab.items()}
        self.id_to_token = {idx: token for token, idx in self.vocab.items()}
        self.vocab_size = len(self.vocab)
        self.pad_token_id = self._resolve_token_id(["<PAD>", "[PAD]"], fallback=0)
        self.unk_token_id = self._resolve_token_id(["<UNK>", "[UNK]"], fallback=3)

        print(f"[TOKENIZER] Loaded vocab tokenizer - Vocab: {self.vocab_size:,}")

    def _resolve_token_id(self, candidates, fallback: int):
        for token in candidates:
            token_id = self.vocab.get(token)
            if token_id is not None:
                return token_id
        return fallback

    def encode(self, text, max_length=128):
        if self.backend is not None:
            ids = self.backend.encode(text).ids[:max_length]
        else:
            words = text.split()
            ids = [self.vocab.get(word, self.unk_token_id) for word in words][:max_length]
        if len(ids) < max_length:
            ids += [self.pad_token_id] * (max_length - len(ids))
        return ids

    def decode(self, ids):
        filtered = [int(i) for i in ids if int(i) != self.pad_token_id]
        if self.backend is not None:
            return self.backend.decode(filtered, skip_special_tokens=False)
        return " ".join(self.id_to_token.get(i, "<UNK>") for i in filtered)


class HFTokenizerAdapter:
    def __init__(self, tokenizer_name: str):
        self.backend = AutoTokenizer.from_pretrained(tokenizer_name, use_fast=True)
        if self.backend.pad_token is None:
            self.backend.pad_token = self.backend.eos_token or self.backend.unk_token
        self.vocab = self.backend.get_vocab()
        self.id_to_token = {idx: token for token, idx in self.vocab.items()}
        self.vocab_size = int(self.backend.vocab_size)
        self.pad_token_id = int(self.backend.pad_token_id)
        self.unk_token_id = int(self.backend.unk_token_id if self.backend.unk_token_id is not None else self.pad_token_id)

    def encode(self, text, max_length=128):
        ids = self.backend.encode(text, add_special_tokens=False, truncation=True, max_length=max_length)
        if len(ids) < max_length:
            ids += [self.pad_token_id] * (max_length - len(ids))
        return ids

    def decode(self, ids):
        filtered = [int(i) for i in ids if int(i) != self.pad_token_id]
        return self.backend.decode(filtered, skip_special_tokens=True)


def build_model(tokenizer):
    return HierarchicalSSM(
        vocab_size=tokenizer.vocab_size,
        d_model=512,
        d_state=32,
        num_blocks=6,
        num_experts=8,
        top_k=2,
        chunk_size=4,
        expert_dim=1024,
    )


def build_hssm_v2_model(tokenizer, checkpoint_config: dict):
    config = HSSMV2Config(
        vocab_size=int(checkpoint_config.get("vocab_size", tokenizer.vocab_size)),
        d_model=int(checkpoint_config.get("d_model", 288)),
        n_layers=int(checkpoint_config.get("n_layers", 10)),
        d_ff=int(checkpoint_config.get("d_ff", 512)),
        state_rank=int(checkpoint_config.get("state_rank", 128)),
        chunk_size=int(checkpoint_config.get("chunk_size", 8)),
        dropout=float(checkpoint_config.get("dropout", 0.0)),
        max_seq_len=int(checkpoint_config.get("max_seq_len", 1024)),
        tie_embeddings=bool(checkpoint_config.get("tie_embeddings", True)),
        num_experts=int(checkpoint_config.get("num_experts", 64)),
        experts_per_token=int(checkpoint_config.get("experts_per_token", 1)),
        expert_dim=int(checkpoint_config.get("expert_dim", 2048)),
        moe_every=int(checkpoint_config.get("moe_every", 4)),
        aux_loss_coef=float(checkpoint_config.get("aux_loss_coef", 1e-2)),
    )
    return HSSMV2LM(config)


def _looks_like_hf_tokenizer_reference(tokenizer_path: str) -> bool:
    path = Path(tokenizer_path)
    return not path.exists()


def _load_tokenizer(tokenizer_path: str):
    if _looks_like_hf_tokenizer_reference(tokenizer_path):
        return HFTokenizerAdapter(tokenizer_path)
    return CompatibleTokenizer(tokenizer_path)


def _is_hssm_v2_checkpoint(checkpoint: dict) -> bool:
    config = checkpoint.get("config") if isinstance(checkpoint, dict) else None
    if not isinstance(config, dict):
        return False
    required_keys = {"d_model", "n_layers", "state_rank", "chunk_size"}
    return required_keys.issubset(config.keys())


def load_pretrained(checkpoint_path: str, tokenizer_path: str, device: str):
    checkpoint_file = Path(checkpoint_path)

    if not checkpoint_file.exists():
        raise FileNotFoundError(f"Checkpoint not found: {checkpoint_file}")

    if not _looks_like_hf_tokenizer_reference(tokenizer_path):
        tokenizer_file = Path(tokenizer_path)
        if not tokenizer_file.exists():
            raise FileNotFoundError(f"Tokenizer not found: {tokenizer_file}")

    tokenizer = _load_tokenizer(tokenizer_path)

    checkpoint = torch.load(str(checkpoint_file), map_location=device, weights_only=False)
    state_dict = checkpoint["model_state_dict"] if "model_state_dict" in checkpoint else checkpoint
    if _is_hssm_v2_checkpoint(checkpoint):
        model = build_hssm_v2_model(tokenizer, checkpoint.get("config", {}))
    else:
        model = build_model(tokenizer)
    missing, unexpected = model.load_state_dict(state_dict, strict=False)

    model = model.to(device)
    model.eval()

    print("Loaded HSSM checkpoint")
    print(f"  Path: {checkpoint_file}")
    print(f"  Missing keys: {len(missing)}")
    print(f"  Unexpected keys: {len(unexpected)}")
    if "epoch" in checkpoint:
        print(f"  Epoch: {checkpoint['epoch']}")
    if "loss" in checkpoint:
        print(f"  Loss: {checkpoint['loss']}")
    print(f"  Model type: {'HSSM v2' if _is_hssm_v2_checkpoint(checkpoint) else 'RubiNet HSSM'}")

    return tokenizer, model


def _model_chunk_size(model) -> int:
    if hasattr(model, "chunk_size"):
        return int(model.chunk_size)
    if hasattr(model, "config") and hasattr(model.config, "chunk_size"):
        return int(model.config.chunk_size)
    return 1


def _next_token_logits(model, input_tensor: torch.Tensor, current_len: int) -> torch.Tensor:
    outputs = model(input_tensor)
    if isinstance(outputs, dict):
        logits = outputs.get("logits")
        if logits is None:
            raise ValueError("Model returned a dict without logits")
        return logits[0, current_len - 1, :].clone()
    chunk_size = _model_chunk_size(model)
    chunk_idx = max((current_len - 1) // chunk_size, 0)
    return outputs[0, chunk_idx, :].clone()


def build_prompt(user_text: str, cot_mode: bool = False) -> str:
    cleaned_user_text = user_text.strip()
    if cot_mode:
        return (
            "system: Reply only in correct English. "
            "Follow English grammar, spelling, punctuation, and sentence structure strictly. "
            "Do not output fragments, corrupted tokens, mixed-language text, or placeholder symbols. "
            "Think step by step briefly and keep the output clean. "
            "Output exactly two lines in this format: "
            "Reasoning: <very short reasoning>. "
            "Answer: <final answer>. "
            "Keep both lines grammatical and concise.\n"
            f"user: {cleaned_user_text}\n"
            "assistant:"
        )
    return (
        "system: Reply only in correct English. "
        "Follow English grammar, spelling, punctuation, and sentence structure strictly. "
        "Use short complete sentences. "
        "Do not output broken words, malformed tokens, mixed-language text, or placeholder symbols.\n"
        f"user: {cleaned_user_text}\n"
        "assistant:"
    )


def safe_print(text: str):
    try:
        print(text)
    except UnicodeEncodeError:
        sanitized = text.encode("utf-8", errors="replace").decode("utf-8", errors="replace")
        print(sanitized)


def _apply_top_p_filter(logits: torch.Tensor, top_p: float) -> torch.Tensor:
    if top_p >= 1.0:
        return logits
    sorted_logits, sorted_indices = torch.sort(logits, descending=True)
    sorted_probs = F.softmax(sorted_logits, dim=-1)
    cumulative_probs = torch.cumsum(sorted_probs, dim=-1)
    sorted_indices_to_remove = cumulative_probs > top_p
    sorted_indices_to_remove[..., 1:] = sorted_indices_to_remove[..., :-1].clone()
    sorted_indices_to_remove[..., 0] = False
    indices_to_remove = sorted_indices[sorted_indices_to_remove]
    logits[indices_to_remove] = float("-inf")
    return logits


def _has_repeat_ngram(token_ids, next_token_id: int, ngram_size: int) -> bool:
    if ngram_size <= 1 or len(token_ids) < ngram_size - 1:
        return False
    candidate = token_ids[-(ngram_size - 1):] + [next_token_id]
    for i in range(len(token_ids) - ngram_size + 1):
        if token_ids[i:i + ngram_size] == candidate:
            return True
    return False


def _normalize_word(word: str) -> str:
    return re.sub(r"[^a-z0-9]+", "", word.lower())


def _recent_word_counts(text: str, window: int = 12):
    words = [_normalize_word(part) for part in text.split()]
    words = [word for word in words if word]
    recent = words[-window:]
    counts = {}
    for word in recent:
        counts[word] = counts.get(word, 0) + 1
    return counts


def _violates_word_repeat(decoded_text: str, candidate_piece: str) -> bool:
    candidate_word = _normalize_word(candidate_piece)
    if not candidate_word:
        return False
    counts = _recent_word_counts(decoded_text, window=12)
    return counts.get(candidate_word, 0) >= 2


def _resolve_special_token_ids(tokenizer):
    special_ids = set()
    for token in ["<BOS>", "[BOS]", "<PAD>", "[PAD]", "<SEP>", "[SEP]", "<EOS>", "[EOS]", "<UNK>", "[UNK]", "<CLS>", "[CLS]", "<MASK>", "[MASK]", "<MASK>"]:
        token_id = tokenizer.vocab.get(token)
        if token_id is not None:
            special_ids.add(int(token_id))
    if getattr(tokenizer, "pad_token_id", None) is not None:
        special_ids.add(int(tokenizer.pad_token_id))
    if getattr(tokenizer, "unk_token_id", None) is not None:
        special_ids.add(int(tokenizer.unk_token_id))
    return special_ids


def _contains_special_marker(text: str) -> bool:
    upper_text = text.upper()
    markers = ["<BOS>", "[BOS]", "<PAD>", "[PAD]", "<SEP>", "[SEP]", "<EOS>", "[EOS]", "<UNK>", "[UNK]", "<CLS>", "[CLS]", "<MASK>", "[MASK]"]
    return any(marker in upper_text for marker in markers)


def _looks_like_artifact(text: str) -> bool:
    stripped = text.strip()
    if not stripped:
        return False
    if "##" in stripped:
        return True
    if stripped.startswith("##") or stripped.endswith("##"):
        return True
    if stripped.count("#") >= 1 and len(stripped) <= 4:
        return True
    if "�" in stripped:
        return True
    if any(ch in stripped for ch in ["�", ""]):
        return True
    if re.search(r"(.)\1{3,}", stripped.lower()):
        return True
    if re.fullmatch(r"[A-Za-z]{1,4}\d{2,}", stripped):
        return True
    if re.fullmatch(r"[#\-_=~`|.]+", stripped):
        return True
    return False


def _strip_special_markers(text: str) -> str:
    cleaned = text
    for pattern in [r"<\s*BOS\s*>", r"\[\s*BOS\s*\]", r"<\s*PAD\s*>", r"\[\s*PAD\s*\]", r"<\s*SEP\s*>", r"\[\s*SEP\s*\]", r"<\s*EOS\s*>", r"\[\s*EOS\s*\]", r"<\s*UNK\s*>", r"\[\s*UNK\s*\]", r"<\s*CLS\s*>", r"\[\s*CLS\s*\]", r"<\s*MASK\s*>", r"\[\s*MASK\s*\]"]:
        cleaned = re.sub(pattern, " ", cleaned, flags=re.IGNORECASE)
    cleaned = re.sub(r"#{2,}", " ", cleaned)
    cleaned = re.sub(r"(?<!\w)#(?!\w)", " ", cleaned)
    cleaned = cleaned.replace("�", " ")
    cleaned = re.sub(r"\b([A-Za-z]+)(\s+\1\b){2,}", r"\1", cleaned, flags=re.IGNORECASE)
    cleaned = re.sub(r"\b(\w{1,20})(\w{1,20})\1\b", r"\1", cleaned)
    cleaned = re.sub(r"\s*([,;:.!?])\s*", r"\1 ", cleaned)
    cleaned = re.sub(r"\s+", " ", cleaned)
    return cleaned.strip()


def _cleanup_english_grammar(text: str) -> str:
    cleaned = text.strip()
    if not cleaned:
        return cleaned

    replacements = {
        " im ": " I'm ",
        " ive ": " I've ",
        " ill ": " I'll ",
        " id ": " I'd ",
        " dont ": " don't ",
        " cant ": " can't ",
        " wont ": " won't ",
        " didnt ": " didn't ",
        " doesnt ": " doesn't ",
        " isnt ": " isn't ",
        " arent ": " aren't ",
        " wasnt ": " wasn't ",
        " werent ": " weren't ",
        " thats ": " that's ",
        " whats ": " what's ",
        " theres ": " there's ",
        " ive ": " I've ",
    }

    padded = f" {cleaned} "
    for source, target in replacements.items():
        padded = re.sub(re.escape(source), target, padded, flags=re.IGNORECASE)
    cleaned = padded.strip()

    cleaned = re.sub(r"\bi\b", "I", cleaned)
    cleaned = re.sub(r"\b([A-Za-z]+)(\s+\1\b){1,}", r"\1", cleaned, flags=re.IGNORECASE)
    cleaned = re.sub(r"\s+([,;:.!?])", r"\1", cleaned)
    cleaned = re.sub(r"([,;:.!?])(?!\s|$)", r"\1 ", cleaned)
    cleaned = re.sub(r"\s+", " ", cleaned).strip()

    if cleaned:
        cleaned = cleaned[0].upper() + cleaned[1:]

    sentences = re.split(r"(?<=[.!?])\s+", cleaned)
    normalized_sentences = []
    for sentence in sentences:
        sentence = sentence.strip()
        if not sentence:
            continue
        if len(sentence) == 1:
            normalized_sentences.append(sentence.upper())
        else:
            normalized_sentences.append(sentence[0].upper() + sentence[1:])
    cleaned = " ".join(normalized_sentences).strip()

    if cleaned and cleaned[-1] not in ".!?":
        cleaned += "."

    return cleaned


def _is_strict_english_output(text: str, cot_mode: bool = False) -> bool:
    cleaned = text.strip()
    if not cleaned:
        return False
    if any(token in cleaned for token in ["[", "]", "{", "}", "|", "<UNK>", "[UNK]", "<PAD>", "[PAD]"]):
        return False
    if re.search(r"[^A-Za-z0-9\s,.;:!?\-\'\"()\n]", cleaned):
        return False
    words = re.findall(r"[A-Za-z']+", cleaned)
    if len(words) < 2:
        return False
    long_weird_words = [word for word in words if len(word) > 18]
    if long_weird_words:
        return False
    if re.search(r"([A-Za-z]{2,})([A-Z][a-z]+)", cleaned):
        return False
    common_markers = {
        "the", "a", "an", "is", "are", "am", "i", "you", "we", "they", "it", "to", "of", "and",
        "that", "this", "can", "will", "do", "not", "yes", "no", "my", "your", "in", "on", "for"
    }
    lowered_words = [word.lower() for word in words]
    if not any(word in common_markers for word in lowered_words):
        return False
    if cot_mode:
        lines = [line.strip() for line in cleaned.splitlines() if line.strip()]
        if len(lines) != 2:
            return False
        if not lines[0].startswith("Reasoning:"):
            return False
        if not lines[1].startswith("Answer:"):
            return False
    sentences = [segment.strip() for segment in re.split(r"(?<=[.!?])\s+", cleaned) if segment.strip()]
    if not sentences:
        return False
    for sentence in sentences:
        if not sentence[0].isupper():
            return False
        if sentence[-1] not in ".!?":
            return False
    return True


def _force_cot_shape(text: str) -> str:
    cleaned = text.strip()
    if not cleaned:
        return cleaned
    lines = [line.strip() for line in cleaned.splitlines() if line.strip()]
    if len(lines) >= 2 and lines[0].startswith("Reasoning:") and lines[1].startswith("Answer:"):
        return f"{lines[0]}\n{lines[1]}"
    parts = re.split(r"(?<=[.!?])\s+", cleaned, maxsplit=1)
    if len(parts) == 2:
        reasoning, answer = parts
    else:
        reasoning, answer = "Reasoning: Briefly considered the request.", f"Answer: {cleaned}"
        return f"{reasoning}\n{answer}"
    reasoning = reasoning if reasoning.startswith("Reasoning:") else f"Reasoning: {reasoning.strip()}"
    answer = answer if answer.startswith("Answer:") else f"Answer: {answer.strip()}"
    return f"{reasoning}\n{answer}"


def _ban_low_quality_candidates(tokenizer, logits: torch.Tensor):
    for token_id in range(logits.size(0)):
        piece = tokenizer.decode([token_id]).strip()
        if not piece:
            continue
        if _contains_special_marker(piece):
            logits[token_id] = float("-inf")


def _select_candidate_id(tokenizer, probs: torch.Tensor, generated, prompt_token_count: int, no_repeat_ngram_size: int):
    candidate_count = min(24, probs.size(0))
    top_probs, top_ids = torch.topk(probs, candidate_count)
    decoded_so_far = tokenizer.decode(generated[prompt_token_count:]).strip()

    fallback_clean_id = None
    fallback_clean_prob = -1.0
    fallback_any_id = None
    fallback_any_prob = -1.0
    for prob_value, candidate_id_tensor in zip(top_probs.tolist(), top_ids.tolist()):
        candidate_id = int(candidate_id_tensor)
        candidate_piece = tokenizer.decode([candidate_id]).strip()
        if not candidate_piece:
            continue
        if _contains_special_marker(candidate_piece):
            continue
        if fallback_any_id is None or prob_value > fallback_any_prob:
            fallback_any_id = candidate_id
            fallback_any_prob = prob_value
        if _looks_like_artifact(candidate_piece):
            continue
        if _violates_word_repeat(decoded_so_far, candidate_piece):
            continue
        if _has_repeat_ngram(generated, candidate_id, max(no_repeat_ngram_size, 4)):
            continue
        normalized_piece = _normalize_word(candidate_piece)
        if normalized_piece and decoded_so_far:
            recent_words = [_normalize_word(part) for part in decoded_so_far.split()[-8:]]
            recent_words = [word for word in recent_words if word]
            if recent_words.count(normalized_piece) >= 1:
                continue
        if fallback_clean_id is None or prob_value > fallback_clean_prob:
            fallback_clean_id = candidate_id
            fallback_clean_prob = prob_value

    if fallback_clean_id is not None:
        return fallback_clean_id
    return fallback_any_id


def _generate_fallback_reply(model, tokenizer, prompt_tokens, blocked_special_ids, max_length: int):
    device = next(model.parameters()).device
    generated = list(prompt_tokens)

    with torch.no_grad():
        for _ in range(min(max_length, 16)):
            current_len = len(generated)
            chunk_size = _model_chunk_size(model)
            pad_len = (chunk_size - current_len % chunk_size) % chunk_size
            padded_input = generated + [tokenizer.pad_token_id] * pad_len
            input_tensor = torch.tensor([padded_input], device=device)
            next_token_logits = _next_token_logits(model, input_tensor, current_len)

            for special_id in blocked_special_ids:
                if 0 <= special_id < next_token_logits.size(0):
                    next_token_logits[special_id] = float("-inf")

            next_token_id = int(torch.argmax(next_token_logits).item())
            if next_token_id == tokenizer.pad_token_id:
                break

            next_piece = tokenizer.decode([next_token_id]).strip()
            if not next_piece or _contains_special_marker(next_piece):
                break

            generated.append(next_token_id)

    return generated


def generate_reply(
    model,
    tokenizer,
    prompt: str,
    max_length: int,
    temperature: float,
    top_k: int,
    top_p: float,
    repetition_penalty: float,
    no_repeat_ngram_size: int,
    cot_mode: bool = False,
):
    model.eval()
    device = next(model.parameters()).device

    formatted_prompt = build_prompt(prompt, cot_mode=cot_mode)
    prompt_ids = tokenizer.encode(formatted_prompt, max_length=128)
    generated = [tok for tok in prompt_ids if tok != tokenizer.pad_token_id]
    prompt_token_count = len(generated)
    blocked_special_ids = _resolve_special_token_ids(tokenizer)
    bos_token_id = tokenizer.vocab.get("<BOS>")
    if not generated:
        generated = [tokenizer.unk_token_id]
        prompt_token_count = len(generated)

    with torch.no_grad():
        for _ in range(max_length):
            current_len = len(generated)
            chunk_size = _model_chunk_size(model)
            pad_len = (chunk_size - current_len % chunk_size) % chunk_size
            padded_input = generated + [tokenizer.pad_token_id] * pad_len
            input_tensor = torch.tensor([padded_input], device=device)

            next_token_logits = _next_token_logits(model, input_tensor, current_len)

            if temperature > 0:
                next_token_logits = next_token_logits / temperature

            for special_id in blocked_special_ids:
                if 0 <= special_id < next_token_logits.size(0):
                    next_token_logits[special_id] = float("-inf")

            if bos_token_id is not None and 0 <= int(bos_token_id) < next_token_logits.size(0):
                next_token_logits[int(bos_token_id)] = float("-inf")

            _ban_low_quality_candidates(tokenizer, next_token_logits)

            recent_tokens = generated[-48:]
            recent_weights = {}
            for idx, token_id in enumerate(recent_tokens):
                distance_weight = 1.0 + (idx / max(len(recent_tokens), 1))
                recent_weights[token_id] = max(recent_weights.get(token_id, 1.0), distance_weight)

            for token_id, distance_weight in recent_weights.items():
                if 0 <= token_id < next_token_logits.size(0):
                    penalty = repetition_penalty * distance_weight
                    if next_token_logits[token_id] > 0:
                        next_token_logits[token_id] /= penalty
                    else:
                        next_token_logits[token_id] *= penalty

            for token_id in range(next_token_logits.size(0)):
                if _has_repeat_ngram(generated, token_id, no_repeat_ngram_size):
                    next_token_logits[token_id] = float("-inf")

            if top_k > 0 and top_k < next_token_logits.size(0):
                threshold = torch.topk(next_token_logits, top_k)[0][..., -1]
                next_token_logits[next_token_logits < threshold] = float("-inf")

            next_token_logits = _apply_top_p_filter(next_token_logits, top_p)
            probs = F.softmax(next_token_logits, dim=-1)

            if torch.isnan(probs).any() or torch.isinf(probs).any() or probs.sum() <= 0:
                break

            next_token = _select_candidate_id(
                tokenizer,
                probs,
                generated,
                prompt_token_count,
                no_repeat_ngram_size,
            )

            if next_token is None:
                break

            if next_token == tokenizer.pad_token_id:
                break
            generated.append(next_token)

            decoded_output = tokenizer.decode(generated[prompt_token_count:]).strip()
            if len(decoded_output.split()) >= 6:
                tail_words = [_normalize_word(part) for part in decoded_output.split()[-4:]]
                tail_words = [word for word in tail_words if word]
                if len(tail_words) >= 4 and len(set(tail_words)) == 1:
                    break

    output_ids = generated[prompt_token_count:]
    cleaned_output = _strip_special_markers(tokenizer.decode(output_ids).strip())
    if cleaned_output:
        normalized_output = _cleanup_english_grammar(cleaned_output)
        if cot_mode:
            normalized_output = _force_cot_shape(normalized_output)
        return normalized_output

    fallback_generated = _generate_fallback_reply(
        model,
        tokenizer,
        generated[:prompt_token_count],
        blocked_special_ids,
        max_length,
    )
    fallback_output_ids = fallback_generated[prompt_token_count:]
    fallback_output = _strip_special_markers(tokenizer.decode(fallback_output_ids).strip())
    normalized_fallback_output = _cleanup_english_grammar(fallback_output)
    if cot_mode:
        normalized_fallback_output = _force_cot_shape(normalized_fallback_output)
    return normalized_fallback_output


def main():
    parser = argparse.ArgumentParser(description="Chat/test with pretrained RubiNet HSSM checkpoint")
    parser.add_argument("--checkpoint", default=DEFAULT_CHECKPOINT)
    parser.add_argument("--tokenizer", default=DEFAULT_TOKENIZER)
    parser.add_argument("--device", default="cuda" if torch.cuda.is_available() else "cpu")
    parser.add_argument("--max-length", type=int, default=40)
    parser.add_argument("--temperature", type=float, default=0.0)
    parser.add_argument("--top-k", type=int, default=4)
    parser.add_argument("--top-p", type=float, default=0.65)
    parser.add_argument("--repetition-penalty", type=float, default=1.9)
    parser.add_argument("--no-repeat-ngram-size", type=int, default=6)
    parser.add_argument("--cot-mode", action="store_true")
    parser.add_argument("--no-cot-mode", action="store_false", dest="cot_mode")
    parser.set_defaults(cot_mode=False)
    parser.add_argument("--message", default="")
    args = parser.parse_args()

    tokenizer, model = load_pretrained(args.checkpoint, args.tokenizer, args.device)

    if args.message:
        output = generate_reply(
            model,
            tokenizer,
            args.message,
            args.max_length,
            args.temperature,
            args.top_k,
            args.top_p,
            args.repetition_penalty,
            args.no_repeat_ngram_size,
            args.cot_mode,
        )
        safe_print(output)
        return

    print("Interactive HSSM chat/test. Type 'exit' to quit.")
    while True:
        user_text = input("You: ").strip()
        if not user_text:
            continue
        if user_text.lower() in {"exit", "quit"}:
            break
        output = generate_reply(
            model,
            tokenizer,
            user_text,
            args.max_length,
            args.temperature,
            args.top_k,
            args.top_p,
            args.repetition_penalty,
            args.no_repeat_ngram_size,
            args.cot_mode,
        )
        safe_print(f"HSSM: {output}\n")


if __name__ == "__main__":
    main()