#!/usr/bin/env python3
"""Test script to validate zip generation functionality"""

import json
import zipfile
import io
import os
from datetime import datetime

def create_requirements():
    """Generate requirements.txt"""
    return "gradio==4.16.0\nrequests==2.31.0"

def create_readme(config):
    """Generate README with deployment instructions"""
    return f"""# {config['name']}

{config['description']}

## 🚀 Quick Deploy to HuggingFace Spaces

### Step 1: Create the Space
1. Go to https://huggingface.co/spaces
2. Click "Create new Space"
3. Choose a name for your assistant
4. Select **Gradio** as the SDK
5. Set visibility (Public/Private)
6. Click "Create Space"

### Step 2: Upload Files
1. In your new Space, click "Files" tab
2. Upload these files from the zip:
   - `app.py`
   - `requirements.txt`
3. Wait for "Building" to complete

### Step 3: Add API Key
1. Go to Settings (⚙️ icon) 
2. Click "Variables and secrets"
3. Click "New secret"
4. Name: `{config['api_key_var']}`
5. Value: Your OpenRouter API key
6. Click "Add"

### Step 4: Get Your API Key
1. Go to https://openrouter.ai/keys
2. Sign up/login if needed
3. Click "Create Key"
4. Copy the key (starts with `sk-or-`)

### Step 5: Test Your Assistant
- Go back to "App" tab
- Your assistant should be running!
- Try the example prompts or ask a question

## 📝 Configuration

- **Model**: {config['model']}
- **Temperature**: {config['temperature']}
- **Max Tokens**: {config['max_tokens']}
- **API Key Variable**: {config['api_key_var']}

Generated on {datetime.now().strftime('%Y-%m-%d %H:%M:%S')} with HF Assistant Converter
"""

# Fixed template based on mvp_simple.py
ASSISTANT_TEMPLATE = '''import gradio as gr
import os
import requests
import json

# Configuration
ASSISTANT_NAME = "{name}"
ASSISTANT_DESCRIPTION = "{description}"
SYSTEM_PROMPT = """{system_prompt}"""
MODEL = "{model}"

# Get API key from environment - customizable variable name
API_KEY = os.environ.get("{api_key_var}")

def generate_response(message, history):
    """Generate response using OpenRouter API"""
    
    if not API_KEY:
        return "Please set your {api_key_var} in the Space settings."
    
    # Build messages array for the API
    messages = [{{"role": "system", "content": SYSTEM_PROMPT}}]
    
    # Add conversation history
    for user_msg, assistant_msg in history:
        messages.append({{"role": "user", "content": user_msg}})
        if assistant_msg:
            messages.append({{"role": "assistant", "content": assistant_msg}})
    
    # Add current message
    messages.append({{"role": "user", "content": message}})
    
    # Make API request
    try:
        response = requests.post(
            url="https://openrouter.ai/api/v1/chat/completions",
            headers={{
                "Authorization": f"Bearer {{API_KEY}}",
                "Content-Type": "application/json"
            }},
            json={{
                "model": MODEL,
                "messages": messages,
                "temperature": {temperature},
                "max_tokens": {max_tokens}
            }}
        )
        
        if response.status_code == 200:
            return response.json()['choices'][0]['message']['content']
        else:
            return f"Error: {{response.status_code}} - {{response.text}}"
            
    except Exception as e:
        return f"Error: {{str(e)}}"

# Create simple Gradio interface using ChatInterface
demo = gr.ChatInterface(
    fn=generate_response,
    title=ASSISTANT_NAME,
    description=ASSISTANT_DESCRIPTION,
    examples={examples}
)

if __name__ == "__main__":
    demo.launch(share=True)
'''

def test_zip_generation():
    """Test the zip generation with sample data"""
    print("🧪 Testing zip generation...")
    
    # Sample configuration
    config = {
        'name': 'Test Assistant',
        'description': 'A helpful test assistant',
        'system_prompt': 'You are a helpful assistant. Provide clear and accurate responses.',
        'model': 'google/gemma-2-9b-it',
        'api_key_var': 'OPENROUTER_API_KEY',
        'temperature': 0.7,
        'max_tokens': 1024,
        'examples': json.dumps([
            "Hello! How can you help me?",
            "Tell me something interesting", 
            "What can you do?"
        ])
    }
    
    print(f"📋 Config: {config}")
    
    try:
        # Generate files
        print("🔄 Generating app.py content...")
        app_content = ASSISTANT_TEMPLATE.format(**config)
        print("✅ App content generated")
        
        print("🔄 Generating README content...")
        readme_content = create_readme(config)
        print("✅ README content generated")
        
        print("🔄 Generating requirements content...")
        requirements_content = create_requirements()
        print("✅ Requirements content generated")
        
        # Create zip file
        print("🔄 Creating zip file...")
        import time
        timestamp = int(time.time())
        filename = f"test_assistant_{timestamp}.zip"
        
        zip_buffer = io.BytesIO()
        with zipfile.ZipFile(zip_buffer, 'w', zipfile.ZIP_DEFLATED) as zip_file:
            zip_file.writestr('app.py', app_content)
            zip_file.writestr('requirements.txt', requirements_content)
            zip_file.writestr('README.md', readme_content)
            zip_file.writestr('config.json', json.dumps(config, indent=2))
        
        # Write to file
        zip_buffer.seek(0)
        with open(filename, 'wb') as f:
            f.write(zip_buffer.getvalue())
        
        # Verify file
        if os.path.exists(filename):
            file_size = os.path.getsize(filename)
            print(f"✅ Zip file created: {filename} ({file_size} bytes)")
            
            # Test zip contents
            with zipfile.ZipFile(filename, 'r') as zip_file:
                files_in_zip = zip_file.namelist()
                print(f"📁 Files in zip: {files_in_zip}")
                
                # Test app.py content
                app_py_content = zip_file.read('app.py').decode('utf-8')
                print(f"📄 app.py first 100 chars: {app_py_content[:100]}...")
                
            return True, filename
        else:
            print("❌ Zip file not created")
            return False, None
            
    except Exception as e:
        print(f"❌ Error during zip generation: {str(e)}")
        import traceback
        traceback.print_exc()
        return False, None

if __name__ == "__main__":
    success, filename = test_zip_generation()
    if success:
        print(f"🎉 Test successful! Generated: {filename}")
    else:
        print("💥 Test failed!")