import gradio as gr
from transformers import AutoModelForCausalLM, AutoTokenizer, PreTrainedModel, PretrainedConfig
import torch
import torch.nn as nn
import torch.nn.functional as F
import math
import os
import sys
import transformers
from dotenv import load_dotenv
from huggingface_hub import login

class RMSNorm(nn.Module):
    def __init__(self, hidden_size, eps=1e-5):
        super().__init__()
        self.weight = nn.Parameter(torch.ones(hidden_size))
        self.eps = eps

    def forward(self, x):
        variance = x.pow(2).mean(-1, keepdim=True)
        x = x * torch.rsqrt(variance + self.eps)
        return self.weight * x

class LlamaAttention(nn.Module):
    def __init__(self, config):
        super().__init__()
        self.hidden_size = config.hidden_size
        self.num_heads = config.num_attention_heads
        self.num_kv_heads = config.num_key_value_heads
        self.head_dim = config.hidden_size // config.num_attention_heads
        
        self.q_proj = nn.Linear(config.hidden_size, self.num_heads * self.head_dim, bias=False)
        self.k_proj = nn.Linear(config.hidden_size, self.num_kv_heads * self.head_dim, bias=False)
        self.v_proj = nn.Linear(config.hidden_size, self.num_kv_heads * self.head_dim, bias=False)
        self.o_proj = nn.Linear(self.num_heads * self.head_dim, config.hidden_size, bias=False)

    def forward(self, hidden_states, attention_mask=None):
        batch_size, seq_length, _ = hidden_states.size()
        
        q = self.q_proj(hidden_states).view(batch_size, seq_length, self.num_heads, self.head_dim)
        k = self.k_proj(hidden_states).view(batch_size, seq_length, self.num_kv_heads, self.head_dim)
        v = self.v_proj(hidden_states).view(batch_size, seq_length, self.num_kv_heads, self.head_dim)
        
        if self.num_kv_heads < self.num_heads:
            k = k.repeat_interleave(self.num_heads // self.num_kv_heads, dim=2)
            v = v.repeat_interleave(self.num_heads // self.num_kv_heads, dim=2)
        
        q = q.transpose(1, 2)
        k = k.transpose(1, 2)
        v = v.transpose(1, 2)
        
        attention_scores = torch.matmul(q, k.transpose(-2, -1)) / math.sqrt(self.head_dim)
        
        if attention_mask is not None:
            attention_scores = attention_scores + attention_mask
            
        attention_probs = F.softmax(attention_scores, dim=-1)
        context = torch.matmul(attention_probs, v)
        
        context = context.transpose(1, 2).contiguous()
        context = context.view(batch_size, seq_length, -1)
        
        return self.o_proj(context)

class LlamaMLP(nn.Module):
    def __init__(self, config):
        super().__init__()
        self.gate_proj = nn.Linear(config.hidden_size, config.intermediate_size, bias=False)
        self.up_proj = nn.Linear(config.hidden_size, config.intermediate_size, bias=False)
        self.down_proj = nn.Linear(config.intermediate_size, config.hidden_size, bias=False)
        self.act_fn = nn.SiLU()

    def forward(self, x):
        gate = self.act_fn(self.gate_proj(x))
        up = self.up_proj(x)
        return self.down_proj(gate * up)

class LlamaDecoderLayer(nn.Module):
    def __init__(self, config):
        super().__init__()
        self.self_attn = LlamaAttention(config)
        self.mlp = LlamaMLP(config)
        self.input_layernorm = RMSNorm(config.hidden_size, config.rms_norm_eps)
        self.post_attention_layernorm = RMSNorm(config.hidden_size, config.rms_norm_eps)

    def forward(self, hidden_states, attention_mask=None):
        residual = hidden_states
        hidden_states = self.input_layernorm(hidden_states)
        hidden_states = self.self_attn(hidden_states, attention_mask)
        hidden_states = residual + hidden_states
        
        residual = hidden_states
        hidden_states = self.post_attention_layernorm(hidden_states)
        hidden_states = self.mlp(hidden_states)
        hidden_states = residual + hidden_states
        
        return hidden_states

class SmolLM2Config(PretrainedConfig):
    model_type = "smollm2"
    
    def __init__(
        self,
        vocab_size=49152,
        hidden_size=576,
        intermediate_size=1536,
        num_hidden_layers=30,
        num_attention_heads=9,
        num_key_value_heads=3,
        hidden_act="silu",
        max_position_embeddings=2048,
        initializer_range=0.041666666666666664,
        rms_norm_eps=1e-5,
        use_cache=True,
        pad_token_id=None,
        bos_token_id=0,
        eos_token_id=0,
        tie_word_embeddings=True,
        rope_theta=10000.0,
        **kwargs
    ):
        self.vocab_size = vocab_size
        self.hidden_size = hidden_size
        self.intermediate_size = intermediate_size
        self.num_hidden_layers = num_hidden_layers
        self.num_attention_heads = num_attention_heads
        self.num_key_value_heads = num_key_value_heads
        self.hidden_act = hidden_act
        self.max_position_embeddings = max_position_embeddings
        self.initializer_range = initializer_range
        self.rms_norm_eps = rms_norm_eps
        self.use_cache = use_cache
        self.rope_theta = rope_theta
        super().__init__(
            pad_token_id=pad_token_id,
            bos_token_id=bos_token_id,
            eos_token_id=eos_token_id,
            tie_word_embeddings=tie_word_embeddings,
            **kwargs
        )

class SmolLM2ForCausalLM(PreTrainedModel):
    config_class = SmolLM2Config
    _no_split_modules = ["LlamaDecoderLayer"]
    
    def __init__(self, config):
        super().__init__(config)
        self.config = config
        
        self.embed_tokens = nn.Embedding(config.vocab_size, config.hidden_size)
        self.layers = nn.ModuleList([LlamaDecoderLayer(config) for _ in range(config.num_hidden_layers)])
        self.norm = RMSNorm(config.hidden_size, config.rms_norm_eps)
        self.lm_head = nn.Linear(config.hidden_size, config.vocab_size, bias=False)
        
        if config.tie_word_embeddings:
            self.lm_head.weight = self.embed_tokens.weight
            
    def forward(self, input_ids, attention_mask=None, labels=None):
        hidden_states = self.embed_tokens(input_ids)
        
        # Create causal attention mask if none provided
        if attention_mask is None:
            attention_mask = torch.triu(
                torch.ones((input_ids.size(1), input_ids.size(1)), dtype=torch.bool, device=input_ids.device),
                diagonal=1
            )
            attention_mask = attention_mask.unsqueeze(0).unsqueeze(0)
            attention_mask = attention_mask * -1e4
        
        for layer in self.layers:
            hidden_states = layer(hidden_states, attention_mask)
            
        hidden_states = self.norm(hidden_states)
        logits = self.lm_head(hidden_states)
        
        loss = None
        if labels is not None:
            loss = F.cross_entropy(logits.view(-1, logits.size(-1)), labels.view(-1))
            
        return logits if loss is None else (loss, logits)

    def prepare_inputs_for_generation(self, input_ids, **kwargs):
        return {
            "input_ids": input_ids,
            "attention_mask": kwargs.get("attention_mask", None)
        }

# Register the model architecture
from transformers import AutoConfig, AutoModelForCausalLM
AutoConfig.register("smollm2", SmolLM2Config)
AutoModelForCausalLM.register(SmolLM2Config, SmolLM2ForCausalLM)

# Load model and tokenizer
model_id = "jatingocodeo/SmolLM2"

def check_huggingface_access():
    try:
        from huggingface_hub import HfApi
        api = HfApi()
        
        # Check if token exists
        try:
            token = os.getenv("HF_TOKEN")
            if not token:
                print("× Warning: HF_TOKEN environment variable not found")
                print("  Please ensure you have set your HuggingFace token")
                return False
        except Exception as e:
            print(f"× Error checking HF_TOKEN: {str(e)}")
            return False
        
        # Check repository access
        try:
            print(f"Checking access to repository: {model_id}")
            repo_info = api.repo_info(model_id, token=token)
            print("✓ Repository access confirmed")
            print(f"Repository info: {repo_info}")
            return True
        except Exception as e:
            print(f"× Error accessing repository: {str(e)}")
            return False
    except Exception as e:
        print(f"× Error in HuggingFace access check: {str(e)}")
        return False

def load_model():
    try:
        print("\n=== Starting model loading process ===")
        print(f"Model ID: {model_id}")
        
        # Load environment variables
        load_dotenv()
        
        # Get HF token
        hf_token = os.getenv('HF_TOKEN')
        if not hf_token:
            raise ValueError("HF_TOKEN not found in environment variables")
            
        # Login to Hugging Face
        login(hf_token)
        
        print("\n1. Loading tokenizer...")
        try:
            tokenizer = AutoTokenizer.from_pretrained(
                model_id,
                use_auth_token=hf_token,
                trust_remote_code=True
            )
            print("✓ Tokenizer loaded successfully")
            print(f"Tokenizer type: {type(tokenizer)}")
            print(f"Vocabulary size: {len(tokenizer)}")
        except Exception as e:
            print(f"× Error loading tokenizer: {str(e)}")
            raise
        
        print("\n2. Adding special tokens...")
        try:
            special_tokens = {
                'pad_token': '[PAD]',
                'eos_token': '</s>',
                'bos_token': '<s>'
            }
            num_added = tokenizer.add_special_tokens(special_tokens)
            print(f"✓ Added {num_added} special tokens")
            print(f"Special tokens: {tokenizer.special_tokens_map}")
        except Exception as e:
            print(f"× Error adding special tokens: {str(e)}")
            raise
        
        print("\n3. Loading model...")
        try:
            model = AutoModelForCausalLM.from_pretrained(
                model_id,
                use_auth_token=hf_token,
                trust_remote_code=True,
                torch_dtype=torch.float16 if torch.cuda.is_available() else torch.float32,
                low_cpu_mem_usage=True
            )
            print("✓ Model loaded successfully")
            print(f"Model type: {type(model)}")
        except Exception as e:
            print(f"× Error loading model: {str(e)}")
            raise
        
        print("\n4. Moving model to device...")
        try:
            device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
            print(f"Selected device: {device}")
            model = model.to(device)
            print("✓ Model moved to device successfully")
        except Exception as e:
            print(f"× Error moving model to device: {str(e)}")
            raise
        
        print("\n=== Model loading completed successfully! ===")
        return model, tokenizer
        
    except Exception as e:
        print("\n!!! ERROR IN MODEL LOADING !!!")
        print(f"Error type: {type(e).__name__}")
        print(f"Error message: {str(e)}")
        print("\nFull traceback:")
        import traceback
        traceback.print_exc()
        print("\nAdditional debug info:")
        print(f"Python version: {sys.version}")
        print(f"PyTorch version: {torch.__version__}")
        print(f"Transformers version: {transformers.__version__}")
        print(f"CUDA available: {torch.cuda.is_available()}")
        if torch.cuda.is_available():
            print(f"CUDA version: {torch.version.cuda}")
        print("\nEnvironment variables:")
        print(f"HF_TOKEN set: {'HF_TOKEN' in os.environ}")
        raise

def generate_text(prompt, max_length=100, temperature=0.7, top_k=50):
    try:
        print("\n=== Starting text generation ===")
        print(f"Input prompt: {prompt}")
        print(f"Parameters: max_length={max_length}, temperature={temperature}, top_k={top_k}")
        
        if not hasattr(generate_text, "model"):
            print("\n1. First call - loading model...")
            generate_text.model, generate_text.tokenizer = load_model()
        
        if not prompt.strip():
            print("× Empty prompt received")
            return "Please enter a prompt."
        
        print("\n2. Processing prompt...")
        if not prompt.startswith(generate_text.tokenizer.bos_token):
            prompt = generate_text.tokenizer.bos_token + prompt
            print("Added BOS token to prompt")
        
        print("\n3. Encoding prompt...")
        try:
            input_ids = generate_text.tokenizer.encode(prompt, return_tensors="pt", truncation=True, max_length=2048)
            print(f"Encoded shape: {input_ids.shape}")
            input_ids = input_ids.to(generate_text.model.device)
            print("✓ Encoding successful")
        except Exception as e:
            print(f"× Error encoding prompt: {str(e)}")
            raise
        
        print("\n4. Generating text...")
        try:
            with torch.no_grad():
                output_ids = generate_text.model.generate(
                    input_ids,
                    max_length=min(max_length + len(input_ids[0]), 2048),
                    temperature=temperature,
                    top_k=top_k,
                    do_sample=True,
                    pad_token_id=generate_text.tokenizer.pad_token_id,
                    eos_token_id=generate_text.tokenizer.eos_token_id,
                    num_return_sequences=1
                )
            print(f"Generation shape: {output_ids.shape}")
        except Exception as e:
            print(f"× Error during generation: {str(e)}")
            raise
        
        print("\n5. Decoding output...")
        try:
            generated_text = generate_text.tokenizer.decode(output_ids[0], skip_special_tokens=True)
            print("✓ Decoding successful")
            print(f"Output length: {len(generated_text)}")
        except Exception as e:
            print(f"× Error decoding output: {str(e)}")
            raise
        
        print("\n=== Generation completed successfully! ===")
        return generated_text.strip()
    
    except Exception as e:
        print("\n!!! ERROR IN TEXT GENERATION !!!")
        print(f"Error type: {type(e).__name__}")
        print(f"Error message: {str(e)}")
        print("\nFull traceback:")
        import traceback
        traceback.print_exc()
        return f"An error occurred: {str(e)}"

# Create Gradio interface
iface = gr.Interface(
    fn=generate_text,
    inputs=[
        gr.Textbox(label="Prompt", placeholder="Enter your prompt here...", lines=2),
        gr.Slider(minimum=10, maximum=200, value=100, step=1, label="Max Length"),
        gr.Slider(minimum=0.1, maximum=1.0, value=0.7, step=0.1, label="Temperature"),
        gr.Slider(minimum=1, maximum=100, value=50, step=1, label="Top K"),
    ],
    outputs=gr.Textbox(label="Generated Text", lines=5),
    title="SmolLM2 Text Generator",
    description="""Generate text using the fine-tuned SmolLM2 model.
    - Max Length: Controls the length of generated text
    - Temperature: Controls randomness (higher = more creative)
    - Top K: Controls diversity of word choices""",
    examples=[
        ["Once upon a time", 100, 0.7, 50],
        ["The quick brown fox", 150, 0.8, 40],
        ["In a galaxy far far away", 200, 0.9, 30],
    ],
    allow_flagging="never"
)

if __name__ == "__main__":
    print("Starting Gradio interface...")
    iface.launch(debug=True)