"""
Chat interface for CosmicFish model downloaded from Hugging Face Hub.
Uses safetensors format only for secure model loading.
"""

import os
import sys
import time
import argparse
import torch
import numpy as np
from termcolor import colored
import logging
import readline
import re
import textwrap
import random
from collections import defaultdict
import json

# Required imports for HF Hub
try:
    from transformers import GPT2Tokenizer
    from huggingface_hub import hf_hub_download, snapshot_download
    HF_AVAILABLE = True
except ImportError:
    HF_AVAILABLE = False
    print("Required libraries not available.")
    print("Install with: pip install transformers huggingface-hub")
    sys.exit(1)

# Required for safetensors
try:
    from safetensors.torch import load_file
    SAFETENSORS_AVAILABLE = True
except ImportError:
    SAFETENSORS_AVAILABLE = False
    print("Safetensors not available. Install with: pip install safetensors")
    sys.exit(1)

# Set up logging
logging.basicConfig(
    level=logging.INFO,
    format='%(asctime)s - %(levelname)s - %(message)s',
    handlers=[logging.StreamHandler(sys.stdout)]
)
logger = logging.getLogger(__name__)

# Default model repository
DEFAULT_MODEL_REPO = "MistyozAI/CosmicFish-90M"

# Default prompt template
DEFAULT_PROMPT_TEMPLATE = "Below is a conversation between a helpful AI assistant and a human. The assistant is knowledgeable, friendly, and provides detailed and accurate responses.\n\n"


class CosmicConfig:
    """Configuration class for CosmicFish."""

    def __init__(self,
                 vocab_size=50257,
                 block_size=512,
                 n_layer=10,
                 n_head=16,
                 n_embd=640,
                 bias=True,
                 dropout=0.0,
                 n_query_groups=4,
                 eps=1e-6,
                 use_rotary=True,
                 use_swiglu=True,
                 use_qk_norm=False,
                 use_gqa=True):
        self.vocab_size = vocab_size
        self.block_size = block_size
        self.n_layer = n_layer
        self.n_head = n_head
        self.n_embd = n_embd
        self.bias = bias
        self.dropout = dropout
        self.eps = eps
        self.use_rotary = use_rotary
        self.use_swiglu = use_swiglu
        self.use_qk_norm = use_qk_norm
        self.use_gqa = use_gqa
        self.n_query_groups = n_query_groups if use_gqa else n_head
        # Ensure n_head is divisible by n_query_groups
        assert n_head % self.n_query_groups == 0, "n_head must be divisible by n_query_groups"


class RMSNorm(torch.nn.Module):
    """Root Mean Square Normalization"""

    def __init__(self, dim, eps=1e-6):
        super().__init__()
        self.eps = eps
        self.weight = torch.nn.Parameter(torch.ones(dim))

    def forward(self, x):
        rms = torch.sqrt(torch.mean(x * x, dim=-1, keepdim=True) + self.eps)
        return self.weight * (x / rms)


