import gradio as gr
import torch
from transformers import AutoModelForCausalLM, AutoTokenizer
import os

# Model configuration
# Since you have all model files in Space root, try loading directly
MODEL_NAME = "."  # Load from current directory with all your uploaded files
CUSTOM_WEIGHTS_PATH = "./model.safetensors"  # Backup: your custom weights

DEVICE = "cuda" if torch.cuda.is_available() else "cpu"

# Global variables for model caching
_model = None
_tokenizer = None

def load_model():
    """Load the model and tokenizer with simple caching"""
    global _model, _tokenizer
    
    # Return cached model if already loaded
    if _model is not None and _tokenizer is not None:
        return _model, _tokenizer
    
    print(f"Loading model from: {MODEL_NAME}")
    print(f"Using device: {DEVICE}")
    
    # List available files for debugging
    import os
    try:
        current_files = os.listdir(".")
        print("Available files in current directory:")
        for f in current_files:
            print(f"  - {f}")
    except Exception as e:
        print(f"Could not list directory: {e}")
    
    try:
        # First try to load directly from your uploaded files
        print("Attempting to load model directly from uploaded files...")
        try:
            tokenizer = AutoTokenizer.from_pretrained(MODEL_NAME)
            model = AutoModelForCausalLM.from_pretrained(MODEL_NAME)
            print("✅ Successfully loaded model directly from your uploaded files!")
        except Exception as direct_load_error:
            print(f"Direct load failed: {direct_load_error}")
            print("Falling back to base model + custom weights...")
            
            # Fallback: Load base model and add custom weights
            tokenizer = AutoTokenizer.from_pretrained("gpt2")
            model = AutoModelForCausalLM.from_pretrained("gpt2")
            
            # Try to load your custom weights
            if os.path.exists(CUSTOM_WEIGHTS_PATH):
                print(f"Loading custom weights from: {CUSTOM_WEIGHTS_PATH}")
                try:
                    from safetensors.torch import load_file
                    custom_weights = load_file(CUSTOM_WEIGHTS_PATH)
                    
                    # Load the weights into the model
                    missing_keys, unexpected_keys = model.load_state_dict(custom_weights, strict=False)
                    
                    if missing_keys:
                        print(f"⚠️ Missing keys: {len(missing_keys)} (this might be normal for LoRA models)")
                    if unexpected_keys:
                        print(f"⚠️ Unexpected keys: {len(unexpected_keys)}")
                    
                    print("✅ Custom weights loaded successfully!")
                except Exception as e:
                    print(f"⚠️ Could not load custom weights: {e}")
                    print("Using base GPT-2 model instead")
        
        # Set pad token if not set
        if tokenizer.pad_token is None:
            tokenizer.pad_token = tokenizer.eos_token
        
        # Move model to device
        model = model.to(DEVICE)
        
        print(f"✅ Model loaded successfully on {DEVICE}!")
        
        # Cache the loaded model and tokenizer
        _model = model
        _tokenizer = tokenizer
        
        return model, tokenizer
    
    except Exception as e:
        print(f"❌ Error loading model: {e}")
        print("\n🔧 Troubleshooting:")
        print("1. Make sure you have uploaded ALL required files:")
        print("   - model.safetensors (✅ you have this)")
        print("   - config.json (❓ might be missing)")
        print("   - tokenizer.json or vocab.json + merges.txt (❓ might be missing)")
        print("   - tokenizer_config.json (✅ you have this)")
        print("2. Files should be in the Space root directory")
        print("3. Check if the model was saved correctly from your notebook")
        raise e

# Initialize model and tokenizer
try:
    model, tokenizer = load_model()
except Exception as e:
    print(f"Failed to load model: {e}")
    # Create dummy objects to prevent further errors
    model, tokenizer = None, None

def generate_code(pseudocode, indent=1, line=1, temperature=0.7, top_p=0.9, max_length=128):
    """
    Generate code from pseudo-code with line and indent information.
    
    Args:
        pseudocode: Input pseudo-code string
        indent: Indentation level (1-10)
        line: Line number (1-100)
        temperature: Sampling temperature (0.1-2.0)
        top_p: Nucleus sampling parameter (0.1-1.0)
        max_length: Maximum length of generated sequence (50-512)
    
    Returns:
        Generated code string
    """
    try:
        # Check if model is loaded
        if model is None or tokenizer is None:
            return """❌ Model not loaded. Please check:
            
1. MODEL_NAME in app.py - should be either:
   - Your HF repository: "username/model-name"
   - Local path: "./model" (if files uploaded to Space)

2. If using HF repository, make sure it exists and is public

3. If using local files, ensure model files are in correct folder

Current MODEL_NAME: """ + MODEL_NAME
        
        # Validate inputs
        if not pseudocode.strip():
            return "❌ Error: Please enter some pseudocode."
        
        # Format input with line and indent information (matches training format)
        prompt = f"Pseudocode: {pseudocode.strip()} | Indent: {indent} | Line: {line}\nCode:"
        
        # Tokenize input
        inputs = tokenizer(prompt, return_tensors='pt', padding=True, truncation=True, max_length=256)
        inputs = {k: v.to(DEVICE) for k, v in inputs.items()}
        
        # Generate with the model
        model.eval()
        with torch.no_grad():
            outputs = model.generate(
                **inputs,
                max_new_tokens=max_length,
                temperature=max(0.1, temperature),  # Ensure minimum temperature
                top_p=max(0.1, top_p),  # Ensure minimum top_p
                do_sample=True,
                pad_token_id=tokenizer.eos_token_id,
                eos_token_id=tokenizer.eos_token_id,
                num_return_sequences=1,
                repetition_penalty=1.1,
                no_repeat_ngram_size=2
            )
        
        # Decode output
        generated_text = tokenizer.decode(outputs[0], skip_special_tokens=True)
        
        # Extract only the code part (remove the prompt)
        if "Code:" in generated_text:
            code = generated_text.split("Code:")[-1].strip()
        else:
            code = generated_text.strip()
        
        # Clean up the output
        if code.startswith(prompt):
            code = code[len(prompt):].strip()
        
        return code if code else "❌ No code generated. Try adjusting the parameters."
        
    except Exception as e:
        return f"❌ Error generating code: {str(e)}"

