import os
import gradio as gr
import requests
import json
from datetime import datetime

# Load secrets
PAYWALLS_API_KEY = os.environ.get("PAYWALLS_API_KEY")
PAYWALLS_API_URL = os.environ.get("PAYWALLS_API_URL", "https://api.paywalls.ai/v1")

# Free tier limits
FREE_MAX_TOKENS = 150
PREMIUM_MAX_TOKENS = 2048

# Style presets (premium feature)
STYLE_PRESETS = {
    "neutral": "You are a creative writing assistant.",
    "poetic": "You are a poetic writing assistant. Write in lyrical, expressive language with vivid imagery.",
    "dramatic": "You are a dramatic writing assistant. Create tension, conflict, and emotional depth.",
    "comedic": "You are a comedic writing assistant. Use humor, wit, and light-hearted tone.",
    "horror": "You are a horror writing assistant. Build suspense and create eerie, unsettling atmospheres."
}

def check_paywall(user_id: str):
    """Check user payment status and generate wallet topup link"""
    headers = {
        "Authorization": f"Bearer {PAYWALLS_API_KEY}",
        "Content-Type": "application/json"
    }
    
    # Check if user has active balance/subscription
    try:
        check_response = requests.get(
            f"{PAYWALLS_API_URL}/payments/check",
            headers=headers,
            params={"user_id": user_id},
            timeout=10
        )
        is_paid = check_response.status_code == 200 and check_response.json().get("paid", False)
    except:
        is_paid = False
    
    # Direct users to AI Wallet topup page
    paywall_url = f"https://wallet.paywalls.ai/topup?user={user_id}"
    
    return is_paid, paywall_url

def free_tier_generation(system_message, history, user_message, max_tokens, temperature):
    try:
        from huggingface_hub import InferenceClient
        client = InferenceClient()
        messages = [{"role": "system", "content": system_message}] + history + [{"role": "user", "content": user_message}]
        
        reply = client.chat_completion(
            messages=messages,
            max_tokens=max_tokens,
            temperature=temperature,
            model="meta-llama/Llama-3.2-3B-Instruct"
        ).choices[0].message.content
        
        return reply
    except Exception as e:
        print(f"[Free Tier] Error: {e}")
        return f"⚠️ Free tier service error: {str(e)}"

def premium_generation(user_id, system_message, history, user_message, max_tokens, temperature, top_p):
    """Premium tier using Paywalls proxy"""
    headers = {
        "Authorization": f"Bearer {PAYWALLS_API_KEY}",
        "Content-Type": "application/json"
    }
    
    messages = [{"role": "system", "content": system_message}] + history + [{"role": "user", "content": user_message}]
    
    payload = {
        "model": "gpt-3.5-turbo",
        "messages": messages,
        "max_tokens": max_tokens,
        "temperature": temperature,
        "top_p": top_p,
        "user": user_id
    }
    
    try:
        response = requests.post(
            f"{PAYWALLS_API_URL}/chat/completions",
            headers=headers,
            json=payload,
            timeout=60
        )
        response.raise_for_status()
        data = response.json()
        
        if data.get("choices"):
            return data["choices"][0]["message"]["content"]
        return "No response from premium model."
        
    except requests.exceptions.HTTPError as e:
        print(f"[Premium] HTTP Error: {e.response.status_code} - {e.response.text}")
        return f"⚠️ Premium service error: {e.response.status_code}"
    except Exception as e:
        print(f"[Premium] Error: {e}")
        return f"⚠️ Premium service error: {str(e)}"

def respond(message, history, style, max_tokens, temperature, top_p):
    user_id = "demo_user"
    is_premium, paywall_url = check_paywall(user_id)
    
    # Check style restriction for free users
    if not is_premium:
        max_tokens = min(max_tokens, FREE_MAX_TOKENS)
        if style != "neutral":
            return history + [{"role": "assistant", "content": f"🔒 **Premium Feature Locked**\n\nStyle customization is available for Premium users only.\n\n[✨ Upgrade to Premium]({paywall_url})"}]
    
    system_message = STYLE_PRESETS.get(style, STYLE_PRESETS["neutral"])
    
    # Generate response based on tier
    if is_premium:
        reply = premium_generation(user_id, system_message, history, message, max_tokens, temperature, top_p)
        tier_badge = "✨ **PREMIUM**"
    else:
        reply = free_tier_generation(system_message, history, message, max_tokens, temperature)
        tier_badge = f"🆓 **FREE TIER** • 150 tokens limit\n\n💎 [Unlock Premium Features]({paywall_url}) • Longer stories • Custom styles • Export"
    
    full_reply = f"{reply}\n\n---\n{tier_badge}"
    return history + [{"role": "user", "content": message}, {"role": "assistant", "content": full_reply}]