def precompute_freqs_cis(dim, end, theta=10000.0):
    """Precompute the frequency tensor for complex exponentials (cis)"""
    freqs = 1.0 / (theta ** (torch.arange(0, dim, 2)[: (dim // 2)].float() / dim))
    t = torch.arange(end, device=freqs.device)
    freqs = torch.outer(t, freqs)
    freqs_cis = torch.polar(torch.ones_like(freqs), freqs)
    return freqs_cis


def apply_rotary_emb(xq, xk, freqs_cis):
    """Apply rotary embeddings to input tensors"""
    xq_ = torch.view_as_complex(xq.float().reshape(*xq.shape[:-1], -1, 2))
    xk_ = torch.view_as_complex(xk.float().reshape(*xk.shape[:-1], -1, 2))

    seq_len = xq_.size(2)
    if freqs_cis.size(0) < seq_len:
        raise ValueError(f"freqs_cis has only {freqs_cis.size(0)} values but sequence length is {seq_len}")

    freqs_cis_seq = freqs_cis[:seq_len]
    xq_out = torch.view_as_real(xq_ * freqs_cis_seq.unsqueeze(0)).flatten(3)
    xk_out = torch.view_as_real(xk_ * freqs_cis_seq.unsqueeze(0)).flatten(3)

    return xq_out.type_as(xq), xk_out.type_as(xk)


class GroupedQueryAttention(torch.nn.Module):
    """Grouped Query Attention (GQA) implementation"""

    def __init__(self, config):
        super().__init__()
        assert config.n_embd % config.n_head == 0

        head_dim = config.n_embd // config.n_head
        self.head_dim = head_dim
        self.n_head = config.n_head
        self.n_embd = config.n_embd
        self.n_query_groups = config.n_query_groups

        self.kv_heads = config.n_head // config.n_query_groups if config.use_gqa else config.n_head
        qkv_proj_size = (config.n_head + 2 * self.kv_heads) * head_dim

        self.c_attn = torch.nn.Linear(config.n_embd, qkv_proj_size, bias=config.bias)
        self.c_proj = torch.nn.Linear(config.n_embd, config.n_embd, bias=config.bias)

        # Flash attention support
        self.flash = hasattr(torch.nn.functional, 'scaled_dot_product_attention')
        if not self.flash:
            self.register_buffer("bias", torch.tril(torch.ones(config.block_size, config.block_size))
                                 .view(1, 1, config.block_size, config.block_size))

        # Query-key normalization
        self.qk_norm = getattr(config, 'use_qk_norm', False)
        if self.qk_norm:
            self.q_norm = RMSNorm(head_dim, eps=getattr(config, 'eps', 1e-6))
            self.k_norm = RMSNorm(head_dim, eps=getattr(config, 'eps', 1e-6))

    def forward(self, x, freqs_cis=None):
        B, T, C = x.size()
        qkv = self.c_attn(x)
        head_dim = C // self.n_head

        q_size = self.n_head * head_dim
        k_size = self.kv_heads * head_dim
        v_size = self.kv_heads * head_dim

        q, k, v = qkv.split([q_size, k_size, v_size], dim=2)

        q = q.view(B, T, self.n_head, head_dim).transpose(1, 2)
        k = k.view(B, T, self.kv_heads, head_dim).transpose(1, 2)
        v = v.view(B, T, self.kv_heads, head_dim).transpose(1, 2)

        # Repeat k and v if needed for GQA
        if self.kv_heads < self.n_head:
            repeats = self.n_head // self.kv_heads
            k = k.repeat_interleave(repeats, dim=1)
            v = v.repeat_interleave(repeats, dim=1)

        # Apply rotary embeddings
        if freqs_cis is not None:
            q, k = apply_rotary_emb(q, k, freqs_cis)

        # Apply query-key normalization
        if self.qk_norm:
            q = self.q_norm(q)
            k = self.k_norm(k)

        # Compute attention
        if self.flash:
            y = torch.nn.functional.scaled_dot_product_attention(
                q, k, v, attn_mask=None, dropout_p=0.0, is_causal=True
            )
        else:
            att = (q @ k.transpose(-2, -1)) * (1.0 / torch.sqrt(torch.tensor(k.size(-1), dtype=torch.float32)))
            att = att.masked_fill(self.bias[:, :, :T, :T] == 0, float('-inf'))
            att = torch.nn.functional.softmax(att, dim=-1)
            y = att @ v

        y = y.transpose(1, 2).contiguous().view(B, T, C)
        y = self.c_proj(y)
        return y


class Block(torch.nn.Module):
    """Transformer block"""

    def __init__(self, config):
        super().__init__()
        self.ln_1 = RMSNorm(config.n_embd, eps=config.eps)
        self.ln_2 = RMSNorm(config.n_embd, eps=config.eps)
        self.attn = GroupedQueryAttention(config)

        # MLP implementation based on configuration
        if config.use_swiglu:
            # SwiGLU MLP
            self.mlp = torch.nn.ModuleDict(dict(
                gate=torch.nn.Linear(config.n_embd, 4 * config.n_embd, bias=config.bias),
                up=torch.nn.Linear(config.n_embd, 4 * config.n_embd, bias=config.bias),
                down=torch.nn.Linear(4 * config.n_embd, config.n_embd, bias=config.bias),
                act=torch.nn.SiLU(),
            ))
            m = self.mlp
            self.mlpf = lambda x: m.down(m.act(m.up(x)) * m.gate(x))
        else:
            # Traditional MLP
            self.mlp = torch.nn.ModuleDict(dict(
                c_fc=torch.nn.Linear(config.n_embd, 4 * config.n_embd, bias=config.bias),
                c_proj=torch.nn.Linear(4 * config.n_embd, config.n_embd, bias=config.bias),
                act=torch.nn.GELU(),
            ))
            m = self.mlp
            self.mlpf = lambda x: m.c_proj(m.act(m.c_fc(x)))

    def forward(self, x, freqs_cis=None):
        x = x + self.attn(self.ln_1(x), freqs_cis)
        x = x + self.mlpf(self.ln_2(x))
        return x


class CosmicFish(torch.nn.Module):
    """
    CosmicFish model for inference only.
    Features: Rotary Positional Embeddings, Grouped-Query Attention, SwiGLU, RMSNorm
    """

    def __init__(self, config):
        super().__init__()
        self.config = config

        self.transformer = torch.nn.ModuleDict(dict(
            wte=torch.nn.Embedding(config.vocab_size, config.n_embd),
            h=torch.nn.ModuleList([Block(config) for _ in range(config.n_layer)]),
            ln_f=RMSNorm(config.n_embd, eps=config.eps),
        ))

        self.lm_head = torch.nn.Linear(config.n_embd, config.vocab_size, bias=False)

        # Share weights between embedding and output
        self.transformer.wte.weight = self.lm_head.weight

        # Precompute rotary embedding frequencies
        if config.use_rotary:
            head_dim = config.n_embd // config.n_head
            self.freqs_cis = precompute_freqs_cis(head_dim, config.block_size)
        else:
            self.freqs_cis = None
            self.transformer.wpe = torch.nn.Embedding(config.block_size, config.n_embd)

    def get_num_params(self, non_embedding=True):
        """Return the number of parameters in the model."""
        n_params = sum(p.numel() for p in self.parameters())
        if non_embedding and hasattr(self.transformer, 'wpe'):
            n_params -= self.transformer.wpe.weight.numel()
        return n_params

    def forward(self, idx, targets=None):
        """Forward pass through the model."""
        device = idx.device
        b, t = idx.size()
        assert t <= self.config.block_size, f"Cannot forward sequence of length {t}, block size is only {self.config.block_size}"

        # Get token embeddings
        tok_emb = self.transformer.wte(idx)

        # Handle positional embeddings
        if self.config.use_rotary:
            x = tok_emb
            freqs_cis = self.freqs_cis.to(device) if self.freqs_cis is not None else None
        else:
            pos = torch.arange(0, t, dtype=torch.long, device=device).unsqueeze(0)
            pos_emb = self.transformer.wpe(pos)
            x = tok_emb + pos_emb
            freqs_cis = None

        # Apply transformer blocks
        for block in self.transformer.h:
            x = block(x, freqs_cis)

        # Apply final normalization
        x = self.transformer.ln_f(x)

        # Calculate outputs
        if targets is not None:
            logits = self.lm_head(x)
            loss = torch.nn.functional.cross_entropy(logits.view(-1, logits.size(-1)), targets.view(-1), ignore_index=-1)
        else:
            # For inference, only compute logits for the last token
            logits = self.lm_head(x[:, [-1], :])
            loss = None

        return logits, loss

    @torch.no_grad()
    def generate(self, idx, max_new_tokens, temperature=1.0, top_k=None):
        """
        Generate text by sampling from the model, token by token.
        """
        for _ in range(max_new_tokens):
            # Crop sequence to block size if needed
            idx_cond = idx if idx.size(1) <= self.config.block_size else idx[:, -self.config.block_size:]

            # Forward pass
            logits, _ = self(idx_cond)
            logits = logits[:, -1, :] / temperature

            # Apply top-k sampling
            if top_k is not None:
                v, _ = torch.topk(logits, top_k)
                logits[logits < v[:, [-1]]] = -float('Inf')

            # Sample next token
            probs = torch.nn.functional.softmax(logits, dim=-1)
            idx_next = torch.multinomial(probs, num_samples=1)

            # Append to sequence
            idx = torch.cat((idx, idx_next), dim=1)

        return idx


class RepetitionPenaltyLogitsProcessor:
    """Apply repetition penalty to prevent repeating tokens."""

    def __init__(self, penalty=1.2):
        self.penalty = penalty

    def __call__(self, input_ids, scores):
        """Apply repetition penalty to logits where input_ids is already seen."""
        score = torch.gather(scores, 1, input_ids)
        # If score > 0, penalize by dividing; if score < 0, penalize by multiplying
        score = torch.where(score > 0, score / self.penalty, score * self.penalty)
        scores.scatter_(1, input_ids, score)
        return scores


class CosmicFishChatSession:
    """Chat session for CosmicFish model from Hugging Face Hub."""

    def __init__(self, model, tokenizer, config):
        """Initialize chat session with model and configuration."""
        self.model = model
        self.tokenizer = tokenizer
        self.config = config
        self.device = next(model.parameters()).device
        self.history = []
        self.history_tokens = []
        self.max_history_tokens = config.max_history_tokens
        self.prompt_template = config.prompt_template
        self.human_prefix = config.human_prefix
        self.assistant_prefix = config.assistant_prefix
        self.end_of_turn = config.end_of_turn
        self.block_size = config.block_size
        self.debug_mode = config.debug_mode
        self.repetition_penalty = config.repetition_penalty
        self.min_tokens_to_generate = config.min_tokens_to_generate
        self.max_retries = 20

        self.fallback_responses = [
            "I'd be happy to help with that. Could you provide more details about what specific information you're looking for?",
            "That's a topic I can provide information about. What specific aspects would you like to know?",
            "I understand your question. I can share factual information on this topic if you could specify what aspects you're interested in.",
            "I can help with your question. To give you the most relevant information, could you clarify what specific details you're looking for?",
            "I'd be glad to address your question. To provide the most helpful response, could you specify what particular aspects of this topic interest you?"
        ]

        self.generation_failure_message = "I'm sorry, but I'm having difficulty generating a response to that prompt. Could you try rephrasing your question or asking something else?"

        # For token counting
        self.total_prompt_tokens = 0
        self.total_generated_tokens = 0

        # End markers for live generation
        self.end_markers = [
            f"{self.human_prefix}",
            "Human:",
            "\nHuman:",
            "\nH:",
            "H:",
            "<|endoftext|>",
            "Below is a conversation",
            "\nA:",
            "A:",
            "</s>",
            "User:",
            "\nUser:"
        ]

        if config.display_welcome:
            self._print_welcome_message()

    def _print_welcome_message(self):
        welcome_text = f"""
{'=' * 80}
Welcome to CosmicFish chat interface

This is a {self.model.get_num_params() / 1e6:.1f}M parameter model.
CosmicFish is an efficient LLM with an advanced architecture.

Type your prompts and CosmicFish will respond.

Special commands:
- /help: Show this help message
- /clear: Clear the conversation history
- /exit or /quit: Exit the chat
- /stats: Show token usage statistics
- /save [filename]: Save the conversation
- /load [filename]: Load a conversation
- /temp [value]: Set temperature (between 0.1 and 2.0)
- /penalty [value]: Set repetition penalty (1.0-2.0)
- /debug: Toggle debug mode


Note: CosmicFIsh may generate incorrect or fictional responses. Verify facts if needed.

Visit https://cosmicfish.ai for more info


Developed by Mistyoz AI (https://www.mistyoz.com)
{'=' * 80}
"""
        print(colored(welcome_text, 'cyan'))

    def _format_prompt(self, user_input):
        """Format the complete prompt with history and current input."""
        # Start with the template
        formatted_prompt = self.prompt_template

        # Add conversation history
        for entry in self.history:
            role, text = entry
            if role == "human":
                formatted_prompt += f"{self.human_prefix}{text}{self.end_of_turn}"
            else:  # assistant
                formatted_prompt += f"{self.assistant_prefix}{text}{self.end_of_turn}"

        # Add the current user input
        formatted_prompt += f"{self.human_prefix}{user_input}{self.end_of_turn}{self.assistant_prefix}"

        return formatted_prompt

    def _tokenize(self, text):
        """Tokenize text and return token IDs."""
        return self.tokenizer.encode(text)

    def _update_history(self, user_input, response):
        """Update conversation history."""
        # Add to text history
        self.history.append(("human", user_input))
        self.history.append(("assistant", response))

        # Update token history for context window management
        user_tokens = self._tokenize(f"{self.human_prefix}{user_input}{self.end_of_turn}")
        response_tokens = self._tokenize(f"{self.assistant_prefix}{response}{self.end_of_turn}")

        self.history_tokens.extend(user_tokens)
        self.history_tokens.extend(response_tokens)

        # Track token usage
        self.total_prompt_tokens += len(user_tokens)
        self.total_generated_tokens += len(response_tokens)

        # Trim history if it gets too long
        self._trim_history_if_needed()

    def _trim_history_if_needed(self):
        """Trim history to fit within the context window."""
        if len(self.history_tokens) > self.max_history_tokens:
            # Remove oldest turns until we're under the limit
            while len(self.history_tokens) > self.max_history_tokens and len(self.history) >= 2:
                # Remove oldest human and assistant turn
                self.history = self.history[2:]

                # Find token boundary for the removed turns
                user_turn = self.history[0][1]
                assistant_turn = self.history[1][1]
                user_tokens = len(self._tokenize(f"{self.human_prefix}{user_turn}{self.end_of_turn}"))
                assistant_tokens = len(self._tokenize(f"{self.assistant_prefix}{assistant_turn}{self.end_of_turn}"))

                # Update token history
                self.history_tokens = self.history_tokens[user_tokens + assistant_tokens:]

    def _should_stop_generation(self, text):
        """Check if generation should stop based on end markers."""
        for marker in self.end_markers:
            if marker in text:
                return True
        return False

    def _clean_token_text(self, text):
        text = text.replace('��', "'")
        text = text.replace('�', "'")
        text = text.replace('\ufffd', "'")
        text = text.replace('\uFFFD', "'")
        text = text.replace('’', "'")
        text = text.replace('â€Å"', "'")
        text = text.replace('�', "'")
        text = text.replace('â€"', "'")
        text = text.replace('â€"', "'")
        return text

    def generate_with_repetition_penalty(self, input_ids, max_new_tokens, temperature, top_k, penalty=1.2, live=False):
        """Custom generate function with repetition penalty and optional live generation."""
        model = self.model
        device = self.device

        # Ensure model is in eval mode
        model.eval()

        # Initialize sequence with input_ids
        generated = input_ids.clone()

        # Initialize live text buffer
        live_buffer = ""

        # Create repetition penalty processor
        rep_processor = RepetitionPenaltyLogitsProcessor(penalty=penalty)

        # Counter for generated tokens
        tokens_generated = 0
        min_tokens = self.min_tokens_to_generate

        # EOT token ID
        eot_token_id = self.tokenizer.eos_token_id if hasattr(self.tokenizer, 'eos_token_id') else 50256

        # Generate tokens one at a time
        for _ in range(max_new_tokens):
            # Get only the last block_size tokens if context is too long
            if generated.size(1) > self.block_size:
                context = generated[:, -self.block_size:]
            else:
                context = generated

            # Forward pass for next token prediction
            with torch.no_grad():
                logits, _ = model(context)

            # Get logits for the next token (last position)
            next_token_logits = logits[:, -1, :]

            # Apply temperature
            next_token_logits = next_token_logits / temperature

            # Apply repetition penalty
            if penalty > 1.0:
                next_token_logits = rep_processor(context, next_token_logits)

            # Optional top-k sampling
            if top_k is not None:
                indices_to_remove = next_token_logits < torch.topk(next_token_logits, top_k)[0][..., -1, None]
                next_token_logits[indices_to_remove] = float('-inf')

            # Convert logits to probabilities
            probs = torch.nn.functional.softmax(next_token_logits, dim=-1)

            # Sample next token
            next_token = torch.multinomial(probs, num_samples=1)

            # Check if the next token is EOT and break immediately if so
            if next_token.item() == eot_token_id:
                if live:
                    yield "", live_buffer, True
                break

            # Append next token to generated sequence
            generated = torch.cat((generated, next_token), dim=1)
            tokens_generated += 1

            # If live generation, decode and yield the token
            if live:
                # Decode the next token
                next_token_text = self.tokenizer.decode([next_token.item()])
                # Clean the token text to fix encoding issues
                next_token_text = self._clean_token_text(next_token_text)
                live_buffer += next_token_text

                # Check if we've hit an end marker in the buffer
                eot_marker_pos = live_buffer.find("<|endoftext|>")
                if eot_marker_pos != -1:
                    # Cut off at the EOT marker
                    live_buffer = live_buffer[:eot_marker_pos]
                    yield "", live_buffer, True
                    break

                # Check other end markers
                should_stop = tokens_generated >= min_tokens and self._should_stop_generation(live_buffer)
                yield next_token_text, live_buffer, should_stop

                if should_stop:
                    break

            # For non-live generation, check if we should stop
            elif tokens_generated >= min_tokens:
                # Check for end markers in the recent generated tokens
                recent_text = self.tokenizer.decode(generated[0, -20:].tolist())
                if self._should_stop_generation(recent_text):
                    break

        # Check if we generated any tokens at all
        if tokens_generated == 0 and not live:
            if self.debug_mode:
                print(colored("\n[No tokens generated in this attempt]", "red"))
            return None

        if not live:
            return generated

    def generate_response(self, user_input):
        """Generate a response to the user input."""
        # Format the complete prompt
        prompt = self._format_prompt(user_input)

        # Tokenize the prompt
        input_ids = torch.tensor(self._tokenize(prompt), dtype=torch.long).unsqueeze(0).to(self.device)

        # Ensure we don't exceed the model's context length
        if input_ids.size(1) > self.block_size:
            # If too long, keep the beginning part with the instruction template and trim the middle
            instruction_tokens = self._tokenize(self.prompt_template)
            # Keep the instruction and the most recent conversation that will fit
            keep_from_beginning = len(instruction_tokens)
            keep_from_end = self.block_size - keep_from_beginning

            # Combine beginning and end, ensuring we don't exceed array bounds
            if keep_from_end < 0:
                # If instruction alone is too long, trim it (shouldn't happen with reasonable templates)
                input_ids = input_ids[:, :self.block_size]
            else:
                # Keep instruction and most recent conversation
                input_ids = torch.cat([
                    input_ids[:, :keep_from_beginning],
                    input_ids[:, -(keep_from_end):]
                ], dim=1)

        # Track generation start time
        start_time = time.time()

        # Always use live generation
        return self._generate_live_response(input_ids, user_input, start_time)

    def _generate_live_response(self, input_ids, user_input, start_time):
        """Generate response with live token-by-token output."""
        # Initialize for live generation
        live_text = ""
        tokens_generated = 0
        retry_count = 0

        # Keep trying until we get a valid response or exhaust retries
        while retry_count <= self.max_retries:
            if retry_count > 0:
                # Calculate temperature for this retry
                if retry_count % 2 == 0:
                    # Even retries: increase temperature
                    temp_adjustment = min(0.2 * (retry_count // 2), 0.8)
                    current_temp = min(self.config.temperature + temp_adjustment, 1.8)
                else:
                    # Odd retries: decrease temperature
                    temp_adjustment = min(0.2 * ((retry_count + 1) // 2), 0.4)
                    current_temp = max(self.config.temperature - temp_adjustment, 0.2)

                if self.debug_mode:
                    print(colored(f"\n[Live retry {retry_count}: Using temperature {current_temp:.2f}]", "yellow"))
            else:
                current_temp = self.config.temperature

            # Reset for this attempt
            live_text = ""
            tokens_generated = 0
            generation_failed = False

            # Try to generate with current settings
            try:
                # Generate with live output
                for token_text, live_buffer, should_stop in self.generate_with_repetition_penalty(
                        input_ids,
                        max_new_tokens=self.config.max_new_tokens,
                        temperature=current_temp,
                        top_k=self.config.top_k,
                        penalty=self.repetition_penalty,
                        live=True
                ):
                    # If we should stop but there's a token, this is the last one
                    if should_stop:
                        # Update with the final clean buffer (will have EOT removed if present)
                        live_text = live_buffer
                        break

                    # Otherwise add the token and continue
                    if token_text:
                        live_text += token_text
                        tokens_generated += 1
                        yield token_text, live_text, False

                # Check if we got a valid response
                if not live_text or len(live_text.strip()) < 10:
                    if self.debug_mode:
                        print(colored("\n[Live generation produced empty or too short response, retrying]", "yellow"))
                    generation_failed = True
                    retry_count += 1
                    # Clear any partial output
                    if retry_count <= self.max_retries:
                        print("\r" + " " * 80 + "\r", end="")  # Clear the line
                else:
                    # We got a valid response, stop retrying
                    break

            except Exception as e:
                if self.debug_mode:
                    print(colored(f"\n[Live generation error: {str(e)}, retrying]", "red"))
                generation_failed = True
                retry_count += 1

        # If we still failed after all retries, use the failure message
        if generation_failed or not live_text or len(live_text.strip()) < 10:
            live_text = self.generation_failure_message
            if self.debug_mode:
                print(colored(f"\n[Returning failure message after {retry_count} live retries]", "red"))

        # Calculate time taken and metrics
        time_taken = time.time() - start_time
        tokens_per_second = tokens_generated / time_taken if time_taken > 0 else 0

        # Update history
        self._update_history(user_input, live_text)

        # Log generation stats
        logger.debug(f"Generated {tokens_generated} tokens in {time_taken:.2f}s ({tokens_per_second:.2f} tokens/s)")

        # Final yield of the complete response
        yield "", live_text, True

    def execute_command(self, command):
        """Execute a special command prefixed with /."""
        command = command.strip()

        if command == '/help':
            self._print_welcome_message()
            return True

        elif command == '/clear':
            self.history = []
            self.history_tokens = []
            print(colored("Conversation history cleared.", 'yellow'))
            return True

        elif command in ['/exit', '/quit']:
            print(colored("Goodbye!", 'cyan'))
            return False  # Signal to exit the chat loop

        elif command == '/stats':
            prompt_tokens = self.total_prompt_tokens
            generated_tokens = self.total_generated_tokens
            total_tokens = prompt_tokens + generated_tokens

            stats = f"""
Token usage statistics:
- Prompt tokens: {prompt_tokens}
- Generated tokens: {generated_tokens}
- Total tokens: {total_tokens}
- Current history length: {len(self.history_tokens)} tokens
- Current repetition penalty: {self.repetition_penalty}
- Current temperature: {self.config.temperature}
- Model: CosmicFish ({self.model.get_num_params() / 1e6:.1f}M parameters)
- Source: {DEFAULT_MODEL_REPO}
- Format: Safetensors (secure)
"""
            print(colored(stats, 'yellow'))
            return True

        elif command == '/debug':
            self.debug_mode = not self.debug_mode
            self.config.debug_mode = self.debug_mode  # Sync with config
            mode = "enabled" if self.debug_mode else "disabled"
            print(colored(f"Debug mode {mode}", 'yellow'))
            return True

        elif command.startswith('/penalty '):
            try:
                penalty = float(command[9:].strip())
                if 1.0 <= penalty <= 2.0:
                    self.repetition_penalty = penalty
                    print(colored(f"Repetition penalty set to {penalty}", 'yellow'))
                else:
                    print(colored("Repetition penalty should be between 1.0 and 2.0", 'red'))
            except ValueError:
                print(colored("Invalid repetition penalty value. Please use a number between 1.0 and 2.0", 'red'))
            return True

        elif command.startswith('/temp '):
            try:
                temp = float(command[6:].strip())
                if 0.1 <= temp <= 2.0:
                    self.config.temperature = temp
                    print(colored(f"Temperature set to {temp}", 'yellow'))
                else:
                    print(colored("Temperature should be between 0.1 and 2.0", 'red'))
            except ValueError:
                print(colored("Invalid temperature value. Please use a number between 0.1 and 2.0", 'red'))
            return True

        elif command.startswith('/save '):
            filename = command[6:].strip()
            if not filename:
                print(colored("Please specify a filename: /save <filename>", 'red'))
                return True

            try:
                # Create conversations directory if it doesn't exist
                os.makedirs('conversations', exist_ok=True)

                # Add .txt extension if not present
                if not filename.endswith('.txt'):
                    filename += '.txt'

                filepath = os.path.join('conversations', filename)

                with open(filepath, 'w', encoding='utf-8') as f:
                    for entry in self.history:
                        role, text = entry
                        prefix = self.human_prefix if role == "human" else self.assistant_prefix
                        f.write(f"{prefix}{text}{self.end_of_turn}")

                print(colored(f"Conversation saved to {filepath}", 'green'))

            except Exception as e:
                print(colored(f"Error saving conversation: {str(e)}", 'red'))

            return True

        elif command.startswith('/load '):
            filename = command[6:].strip()
            if not filename:
                print(colored("Please specify a filename: /load <filename>", 'red'))
                return True

            try:
                # Add .txt extension if not present
                if not filename.endswith('.txt'):
                    filename += '.txt'

                filepath = os.path.join('conversations', filename)

                if not os.path.exists(filepath):
                    print(colored(f"File not found: {filepath}", 'red'))
                    return True

                with open(filepath, 'r', encoding='utf-8') as f:
                    content = f.read()

                # Parse conversation turns
                self.history = []
                self.history_tokens = []

                # Split by end of turn marker
                turns = content.split(self.end_of_turn)
                for turn in turns:
                    turn = turn.strip()
                    if not turn:
                        continue

                    if turn.startswith(self.human_prefix):
                        text = turn[len(self.human_prefix):].strip()
                        self.history.append(("human", text))
                    elif turn.startswith(self.assistant_prefix):
                        text = turn[len(self.assistant_prefix):].strip()
                        self.history.append(("assistant", text))

                # Recalculate token counts
                self.history_tokens = []
                for entry in self.history:
                    role, text = entry
                    if role == "human":
                        self.history_tokens.extend(self._tokenize(f"{self.human_prefix}{text}{self.end_of_turn}"))
                    else:
                        self.history_tokens.extend(self._tokenize(f"{self.assistant_prefix}{text}{self.end_of_turn}"))

                print(colored(f"Loaded conversation from {filepath} ({len(self.history) // 2} turns)", 'green'))

                # Print the conversation
                for i in range(0, len(self.history), 2):
                    if i < len(self.history):
                        user_text = self.history[i][1]
                        print(colored(f"\nYou: {user_text}", 'green'))

                    if i + 1 < len(self.history):
                        assistant_text = self.history[i + 1][1]
                        print(colored("CosmicFish: ", 'blue'), end="")
                        for line in assistant_text.split('\n'):
                            wrapped_lines = textwrap.wrap(line, width=100) if line.strip() else ['']
                            for wrapped_line in wrapped_lines:
                                print(wrapped_line)

            except Exception as e:
                print(colored(f"Error loading conversation: {str(e)}", 'red'))

            return True

        else:
            print(colored(f"Unknown command: {command}. Type /help for available commands.", 'red'))
            return True


def download_cosmicfish_from_hub(model_repo=DEFAULT_MODEL_REPO, device='cpu'):
    """Download and load CosmicFish model from Hugging Face Hub (safetensors only)"""
    print(colored(f"Downloading CosmicFish from Hugging Face: {model_repo}", "cyan"))

    try:
        # Download the model files to local cache
        print("Downloading model files...")
        cache_dir = snapshot_download(repo_id=model_repo, cache_dir=None)
        print(f"Model cached at: {cache_dir}")

        # Load config
        config_path = os.path.join(cache_dir, "config.json")
        with open(config_path, "r") as f:
            config_dict = json.load(f)

        # Create CosmicConfig
        config = CosmicConfig(
            vocab_size=config_dict["vocab_size"],
            block_size=config_dict["block_size"],
            n_layer=config_dict["n_layer"],
            n_head=config_dict["n_head"],
            n_embd=config_dict["n_embd"],
            bias=config_dict["bias"],
            dropout=0.0,  # Set to 0 for inference
            eps=config_dict.get("eps", 1e-6),
            use_rotary=config_dict["use_rotary"],
            use_swiglu=config_dict["use_swiglu"],
            use_gqa=config_dict["use_gqa"],
            n_query_groups=config_dict["n_query_groups"],
            use_qk_norm=config_dict.get("use_qk_norm", False)
        )

        # Create model
        print("Creating model...")
        model = CosmicFish(config)

        # Load weights from safetensors ONLY
        print("Loading weights from safetensors...")
        safetensors_path = os.path.join(cache_dir, "model.safetensors")

        if not os.path.exists(safetensors_path):
            raise FileNotFoundError(f"model.safetensors not found in {cache_dir}. This model requires safetensors format.")

        state_dict = load_file(safetensors_path)

        # Handle weight sharing: lm_head.weight shares with transformer.wte.weight
        if 'lm_head.weight' not in state_dict and 'transformer.wte.weight' in state_dict:
            state_dict['lm_head.weight'] = state_dict['transformer.wte.weight']

        model.load_state_dict(state_dict)
        model.to(device)
        model.eval()

        print(f"Model loaded: {model.get_num_params() / 1e6:.1f}M parameters")
        print(f"Device: {device}")
        return model, config

    except Exception as e:
        print(colored(f"Error downloading/loading model: {str(e)}", "red"))
        print(colored("Make sure you have internet connection and the model repo exists", "yellow"))
        sys.exit(1)


def load_tokenizer():
    tokenizer = GPT2Tokenizer.from_pretrained("gpt2")
    return tokenizer


def main():
    parser = argparse.ArgumentParser(description="Chat with CosmicFish")

    # Model parameters
    parser.add_argument("--model_repo", type=str, default=DEFAULT_MODEL_REPO,
                        help=f"Hugging Face model repository (default: {DEFAULT_MODEL_REPO})")
    parser.add_argument("--device", type=str, default="cuda" if torch.cuda.is_available() else "cpu",
                        help="Device to use (cuda or cpu)")

    # Generation parameters
    parser.add_argument("--temperature", type=float, default=0.5,
                        help="Temperature for sampling (default: 0.7)")
    parser.add_argument("--max_tokens", type=int, default=512,
                        help="Maximum number of tokens to generate per response")
    parser.add_argument("--min_tokens", type=int, default=10,
                        help="Minimum number of tokens to generate per response")
    parser.add_argument("--top_k", type=int, default=40,
                        help="Top-k sampling (0 to disable)")
    parser.add_argument("--repetition_penalty", type=float, default=1.2,
                        help="Repetition penalty (1.0 = no penalty, 1.2 = mild, 1.5 = moderate)")

    # Chat parameters
    parser.add_argument("--human_prefix", type=str, default="Human: ",
                        help="Prefix for human messages")
    parser.add_argument("--assistant_prefix", type=str, default="Assistant: ",
                        help="Prefix for assistant messages")
    parser.add_argument("--end_of_turn", type=str, default="\n\n",
                        help="Delimiter between conversation turns")
    parser.add_argument("--instruction", type=str,
                        default=DEFAULT_PROMPT_TEMPLATE,
                        help="Instruction prompt to prepend to the conversation")
    parser.add_argument("--max_history", type=int, default=512,
                        help="Maximum number of tokens to keep in history")

    # UI parameters
    parser.add_argument("--no_welcome", action="store_true",
                        help="Don't display the welcome message")
    parser.add_argument("--debug", action="store_true",
                        help="Enable debug mode")

    args = parser.parse_args()

    # Configure device
    device = args.device
    if device == "cuda" and not torch.cuda.is_available():
        print(colored("CUDA is not available, falling back to CPU", "yellow"))
        device = "cpu"

    try:
        # Download and load the model from HF Hub
        model, model_config = download_cosmicfish_from_hub(args.model_repo, device)

        # Load tokenizer
        tokenizer = load_tokenizer()

        # Create a config object with all the necessary parameters
        class ChatConfig:
            def __init__(self, args, block_size):
                self.device = device
                self.temperature = args.temperature
                self.max_new_tokens = args.max_tokens
                self.min_tokens_to_generate = args.min_tokens
                self.top_k = args.top_k
                self.human_prefix = args.human_prefix
                self.assistant_prefix = args.assistant_prefix
                self.end_of_turn = args.end_of_turn
                self.prompt_template = args.instruction
                self.max_history_tokens = args.max_history
                self.display_welcome = not args.no_welcome
                self.block_size = block_size
                self.debug_mode = args.debug
                self.repetition_penalty = args.repetition_penalty

        config = ChatConfig(args, model_config.block_size)

        # Initialize chat session
        chat = CosmicFishChatSession(model, tokenizer, config)

        # Main chat loop
        print(colored("\nCosmicFish initialized from Hugging Face! Type your message (or /help for commands).\n", 'cyan'))

        while True:
            try:
                # Get user input
                user_input = input(colored("You: ", 'green'))

                # Check if it's a command
                if user_input.startswith('/'):
                    # Execute command, continue loop if True, exit if False
                    if not chat.execute_command(user_input):
                        break
                    continue

                # Skip if empty input
                if not user_input.strip():
                    continue

                # Generate response using live generation
                live_buffer = ""
                final_response = None

                # Use the generator version
                response_generator = chat.generate_response(user_input)

                try:
                    # First print the assistant prefix
                    print(colored("CosmicFish: ", 'blue'), end="")
                    sys.stdout.flush()

                    for token, live_text, is_done in response_generator:
                        # If this is the final clean response
                        if is_done:
                            final_response = live_text
                            # Print the final response directly if we didn't get any tokens yet
                            if not live_buffer:
                                print(final_response, end="")
                            break
                        if token:
                            # Check if token contains <|endoftext|> and remove it if present
                            if "<|endoftext|>" in token:
                                token = token.replace("<|endoftext|>", "")
                                if token:  # Only print if there's anything left
                                    print(token, end="", flush=True)
                                break

                            # Display it
                            print(token, end="", flush=True)
                            live_buffer += token

                except KeyboardInterrupt:
                    # Allow user to interrupt generation
                    print("\n[Generation interrupted]")
                    final_response = "I was going to respond, but I'll stop here since you interrupted."

                # Add an extra line for readability
                print()

            except KeyboardInterrupt:
                print("\n\nKeyboard interrupt detected. Type /exit to quit or continue chatting.")

            except Exception as e:
                print(colored(f"\nError: {str(e)}", 'red'))
                logger.error(f"Error in chat loop: {str(e)}", exc_info=True)

    except Exception as e:
        print(colored(f"Error setting up chat: {str(e)}", 'red'))
        logger.error(f"Error setting up chat: {str(e)}", exc_info=True)
        sys.exit(1)


if __name__ == "__main__":
    try:
        main()
    except Exception as e:
        logger.error(f"Fatal error: {str(e)}", exc_info=True)
        sys.exit(1)