def create_examples():
    """Create example inputs for the interface"""
    return [
        ["create string s", 1, 1, 0.7, 0.9, 100],
        ["read input from user", 1, 2, 0.7, 0.9, 100],
        ["if s is empty", 1, 3, 0.7, 0.9, 100],
        ["print hello world", 2, 4, 0.7, 0.9, 100],
        ["for i from 0 to n", 1, 5, 0.7, 0.9, 100],
        ["declare integer array", 1, 1, 0.5, 0.9, 80],
        ["while condition is true", 2, 10, 0.8, 0.95, 120]
    ]

# Create Gradio interface
with gr.Blocks(
    theme=gr.themes.Soft(),
    title="🐍 Pseudo-Code to Code Generator",
    css="""
    .gradio-container {
        max-width: 1200px;
        margin: auto;
    }
    .header {
        text-align: center;
        margin-bottom: 30px;
    }
    .info-box {
        background-color: #f0f8f;
        color: #FF0000;
        padding: 15px;
        border-radius: 10px;
        margin: 10px 0;
        border: 1px solid #ddd;
    }
    """
) as demo:
    
    # Header
    gr.HTML("""
    <div class="header">
        <h1>🐍 Pseudo-Code to Code Generator</h1>
        <p>Convert natural language pseudo-code to executable code using fine-tuned GPT-2</p>
    </div>
    """)
    
    # Info box
    gr.HTML("""
    <div class="info-box">
        <h3>📋 How to use:</h3>
        <ol>
            <li><strong>Enter pseudocode:</strong> Describe what you want the code to do in natural language</li>
            <li><strong>Set context:</strong> Adjust indent level and line number for better structure</li>
            <li><strong>Tune generation:</strong> Modify temperature and top_p for different creativity levels</li>
            <li><strong>Generate:</strong> Click submit to get your code!</li>
        </ol>
        <p><strong>Note:</strong> This model was trained on the SPOC dataset containing C++ code examples.</p>
    </div>
    """)
    
    with gr.Row():
        # Left column - Inputs
        with gr.Column(scale=1):
            pseudocode_input = gr.Textbox(
                label="📝 Pseudocode",
                placeholder="Enter your pseudocode here... (e.g., 'create string variable s')",
                lines=3,
                value="create string s"
            )
            
            with gr.Row():
                indent_input = gr.Slider(
                    minimum=1, maximum=10, value=1, step=1,
                    label="🔢 Indent Level",
                    info="Indentation level for the code"
                )
                line_input = gr.Slider(
                    minimum=1, maximum=100, value=1, step=1,
                    label="📍 Line Number",
                    info="Line number in the program"
                )
            
            gr.Markdown("### 🎛️ Generation Parameters")
            
            with gr.Row():
                temperature_input = gr.Slider(
                    minimum=0.1, maximum=2.0, value=0.7, step=0.1,
                    label="🌡️ Temperature",
                    info="Higher = more creative, Lower = more focused"
                )
                top_p_input = gr.Slider(
                    minimum=0.1, maximum=1.0, value=0.9, step=0.05,
                    label="🎯 Top-p",
                    info="Nucleus sampling parameter"
                )
            
            max_length_input = gr.Slider(
                minimum=50, maximum=512, value=128, step=10,
                label="📏 Max Length",
                info="Maximum number of tokens to generate"
            )
            
            generate_btn = gr.Button("🚀 Generate Code", variant="primary", size="lg")
        
        # Right column - Output
        with gr.Column(scale=1):
            output = gr.Textbox(
                label="💻 Generated Code",
                lines=15,
                placeholder="Generated code will appear here...",
                show_copy_button=True
            )
    
    # Examples section
    gr.Markdown("### 📚 Example Inputs")
    examples = gr.Examples(
        examples=create_examples(),
        inputs=[pseudocode_input, indent_input, line_input, temperature_input, top_p_input, max_length_input],
        outputs=output,
        fn=generate_code,
        cache_examples=False
    )
    
    # Event handlers
    generate_btn.click(
        fn=generate_code,
        inputs=[pseudocode_input, indent_input, line_input, temperature_input, top_p_input, max_length_input],
        outputs=output
    )
    
    # Also allow Enter key to generate
    pseudocode_input.submit(
        fn=generate_code,
        inputs=[pseudocode_input, indent_input, line_input, temperature_input, top_p_input, max_length_input],
        outputs=output
    )
    
    # Footer
    gr.HTML("""
    <div style="text-align: center; margin-top: 30px; padding: 20px; border-top: 1px solid #eee;">
        <p>🤖 <strong>Model Details:</strong> Fine-tuned GPT-2 with LoRA on SPOC dataset</p>
        <p>📊 <strong>Training:</strong> Pseudo-code to C++ code generation with structural information</p>
        <p>⚡ <strong>Powered by:</strong> Transformers, Safetensors, and Gradio</p>
    </div>
    """)

# Launch configuration
if __name__ == "__main__":
    demo.launch(
        server_name="0.0.0.0",  # Required for Hugging Face Spaces
        server_port=7860,       # Default port for Spaces
        share=False,            # Don't create public links in Spaces
        show_api=False,         # Disable API docs for cleaner interface
        show_error=True,        # Show errors for debugging
        quiet=False             # Show startup logs
    )