def export_conversation(history):
    user_id = "demo_user"
    is_premium, paywall_url = check_paywall(user_id)
    
    if not is_premium:
        return None, f"🔒 Export feature requires Premium.\n\n[Upgrade Now]({paywall_url})"
    
    content = f"Creative Writing Session - {datetime.now().strftime('%Y-%m-%d %H:%M')}\n\n"
    for msg in history:
        role = msg["role"].upper()
        content += f"{role}: {msg['content']}\n\n"
    
    filepath = "/tmp/creative_writing.txt"
    with open(filepath, "w") as f:
        f.write(content)
    
    return filepath, "✅ Exported successfully!"

# Custom CSS for elegant design
custom_css = """
@import url('https://fonts.googleapis.com/css2?family=Inter:wght@400;500;600;700&display=swap');

* {
    font-family: 'Inter', sans-serif !important;
}

.gradio-container {
    background: linear-gradient(135deg, #667eea 0%, #764ba2 100%) !important;
}

.main-header {
    background: rgba(255, 255, 255, 0.95);
    backdrop-filter: blur(10px);
    border-radius: 20px;
    padding: 30px;
    margin-bottom: 25px;
    box-shadow: 0 8px 32px rgba(0, 0, 0, 0.1);
    text-align: center;
}

.main-header h1 {
    background: linear-gradient(135deg, #667eea 0%, #764ba2 100%);
    -webkit-background-clip: text;
    -webkit-text-fill-color: transparent;
    font-size: 2.5em;
    font-weight: 700;
    margin-bottom: 10px;
}

.main-header p {
    color: #64748b;
    font-size: 1.1em;
}

.chat-container {
    background: rgba(255, 255, 255, 0.98);
    border-radius: 20px;
    padding: 25px;
    box-shadow: 0 8px 32px rgba(0, 0, 0, 0.1);
}

.controls-container {
    background: rgba(255, 255, 255, 0.98);
    border-radius: 20px;
    padding: 25px;
    box-shadow: 0 8px 32px rgba(0, 0, 0, 0.1);
}

.premium-badge {
    background: linear-gradient(135deg, #f093fb 0%, #f5576c 100%);
    color: white;
    padding: 8px 16px;
    border-radius: 20px;
    font-weight: 600;
    font-size: 0.85em;
    display: inline-block;
    margin-left: 10px;
}

.feature-card {
    background: rgba(102, 126, 234, 0.1);
    border-radius: 12px;
    padding: 15px;
    margin-bottom: 15px;
    border-left: 4px solid #667eea;
}

button {
    border-radius: 12px !important;
    font-weight: 600 !important;
    transition: all 0.3s ease !important;
}

button:hover {
    transform: translateY(-2px) !important;
    box-shadow: 0 8px 20px rgba(0, 0, 0, 0.15) !important;
}

.primary-btn {
    background: linear-gradient(135deg, #667eea 0%, #764ba2 100%) !important;
}

input, textarea, select {
    border-radius: 10px !important;
    border: 2px solid #e2e8f0 !important;
    transition: all 0.3s ease !important;
}

input:focus, textarea:focus, select:focus {
    border-color: #667eea !important;
    box-shadow: 0 0 0 3px rgba(102, 126, 234, 0.1) !important;
}

.slider {
    accent-color: #667eea !important;
}
"""

