import gradio as gr
import tensorflow as tf
import numpy as np
import os
import json
from model import VedaProgrammingLLM, create_veda_model
from tokenizer import VedaTokenizer
from train import VedaTrainer

# Global variables for model and tokenizer
model = None
tokenizer = None
trainer = None

def load_or_create_model():
    """Load existing model or create new one"""
    global model, tokenizer, trainer
    
    model_path = "veda_model"
    
    if os.path.exists(os.path.join(model_path, "config.json")):
        print("Loading existing model...")
        
        # Load config
        with open(os.path.join(model_path, "config.json"), 'r') as f:
            config = json.load(f)
        
        # Load tokenizer
        tokenizer = VedaTokenizer()
        tokenizer.load(os.path.join(model_path, "tokenizer.json"))
        
        # Create model
        model = create_veda_model(
            vocab_size=tokenizer.vocabulary_size,
            max_length=config.get('max_length', 256),
            model_size="small"
        )
        
        # Build model
        dummy_input = tf.zeros((1, config.get('max_length', 256)), dtype=tf.int32)
        model(dummy_input)
        
        # Load weights
        model.load_weights(os.path.join(model_path, "model_weights.h5"))
        print("Model loaded successfully!")
    else:
        print("No existing model found. Creating and training new model...")
        trainer = VedaTrainer(
            data_path="programming.txt",
            vocab_size=5000,
            max_length=256,
            batch_size=8,
            model_size="small"
        )
        trainer.train(epochs=5, save_path=model_path)
        model = trainer.model
        tokenizer = trainer.tokenizer
        print("Model trained and ready!")

def generate_code(
    prompt: str,
    max_tokens: int,
    temperature: float,
    top_k: int,
    top_p: float
) -> str:
    """Generate code based on prompt"""
    global model, tokenizer
    
    if model is None or tokenizer is None:
        return "Error: Model not loaded. Please wait for initialization."
    
    try:
        # Encode prompt
        prompt_tokens = tokenizer.encode(prompt)
        
        if len(prompt_tokens) == 0:
            return "Error: Could not tokenize prompt."
        
        # Generate
        generated_tokens = model.generate(
            prompt_tokens,
            max_new_tokens=int(max_tokens),
            temperature=float(temperature),
            top_k=int(top_k),
            top_p=float(top_p)
        )
        
        # Decode
        generated_text = tokenizer.decode(generated_tokens)
        
        return generated_text
    
    except Exception as e:
        return f"Error generating code: {str(e)}"

def train_model(
    training_data: str,
    epochs: int,
    learning_rate: float
) -> str:
    """Train or fine-tune the model"""
    global model, tokenizer, trainer
    
    try:
        # Save training data
        with open("programming.txt", 'w', encoding='utf-8') as f:
            f.write(training_data)
        
        # Create trainer
        trainer = VedaTrainer(
            data_path="programming.txt",
            vocab_size=5000,
            max_length=256,
            batch_size=8,
            model_size="small"
        )
        
        # Train
        history = trainer.train(epochs=int(epochs), save_path="veda_model")
        
        # Update global references
        model = trainer.model
        tokenizer = trainer.tokenizer
        
        # Get final metrics
        final_loss = history.history['loss'][-1]
        final_acc = history.history.get('accuracy', [0])[-1]
        
        return f"""Training completed successfully!
        
Final Loss: {final_loss:.4f}
Final Accuracy: {final_acc:.4f}
Epochs: {epochs}
Vocabulary Size: {tokenizer.vocabulary_size}
"""
    
    except Exception as e:
        return f"Error during training: {str(e)}"

def get_model_info() -> str:
    """Get information about the current model"""
    global model, tokenizer
    
    if model is None:
        return "No model loaded."
    
    config = model.get_config()
    
    info = f"""
**Veda Programming LLM**

- Vocabulary Size: {config['vocab_size']}
- Max Sequence Length: {config['max_length']}
- Model Dimension: {config['d_model']}
- Number of Heads: {config['num_heads']}
- Number of Layers: {config['num_layers']}
- Feed Forward Dimension: {config['ff_dim']}

Total Parameters: {model.count_params():,}
"""
    return info