# Build UI
with gr.Blocks(css=custom_css, title="Creative Writing Assistant", theme=gr.themes.Soft()) as demo:
    
    # Header
    with gr.Row():
        with gr.Column():
            gr.HTML("""
                <div class="main-header">
                    <h1>✨ Creative Writing Assistant</h1>
                    <p>Transform your ideas into captivating stories with AI-powered creativity</p>
                    <div style="margin-top: 20px;">
                        <span style="background: rgba(102, 126, 234, 0.1); padding: 10px 20px; border-radius: 25px; color: #667eea; font-weight: 600; margin: 0 10px;">
                            🆓 Free: 150 tokens
                        </span>
                        <span style="background: linear-gradient(135deg, #f093fb 0%, #f5576c 100%); padding: 10px 20px; border-radius: 25px; color: white; font-weight: 600; margin: 0 10px;">
                            ✨ Premium: 2048 tokens + Styles + Export
                        </span>
                    </div>
                </div>
            """)
    
    # Main Content
    with gr.Row(equal_height=True):
        # Chat Section
        with gr.Column(scale=2, elem_classes="chat-container"):
            gr.Markdown("### 💬 Your Creative Space")
            chatbot = gr.Chatbot(
                type="messages", 
                height=550,
                show_copy_button=True,
                avatar_images=(None, "https://em-content.zobj.net/source/twitter/376/sparkles_2728.png")
            )
            
            with gr.Row():
                msg = gr.Textbox(
                    label="",
                    placeholder="🎨 Write your creative prompt here... (e.g., 'Write a story about a magical garden')",
                    lines=3,
                    scale=9
                )
            
            with gr.Row():
                submit = gr.Button("✨ Generate Story", variant="primary", scale=2, size="lg")
                clear = gr.Button("🗑️ Clear", scale=1, size="lg")
        
        # Controls Section
        with gr.Column(scale=1, elem_classes="controls-container"):
            gr.Markdown("### 🎛️ Creative Controls")
            
            gr.HTML('<div class="feature-card"><b>🎨 Writing Style</b><span class="premium-badge">PREMIUM</span></div>')
            style = gr.Dropdown(
                choices=list(STYLE_PRESETS.keys()),
                value="neutral",
                label="",
                info="Choose your narrative voice"
            )
            
            gr.Markdown("---")
            
            gr.Markdown("**📏 Generation Length**")
            max_tokens = gr.Slider(
                FREE_MAX_TOKENS, 
                PREMIUM_MAX_TOKENS, 
                value=FREE_MAX_TOKENS, 
                step=50, 
                label="Max Tokens",
                info="Free tier limited to 150 tokens",
                elem_classes="slider"
            )
            
            gr.Markdown("**🌡️ Creativity Level**")
            temperature = gr.Slider(
                0.1, 
                2.0, 
                value=0.8, 
                step=0.1, 
                label="Temperature",
                info="Higher = more creative",
                elem_classes="slider"
            )
            
            gr.Markdown("**🎲 Diversity**")
            top_p = gr.Slider(
                0.1, 
                1.0, 
                value=0.95, 
                step=0.05, 
                label="Top-p",
                info="Response variety",
                elem_classes="slider"
            )
            
            gr.Markdown("---")
            
            gr.HTML('<div class="feature-card"><b>💾 Export</b><span class="premium-badge">PREMIUM</span></div>')
            export_btn = gr.Button("📥 Download Story", variant="secondary", size="lg")
            download_file = gr.File(label="Your File", visible=False)
            export_status = gr.Textbox(label="Status", interactive=False, visible=False)
    
    # Footer
    gr.HTML("""
        <div style="text-align: center; margin-top: 30px; padding: 20px; background: rgba(255, 255, 255, 0.9); border-radius: 15px;">
            <p style="color: #64748b; font-size: 0.95em;">
                Powered by AI • Built with ❤️ using Gradio & Paywalls.ai
            </p>
        </div>
    """)
    
    # Event handlers
    submit.click(respond, [msg, chatbot, style, max_tokens, temperature, top_p], [chatbot]).then(lambda: "", None, [msg])
    msg.submit(respond, [msg, chatbot, style, max_tokens, temperature, top_p], [chatbot]).then(lambda: "", None, [msg])
    clear.click(lambda: None, None, [chatbot])
    export_btn.click(export_conversation, [chatbot], [download_file, export_status]).then(
        lambda: (gr.update(visible=True), gr.update(visible=True)), None, [download_file, export_status]
    )

if __name__ == "__main__":
    demo.launch()