# Create Gradio interface
def create_interface():
    """Create the Gradio interface"""
    
    with gr.Blocks(
        title="Veda Programming LLM",
        theme=gr.themes.Soft()
    ) as demo:
        
        gr.Markdown("""
        # 🕉️ Veda Programming LLM
        
        A TensorFlow-based language model for code generation.
        
        Train on your own code or use the pre-trained model to generate programming code.
        """)
        
        with gr.Tabs():
            # Generation Tab
            with gr.TabItem("💻 Code Generation"):
                with gr.Row():
                    with gr.Column(scale=1):
                        prompt_input = gr.Textbox(
                            label="Code Prompt",
                            placeholder="def fibonacci(",
                            lines=5,
                            value="def calculate_sum("
                        )
                        
                        with gr.Row():
                            max_tokens = gr.Slider(
                                minimum=10,
                                maximum=500,
                                value=100,
                                step=10,
                                label="Max New Tokens"
                            )
                            temperature = gr.Slider(
                                minimum=0.1,
                                maximum=2.0,
                                value=0.7,
                                step=0.1,
                                label="Temperature"
                            )
                        
                        with gr.Row():
                            top_k = gr.Slider(
                                minimum=1,
                                maximum=100,
                                value=50,
                                step=1,
                                label="Top-K"
                            )
                            top_p = gr.Slider(
                                minimum=0.1,
                                maximum=1.0,
                                value=0.9,
                                step=0.05,
                                label="Top-P (Nucleus)"
                            )
                        
                        generate_btn = gr.Button(
                            "🚀 Generate Code",
                            variant="primary"
                        )
                    
                    with gr.Column(scale=1):
                        output_code = gr.Code(
                            label="Generated Code",
                            language="python",
                            lines=15
                        )
                
                generate_btn.click(
                    fn=generate_code,
                    inputs=[prompt_input, max_tokens, temperature, top_k, top_p],
                    outputs=output_code
                )
            
            # Training Tab
            with gr.TabItem("🎓 Training"):
                gr.Markdown("""
                ### Train the Model
                
                Paste your programming code below to train or fine-tune the model.
                """)
                
                training_data = gr.Textbox(
                    label="Training Data (Python Code)",
                    placeholder="Paste your Python code here...",
                    lines=15,
                    value='''def hello_world():
    print("Hello, World!")
    return True

def add_numbers(a, b):
    return a + b

class MyClass:
    def __init__(self, value):
        self.value = value
    
    def get_value(self):
        return self.value
'''
                )
                
                with gr.Row():
                    epochs_input = gr.Slider(
                        minimum=1,
                        maximum=50,
                        value=10,
                        step=1,
                        label="Epochs"
                    )
                    lr_input = gr.Slider(
                        minimum=0.00001,
                        maximum=0.01,
                        value=0.0001,
                        step=0.00001,
                        label="Learning Rate"
                    )
                
                train_btn = gr.Button(
                    "🎯 Start Training",
                    variant="primary"
                )
                
                training_output = gr.Textbox(
                    label="Training Output",
                    lines=10
                )
                
                train_btn.click(
                    fn=train_model,
                    inputs=[training_data, epochs_input, lr_input],
                    outputs=training_output
                )
            
            # Model Info Tab
            with gr.TabItem("ℹ️ Model Info"):
                model_info = gr.Markdown()
                refresh_btn = gr.Button("🔄 Refresh Info")
                refresh_btn.click(
                    fn=get_model_info,
                    outputs=model_info
                )
            
            # Examples Tab
            with gr.TabItem("📝 Examples"):
                gr.Markdown("""
                ### Example Prompts
                
                Try these prompts for code generation:
                
                1. `def fibonacci(`
                2. `class Calculator:`
                3. `def binary_search(arr, target):`
                4. `async def fetch_data(`
                5. `def quicksort(arr):`
                6. `class LinkedList:`
                """)
                
                gr.Examples(
                    examples=[
                        ["def fibonacci(", 100, 0.7, 50, 0.9],
                        ["class Calculator:", 150, 0.8, 40, 0.95],
                        ["def binary_search(arr, target):", 100, 0.6, 50, 0.9],
                        ["def bubble_sort(", 100, 0.7, 50, 0.9],
                    ],
                    inputs=[prompt_input, max_tokens, temperature, top_k, top_p],
                    outputs=output_code,
                    fn=generate_code
                )
        
        gr.Markdown("""
        ---
        **Veda Programming LLM** - Built with TensorFlow & Gradio
        
        *Note: This is a demonstration model. For production use, train on larger datasets.*
        """)
    
    return demo

# Main execution
if __name__ == "__main__":
    print("Initializing Veda Programming LLM...")
    load_or_create_model()
    
    print("Starting Gradio interface...")
    demo = create_interface()
    demo.launch(
        server_name="0.0.0.0",
        server_port=7860,
        share=False
    )