import os
import functools
import re
import json
import time
from datetime import datetime
from typing import Dict, List, Any, Tuple
from openai import OpenAI
import tiktoken
import gradio as gr

# Configure API client
af_key = os.getenv("OPENAI_API_KEY")
if not af_key:
    raise ValueError("Please set the OPENAI_API_KEY environment variable.")
client = OpenAI(api_key=af_key)

# Available models
_env_models = os.getenv("OPENAI_MODEL_LIST", "gpt-3.5-turbo,gpt-4,gpt-4o,gpt-4o-mini")
ALL_MODELS = [m.strip() for m in _env_models.split(",") if m.strip()]
if not ALL_MODELS:
    ALL_MODELS = ["gpt-3.5-turbo"]

# Token encoder
@functools.lru_cache(maxsize=64)
def _get_encoding(model: str):
    try:
        return tiktoken.encoding_for_model(model)
    except KeyError:
        return tiktoken.get_encoding("cl100k_base")

def count_tokens(text: str, model: str) -> int:
    enc = _get_encoding(model)
    return len(enc.encode(text))

# Advanced prompt templates
PROMPT_TEMPLATES = {
    "Writing & Content": {
        "Blog Post Writer": {
            "description": "Create engaging blog posts with SEO optimization",
            "system": "You are an expert content writer and SEO specialist.",
            "template": "Write a {tone} blog post about {topic} for {audience}. Include {word_count} words, SEO keywords, and a compelling hook.",
            "variables": ["tone", "topic", "audience", "word_count"],
            "examples": [
                {"input": "Topic: AI in healthcare, Tone: Professional, Audience: Medical professionals", "output": "# Revolutionizing Patient Care: How AI is Transforming Modern Healthcare\n\nThe stethoscope around a doctor's neck might soon be joined by..."}
            ]
        },
        "Email Composer": {
            "description": "Professional email writing for various contexts",
            "system": "You are a professional communication expert.",
            "template": "Write a {tone} email to {recipient} about {subject}. Purpose: {purpose}",
            "variables": ["tone", "recipient", "subject", "purpose"],
            "examples": []
        },
        "Social Media Creator": {
            "description": "Engaging social media content for all platforms",
            "system": "You are a social media expert who creates viral content.",
            "template": "Create a {platform} post about {topic} that is {tone} and includes relevant hashtags. Target audience: {audience}",
            "variables": ["platform", "topic", "tone", "audience"],
            "examples": []
        }
    },
    "Analysis & Research": {
        "Data Analyst": {
            "description": "Comprehensive data analysis and insights",
            "system": "You are a senior data analyst with expertise in statistical analysis and business intelligence.",
            "template": "Analyze the following data about {subject} and provide insights on {focus_area}. Include trends, patterns, and actionable recommendations.",
            "variables": ["subject", "focus_area"],
            "examples": []
        },
        "Competitive Analysis": {
            "description": "In-depth competitor research and positioning",
            "system": "You are a strategic business analyst specializing in competitive intelligence.",
            "template": "Conduct a competitive analysis of {company} in the {industry} market. Focus on {analysis_areas} and provide strategic recommendations.",
            "variables": ["company", "industry", "analysis_areas"],
            "examples": []
        },
        "Research Synthesizer": {
            "description": "Synthesize complex research into clear insights",
            "system": "You are a research expert who excels at synthesizing complex information.",
            "template": "Synthesize research on {topic} focusing on {key_questions}. Provide evidence-based conclusions and identify knowledge gaps.",
            "variables": ["topic", "key_questions"],
            "examples": []
        }
    },
    "Business & Strategy": {
        "Business Plan Creator": {
            "description": "Comprehensive business planning assistance",
            "system": "You are a business strategy consultant with expertise in business plan development.",
            "template": "Create a business plan section for {business_type} focusing on {section}. Include market analysis, financial projections, and risk assessment.",
            "variables": ["business_type", "section"],
            "examples": []
        },
        "SWOT Analyzer": {
            "description": "Strategic SWOT analysis for any situation",
            "system": "You are a strategic planning expert who conducts thorough SWOT analyses.",
            "template": "Conduct a SWOT analysis for {entity} in the context of {situation}. Provide specific, actionable insights for each quadrant.",
            "variables": ["entity", "situation"],
            "examples": []
        },
        "Project Manager": {
            "description": "Project planning and management guidance",
            "system": "You are an experienced project manager with PMP certification.",
            "template": "Create a project plan for {project_name} with timeline {duration}. Include milestones, risks, and resource requirements for {team_size} team members.",
            "variables": ["project_name", "duration", "team_size"],
            "examples": []
        }
    },
    "Technical & Code": {
        "Code Reviewer": {
            "description": "Professional code review and optimization",
            "system": "You are a senior software engineer and code review expert.",
            "template": "Review this {language} code for {purpose}. Focus on {review_areas} and provide specific improvement suggestions.",
            "variables": ["language", "purpose", "review_areas"],
            "examples": []
        },
        "Architecture Designer": {
            "description": "Software architecture and system design",
            "system": "You are a solutions architect with expertise in scalable system design.",
            "template": "Design a {system_type} architecture for {use_case}. Consider {requirements} and provide detailed technical specifications.",
            "variables": ["system_type", "use_case", "requirements"],
            "examples": []
        },
        "Debug Assistant": {
            "description": "Advanced debugging and troubleshooting",
            "system": "You are a debugging expert who can identify and solve complex technical issues.",
            "template": "Help debug this {technology} issue: {problem_description}. Provide step-by-step troubleshooting and root cause analysis.",
            "variables": ["technology", "problem_description"],
            "examples": []
        }
    },
    "Creative & Innovation": {
        "Creative Brainstormer": {
            "description": "Generate innovative ideas and solutions",
            "system": "You are a creative thinking expert who generates innovative solutions.",
            "template": "Generate {number} creative ideas for {challenge} targeting {target_audience}. Use {creativity_method} thinking approach.",
            "variables": ["number", "challenge", "target_audience", "creativity_method"],
            "examples": []
        },
        "Story Writer": {
            "description": "Creative storytelling and narrative development",
            "system": "You are a professional storyteller and creative writer.",
            "template": "Write a {genre} story about {theme} set in {setting}. Target audience: {audience}. Length: {length}",
            "variables": ["genre", "theme", "setting", "audience", "length"],
            "examples": []
        },
        "Innovation Consultant": {
            "description": "Innovation strategy and ideation",
            "system": "You are an innovation consultant who helps organizations think differently.",
            "template": "Develop innovative solutions for {industry} to address {challenge}. Use {innovation_framework} methodology.",
            "variables": ["industry", "challenge", "innovation_framework"],
            "examples": []
        }
    }
}

# Role personas
ROLE_PERSONAS = {
    "Expert Consultant": "You are a world-class expert consultant with 20+ years of experience. You provide authoritative, evidence-based advice.",
    "Creative Collaborator": "You are a creative partner who thinks outside the box and offers innovative perspectives.",
    "Analytical Thinker": "You are a logical, systematic thinker who breaks down complex problems methodically.",
    "Supportive Mentor": "You are a patient, encouraging mentor who guides others to discover insights themselves.",
    "Critical Reviewer": "You are a constructive critic who identifies weaknesses and suggests improvements.",
    "Visionary Strategist": "You are a forward-thinking strategist who sees big-picture opportunities and trends.",
    "Practical Implementer": "You are focused on practical, actionable solutions that can be implemented immediately.",
    "Research Scholar": "You are a meticulous researcher who values accuracy and comprehensive analysis."
}

# Advanced instruction categories
ADVANCED_CATEGORIES = {
    "Cognitive Techniques": [
        "Use chain-of-thought reasoning, showing your step-by-step thinking",
        "Apply the six thinking hats method (white, red, black, yellow, green, blue)",
        "Use the SCAMPER technique for creative problem-solving",
        "Apply first principles thinking to break down complex problems",
        "Use analogical reasoning to explain complex concepts"
    ],
    "Output Structure": [
        "Provide an executive summary followed by detailed analysis",
        "Use the pyramid principle: conclusion first, then supporting arguments",
        "Structure response with clear headings and subheadings",
        "Include a TL;DR section at the beginning",
        "End with specific action items and next steps",
        "Use bullet points for key insights and numbered lists for processes"
    ],
    "Quality Assurance": [
        "Fact-check all claims and provide sources when possible",
        "Include confidence levels for uncertain statements",
        "Acknowledge limitations and assumptions in your analysis",
        "Provide alternative perspectives or counterarguments",
        "Include risk assessment for recommendations"
    ],
    "Audience Adaptation": [
        "Adjust technical depth based on audience expertise level",
        "Use relevant examples and analogies for the target audience",
        "Consider cultural context and sensitivities",
        "Tailor tone and formality to the professional setting",
        "Include stakeholder impact analysis"
    ]
}

# Default filtering
DEFAULT_BANNED_WORDS = [
    "Hurdles", "Tapestry", "Bustling", "Harnessing", "Unveiling the power",
    "Realm", "Depicted", "Demystify", "Insurmountable", "New Era",
    "Poised", "Unravel", "Entanglement", "Unprecedented", "Beacon",
    "Unleash", "Delve", "Enrich", "Multifaceted", "Discover", "Unlock",
    "Tailored", "Elegant", "Dive", "Ever-evolving", "Adventure",
    "Journey", "Navigate", "Navigation"
]

DEFAULT_FILTER_PATTERNS = [
    r"As an AI language model,?\s*",
    r"I(?:'m|'m|\s+am)\s+sorry,?\s*",
    r"I\s+apologize,?\s*",
    r"In\s+conclusion,?\s*",
    r"At\s+the\s+end\s+of\s+the\s+day,?\s*"
]

# Global state for saved prompts and analytics
class PromptState:
    def __init__(self):
        self.saved_prompts = {}
        self.analytics = []
        self.ab_tests = {}
        self.custom_categories = {}
        
    def save_prompt(self, name: str, prompt_data: dict):
        self.saved_prompts[name] = {
            **prompt_data,
            "created_at": datetime.now().isoformat(),
            "id": len(self.saved_prompts)
        }
    
    def log_usage(self, prompt: str, response: str, model: str, tokens: int, rating: int = None):
        self.analytics.append({
            "timestamp": datetime.now().isoformat(),
            "prompt_length": len(prompt),
            "response_length": len(response),
            "model": model,
            "tokens": tokens,
            "rating": rating
        })

state = PromptState()

def validate_regex_patterns(patterns: list) -> tuple[list, list]:
    """Validate regex patterns and return valid ones plus any errors"""
    valid_patterns = []
    errors = []
    
    for pattern in patterns:
        if not pattern.strip():
            continue
        try:
            re.compile(pattern.strip())
            valid_patterns.append(pattern.strip())
        except re.error as e:
            errors.append(f"Invalid pattern '{pattern}': {str(e)}")
    
    return valid_patterns, errors

def smart_filter_text(text: str, banned_words: list, banned_patterns: list) -> str:
    """Apply smart filtering that preserves sentence structure"""
    result = text
    
    # Filter patterns first (more specific) - patterns are pre-validated
    for pattern in banned_patterns:
        if pattern.strip():
            result = re.sub(pattern.strip(), "", result, flags=re.IGNORECASE)
    
    # Filter banned words (preserve word boundaries)
    for word in banned_words:
        if word.strip():
            escaped_word = re.escape(word.strip())
            pattern = rf"\b{escaped_word}\b"
            result = re.sub(pattern, "", result, flags=re.IGNORECASE)
    
    # Clean up extra whitespace and punctuation
    result = re.sub(r'\s+', ' ', result)
    result = re.sub(r'\s+([.!?])', r'\1', result)
    result = re.sub(r'([.!?])\s*([.!?])', r'\1\2', result)
    result = re.sub(r'^[.!?]\s*', '', result)
    
    return result.strip()

def substitute_variables(template: str, variables: dict) -> str:
    """Replace {variable} placeholders with actual values"""
    result = template
    for var, value in variables.items():
        if value:  # Only substitute if value is provided
            result = result.replace(f"{{{var}}}", str(value))
    return result

def optimize_prompt(prompt: str) -> str:
    """Provide optimization suggestions for prompts"""
    suggestions = []
    
    # Check length
    if len(prompt) < 50:
        suggestions.append("Consider adding more context and specific requirements")
    elif len(prompt) > 2000:
        suggestions.append("Consider breaking this into smaller, focused prompts")
    
    # Check for vague terms
    vague_terms = ["good", "better", "nice", "improve", "optimize", "enhance"]
    found_vague = [term for term in vague_terms if term.lower() in prompt.lower()]
    if found_vague:
        suggestions.append(f"Replace vague terms ({', '.join(found_vague)}) with specific criteria")
    
    # Check for examples
    if "example" not in prompt.lower() and len(prompt) > 100:
        suggestions.append("Consider adding specific examples to clarify your expectations")
    
    # Check for output format specification
    format_words = ["format", "structure", "organize", "layout"]
    if not any(word in prompt.lower() for word in format_words):
        suggestions.append("Specify the desired output format (bullet points, paragraphs, etc.)")
    
    if not suggestions:
        return "✅ Your prompt looks well-optimized!"
    
    return "💡 Optimization suggestions:\n• " + "\n• ".join(suggestions)

def make_api_call(model: str, prompt: str, temperature: float = 0.7) -> tuple[str, str, int]:
    """Make API call with error handling and token tracking"""
    try:
        response = client.chat.completions.create(
            model=model,
            messages=[{"role": "user", "content": prompt}],
            temperature=temperature,
            max_tokens=2000
        )
        content = response.choices[0].message.content
        tokens = response.usage.total_tokens if response.usage else count_tokens(prompt + content, model)
        return content, "", tokens
    except Exception as e:
        return "", f"API Error: {str(e)}", 0

def calculate_analytics() -> dict:
    """Calculate usage analytics"""
    if not state.analytics:
        return {}
    
    total_prompts = len(state.analytics)
    avg_tokens = sum(a["tokens"] for a in state.analytics) / total_prompts if total_prompts > 0 else 0
    rated_prompts = [a for a in state.analytics if a["rating"] is not None]
    avg_rating = sum(a["rating"] for a in rated_prompts) / len(rated_prompts) if rated_prompts else 0
    
    return {
        "total_prompts": total_prompts,
        "avg_tokens": round(avg_tokens, 1),
        "avg_rating": round(avg_rating, 2),
        "models_used": len(set(a["model"] for a in state.analytics))
    }

# Custom CSS
custom_css = """
.gradio-container {
    max-width: 1400px !important;
}
.template-card {
    background: linear-gradient(135deg, #667eea 0%, #764ba2 100%);
    border-radius: 12px;
    padding: 20px;
    color: white;
    margin: 10px 0;
}
.analytics-card {
    background: linear-gradient(135deg, #f093fb 0%, #f5576c 100%);
    border-radius: 12px;
    padding: 15px;
    color: white;
    text-align: center;
}
.prompt-preview {
    background-color: #f8f9fa;
    border: 2px dashed #dee2e6;
    border-radius: 8px;
    padding: 15px;
    font-family: monospace;
}
.feature-highlight {
    background: linear-gradient(45deg, #FE6B8B 30%, #FF8E53 90%);
    border-radius: 8px;
    padding: 10px;
    color: white;
    margin: 5px 0;
}
"""

with gr.Blocks(css=custom_css, title="Ultimate LLM Prompt Builder", theme=gr.themes.Soft()) as demo:
    
    # Header
    gr.Markdown("""
    # 🚀 Ultimate LLM Prompt Builder v2.0
    
    **The most comprehensive prompt engineering tool** with templates, variables, A/B testing, analytics, and more!
    
    [![Features](https://img.shields.io/badge/Features-10/10-success)]() [![Templates](https://img.shields.io/badge/Templates-15+-blue)]() [![AI Models](https://img.shields.io/badge/Models-Multiple-orange)]()
    """)
    
    # Quick stats
    with gr.Row():
        with gr.Column(scale=1):
            analytics_display = gr.HTML(value="<div class='analytics-card'>📊 No usage data yet</div>")
        with gr.Column(scale=3):
            gr.Markdown("""
            **🎯 Quick Start:** Choose a template → Customize variables → Add instructions → Generate!
            **💡 Pro Tip:** Use the A/B testing feature to optimize your prompts for better results.
            """, elem_classes=["feature-highlight"])

    # Main interface
    with gr.Tabs() as main_tabs:
        
        # Template Library Tab
        with gr.TabItem("📚 Template Library"):
            gr.Markdown("### Choose from professional prompt templates")
            
            with gr.Row():
                template_category = gr.Dropdown(
                    choices=list(PROMPT_TEMPLATES.keys()),
                    label="🗂️ Category",
                    value=list(PROMPT_TEMPLATES.keys())[0]
                )
                template_name = gr.Dropdown(
                    choices=list(PROMPT_TEMPLATES[list(PROMPT_TEMPLATES.keys())[0]].keys()),
                    label="📄 Template",
                    value=list(PROMPT_TEMPLATES[list(PROMPT_TEMPLATES.keys())[0]].keys())[0]
                )
            
            template_info = gr.Markdown(value="")
            
            # Variable inputs (dynamically generated)
            variables_section = gr.Column(visible=False)
            with variables_section:
                gr.Markdown("### 🔧 Customize Variables")
                variable_inputs = {}
                for i in range(10):  # Support up to 10 variables
                    variable_inputs[f"var_{i}"] = gr.Textbox(label="", visible=False)
        
        # Advanced Prompt Builder Tab  
        with gr.TabItem("⚙️ Advanced Builder"):
            with gr.Row():
                with gr.Column(scale=2):
                    gr.Markdown("### 🎭 Role & System Prompt")
                    role_persona = gr.Dropdown(
                        choices=list(ROLE_PERSONAS.keys()),
                        label="🎭 Role Persona",
                        value="Expert Consultant"
                    )
                    custom_system = gr.Textbox(
                        label="🛠️ Custom System Prompt",
                        placeholder="Override the role persona with a custom system prompt...",
                        lines=3
                    )
                    
                    gr.Markdown("### 📝 Main Prompt")
                    main_prompt = gr.Textbox(
                        label="✍️ Your Prompt",
                        placeholder="Enter your main prompt here...",
                        lines=6
                    )
                    
                    # Few-shot examples
                    gr.Markdown("### 🎯 Few-Shot Examples (Optional)")
                    with gr.Accordion("Add Examples", open=False):
                        example_1_input = gr.Textbox(label="Example 1 - Input", lines=2)
                        example_1_output = gr.Textbox(label="Example 1 - Output", lines=2)
                        example_2_input = gr.Textbox(label="Example 2 - Input", lines=2)
                        example_2_output = gr.Textbox(label="Example 2 - Output", lines=2)
                
                with gr.Column(scale=1):
                    gr.Markdown("### 🧠 Cognitive Instructions")
                    # Create individual components instead of a list
                    with gr.Accordion("📌 Cognitive Techniques", open=False):
                        cognitive_techniques = gr.CheckboxGroup(
                            choices=ADVANCED_CATEGORIES["Cognitive Techniques"],
                            label="Cognitive Techniques",
                            value=[]
                        )
                    with gr.Accordion("📌 Output Structure", open=False):
                        output_structure = gr.CheckboxGroup(
                            choices=ADVANCED_CATEGORIES["Output Structure"],
                            label="Output Structure",
                            value=[]
                        )
                    with gr.Accordion("📌 Quality Assurance", open=False):
                        quality_assurance = gr.CheckboxGroup(
                            choices=ADVANCED_CATEGORIES["Quality Assurance"],
                            label="Quality Assurance",
                            value=[]
                        )
                    with gr.Accordion("📌 Audience Adaptation", open=False):
                        audience_adaptation = gr.CheckboxGroup(
                            choices=ADVANCED_CATEGORIES["Audience Adaptation"],
                            label="Audience Adaptation",
                            value=[]
                        )
        
        # Style & Filtering Tab
        with gr.TabItem("🎨 Style & Filtering"):
            with gr.Row():
                with gr.Column():
                    gr.Markdown("### ✨ Response Style")
                    tone_style = gr.Dropdown(
                        choices=["Professional", "Conversational", "Academic", "Creative", "Technical", "Friendly"],
                        label="🎨 Tone",
                        value="Professional"
                    )
                    length_preference = gr.Dropdown(
                        choices=["Concise", "Moderate", "Detailed", "Comprehensive"],
                        label="📏 Length",
                        value="Moderate"
                    )
                    format_preference = gr.Dropdown(
                        choices=["Paragraph", "Bullet Points", "Numbered List", "Mixed", "Custom"],
                        label="📋 Format",
                        value="Mixed"
                    )
                
                with gr.Column():
                    gr.Markdown("### 🚫 Content Filtering")
                    banned_words = gr.Textbox(
                        label="🚫 Banned Words (comma-separated)",
                        value=",".join(DEFAULT_BANNED_WORDS[:10]),  # Show fewer by default
                        lines=3
                    )
                    banned_patterns = gr.Textbox(
                        label="🔍 Banned Patterns (regex)",
                        value=",".join(DEFAULT_FILTER_PATTERNS),
                        lines=2
                    )
        
        # Multi-Model Testing Tab
        with gr.TabItem("🔬 A/B Testing"):
            gr.Markdown("### Compare responses across different models and prompts")
            
            with gr.Row():
                model_a = gr.Dropdown(choices=ALL_MODELS, label="🤖 Model A", value=ALL_MODELS[0])
                model_b = gr.Dropdown(choices=ALL_MODELS, label="🤖 Model B", value=ALL_MODELS[1] if len(ALL_MODELS) > 1 else ALL_MODELS[0])
            
            with gr.Row():
                prompt_a = gr.Textbox(label="📝 Prompt A", lines=4)
                prompt_b = gr.Textbox(label="📝 Prompt B", lines=4)
            
            test_btn = gr.Button("🚀 Run A/B Test", variant="primary", size="lg")
            
            with gr.Row():
                response_a = gr.Textbox(label="Response A", lines=8, interactive=False)
                response_b = gr.Textbox(label="Response B", lines=8, interactive=False)
            
            with gr.Row():
                rating_a = gr.Slider(minimum=1, maximum=5, step=1, label="Rate Response A", value=3)
                rating_b = gr.Slider(minimum=1, maximum=5, step=1, label="Rate Response B", value=3)
        
        # Chat Interface Tab
        with gr.TabItem("💬 Smart Chat"):
            with gr.Row():
                with gr.Column(scale=2):
                    selected_model = gr.Dropdown(
                        choices=ALL_MODELS,
                        label="🤖 Model",
                        value=ALL_MODELS[0]
                    )
                    temperature = gr.Slider(
                        minimum=0.0,
                        maximum=2.0,
                        step=0.1,
                        value=0.7,
                        label="🌡️ Temperature (creativity)"
                    )
                
                with gr.Column(scale=1):
                    optimization_display = gr.Textbox(
                        label="💡 Prompt Optimization",
                        lines=4,
                        interactive=False
                    )
            
            chat_input = gr.Textbox(
                label="✍️ Your Message",
                placeholder="Ask anything or use variables like {topic}, {style}...",
                lines=3
            )
            
            with gr.Row():
                send_btn = gr.Button("🚀 Send", variant="primary", size="lg")
                clear_chat_btn = gr.Button("🗑️ Clear", variant="secondary")
                save_prompt_btn = gr.Button("💾 Save Prompt", variant="secondary")
            
            save_status = gr.Textbox(label="💾 Save Status", visible=False, interactive=False)
            
            chatbot = gr.Chatbot(
                label="🤖 AI Assistant",
                type="messages",
                height=500,
                show_copy_button=True
            )
            
            # Prompt preview and stats
            with gr.Row():
                final_prompt_preview = gr.Textbox(
                    label="👀 Final Prompt Preview",
                    lines=5,
                    interactive=False
                )
                prompt_stats = gr.JSON(
                    label="📊 Prompt Statistics",
                    value={}
                )
        
        # Saved Prompts & Analytics Tab
        with gr.TabItem("📊 Analytics & Library"):
            with gr.Row():
                with gr.Column():
                    gr.Markdown("### 💾 Saved Prompts")
                    saved_prompts_list = gr.Dropdown(
                        choices=[],
                        label="📚 Select Saved Prompt",
                        allow_custom_value=False
                    )
                    
                    with gr.Row():
                        load_prompt_btn = gr.Button("📥 Load Prompt")
                        delete_prompt_btn = gr.Button("🗑️ Delete", variant="stop")
                    
                    saved_prompt_display = gr.JSON(label="📄 Prompt Details", value={})
                
                with gr.Column():
                    gr.Markdown("### 📈 Usage Analytics")
                    refresh_analytics_btn = gr.Button("🔄 Refresh Analytics")
                    analytics_json = gr.JSON(label="📊 Statistics", value={})
                    
                    gr.Markdown("### 📤 Export/Import")
                    with gr.Row():
                        export_btn = gr.Button("📤 Export Data")
                        import_file = gr.File(label="📥 Import Data", file_types=[".json"])
                    
                    export_data_display = gr.Textbox(
                        label="📋 Export Data (Copy this)",
                        lines=8,
                        interactive=False
                    )
                    
                    import_status = gr.Textbox(
                        label="📥 Import Status",
                        interactive=False
                    )

    # State management
    chat_history = gr.State([])

    # Helper functions
    def update_template_info(category, template_name):
        """Update template information display safely"""
        try:
            if category in PROMPT_TEMPLATES and template_name in PROMPT_TEMPLATES[category]:
                template = PROMPT_TEMPLATES[category][template_name]
                info = f"**{template['description']}**\n\n"
                info += f"**System:** {template['system']}\n\n"
                info += f"**Template:** {template['template']}\n\n"
                if template.get('variables'):
                    info += f"**Variables:** {', '.join(template['variables'])}"
                return info
        except Exception as e:
            return f"Error loading template: {str(e)}"
        return "Select a template to see details"

    def update_template_dropdown(category):
        """Update template dropdown safely"""
        try:
            if category in PROMPT_TEMPLATES:
                choices = list(PROMPT_TEMPLATES[category].keys())
                return gr.update(choices=choices, value=choices[0] if choices else None)
        except Exception:
            pass
        return gr.update(choices=[], value=None)

    def show_template_variables(category, template_name):
        """Show template variables safely"""
        try:
            if category in PROMPT_TEMPLATES and template_name in PROMPT_TEMPLATES[category]:
                template = PROMPT_TEMPLATES[category][template_name]
                variables = template.get('variables', [])
                
                updates = []
                visible_count = 0
                
                for i in range(10):
                    if i < len(variables):
                        var_name = variables[i].replace('_', ' ').title()
                        updates.append(gr.update(visible=True, label=f"📝 {var_name}"))
                        visible_count += 1
                    else:
                        updates.append(gr.update(visible=False))
                
                section_visible = visible_count > 0
                return [gr.update(visible=section_visible)] + updates
        except Exception:
            pass
        
        return [gr.update(visible=False)] + [gr.update(visible=False) for _ in range(10)]

    def build_prompt_from_template(category, template_name, *var_values):
        if category in PROMPT_TEMPLATES and template_name in PROMPT_TEMPLATES[category]:
            template = PROMPT_TEMPLATES[category][template_name]
            variables = template.get('variables', [])
            
            # Create variable dictionary
            var_dict = {}
            for i, var in enumerate(variables):
                if i < len(var_values) and var_values[i]:
                    var_dict[var] = var_values[i]
            
            # Build prompt
            system = template['system']
            main_prompt = substitute_variables(template['template'], var_dict)
            
            return system, main_prompt, var_dict
        
        return "", "", {}

    def create_final_prompt(system, main_prompt, cognitive_tech, output_struct, quality_assur, audience_adapt, examples, tone, length, format_pref, role_persona, custom_system):
        components = []
        
        # System prompt
        if custom_system.strip():
            components.append(f"System: {custom_system.strip()}")
        elif system.strip():
            components.append(f"System: {system}")
        elif role_persona in ROLE_PERSONAS:
            components.append(f"System: {ROLE_PERSONAS[role_persona]}")
        
        # Add cognitive instructions
        all_selected = []
        for selected_list in [cognitive_tech, output_struct, quality_assur, audience_adapt]:
            if selected_list:
                all_selected.extend(selected_list)
        
        if all_selected:
            components.append("Instructions:")
            for i, instruction in enumerate(all_selected, 1):
                components.append(f"{i}. {instruction}")
        
        # Add style preferences
        style_instructions = []
        if tone != "Professional":
            style_instructions.append(f"Tone: {tone}")
        if length != "Moderate":
            style_instructions.append(f"Length: {length}")
        if format_pref != "Mixed":
            style_instructions.append(f"Format: {format_pref}")
        
        if style_instructions:
            components.append("Style: " + ", ".join(style_instructions))
        
        # Add examples if provided
        if examples and any(examples):
            components.append("Examples:")
            for i in range(0, len(examples), 2):
                if i+1 < len(examples) and examples[i] and examples[i+1]:
                    components.append(f"Input: {examples[i]}")
                    components.append(f"Output: {examples[i+1]}")
        
        # Main prompt
        components.append(f"Task: {main_prompt}")
        
        return "\n\n".join(components)

    def process_chat_message(message, history, model, temp, system, cognitive_tech, output_struct, quality_assur, audience_adapt, tone, length, format_pref, role, custom_sys, banned_w, banned_p, *examples):
        """Process chat message with comprehensive error handling"""
        try:
            if not message.strip():
                return history or [], history or [], "", {}, ""
            
            # Create final prompt
            final_prompt = create_final_prompt(
                system, message, cognitive_tech or [], output_struct or [], 
                quality_assur or [], audience_adapt or [], examples, 
                tone, length, format_pref, role, custom_sys
            )
            
            # Get optimization suggestions
            optimization = optimize_prompt(final_prompt)
            
            # Make API call
            response, error, tokens = make_api_call(model, final_prompt, temp)
            
            if error:
                return history or [], history or [], final_prompt, {"error": error}, optimization
            
            # Filter response
            banned_words_list = [w.strip() for w in (banned_w or "").split(",") if w.strip()]
            banned_patterns_input = [p.strip() for p in (banned_p or "").split(",") if p.strip()]
            banned_patterns_list, _ = validate_regex_patterns(banned_patterns_input)
            
            filtered_response = smart_filter_text(response, banned_words_list, banned_patterns_list)
            
            # Log analytics
            state.log_usage(final_prompt, filtered_response, model, tokens)
            
            # Update history
            if history is None:
                history = []
            
            new_history = history + [
                {"role": "user", "content": message},
                {"role": "assistant", "content": filtered_response}
            ]
            
            # Stats
            stats = {
                "tokens": tokens,
                "model": model,
                "temperature": temp,
                "prompt_length": len(final_prompt),
                "response_length": len(filtered_response)
            }
            
            return new_history, new_history, final_prompt, stats, optimization
            
        except Exception as e:
            error_msg = f"Processing error: {str(e)}"
            return history or [], history or [], "", {"error": error_msg}, error_msg

    def run_ab_test(model_a, model_b, prompt_a, prompt_b):
        """Run A/B test with error handling"""
        try:
            if not prompt_a.strip() or not prompt_b.strip():
                return "⚠️ Please enter both prompts", "⚠️ Please enter both prompts"
            
            response_a, error_a, tokens_a = make_api_call(model_a, prompt_a)
            response_b, error_b, tokens_b = make_api_call(model_b, prompt_b)
            
            if error_a:
                response_a = f"❌ Error: {error_a}"
            if error_b:
                response_b = f"❌ Error: {error_b}"
            
            return response_a, response_b
            
        except Exception as e:
            error_msg = f"❌ A/B Test Error: {str(e)}"
            return error_msg, error_msg

    def save_current_prompt(final_prompt, model, stats):
        # Simple prompt saving - in a real app you'd want a name input
        prompt_name = f"Prompt_{len(state.saved_prompts)+1}_{datetime.now().strftime('%H%M')}"
        
        if not final_prompt:
            return gr.update(value="❌ No prompt to save", visible=True)
        
        prompt_data = {
            "prompt": final_prompt,
            "model": model,
            "stats": stats
        }
        
        state.save_prompt(prompt_name, prompt_data)
        
        return gr.update(value=f"✅ Saved as '{prompt_name}'", visible=True)

    def load_saved_prompt(prompt_name):
        if prompt_name in state.saved_prompts:
            return state.saved_prompts[prompt_name]
        return {}
    
    def delete_saved_prompt(prompt_name):
        if prompt_name in state.saved_prompts:
            del state.saved_prompts[prompt_name]
            choices = list(state.saved_prompts.keys())
            return gr.update(choices=choices, value=None), {}
        return gr.update(), {}
    
    def update_saved_prompts_dropdown():
        choices = list(state.saved_prompts.keys())
        return gr.update(choices=choices)

    def refresh_analytics():
        """Refresh analytics with error handling"""
        try:
            analytics = calculate_analytics()
            
            # Update analytics display
            if analytics:
                html = f"""
                <div class='analytics-card'>
                    <h3>📊 Usage Statistics</h3>
                    <p><strong>{analytics['total_prompts']}</strong> prompts generated</p>
                    <p><strong>{analytics['avg_tokens']}</strong> avg tokens per prompt</p>
                    <p><strong>{analytics['avg_rating']}/5</strong> avg rating</p>
                    <p><strong>{analytics['models_used']}</strong> different models used</p>
                </div>
                """
            else:
                html = "<div class='analytics-card'>📊 No usage data yet</div>"
            
            return html, analytics
            
        except Exception as e:
            error_html = f"<div class='analytics-card'>❌ Analytics Error: {str(e)}</div>"
            return error_html, {}

    def export_data():
        """Export data with error handling"""
        try:
            export = {
                "saved_prompts": state.saved_prompts,
                "analytics": state.analytics,
                "export_date": datetime.now().isoformat()
            }
            return json.dumps(export, indent=2)
        except Exception as e:
            return f'{{"error": "Export failed: {str(e)}"}}'

    def import_data_handler(file):
        """Handle file import with proper error handling"""
        if file is None:
            return "❌ No file selected"
        
        try:
            # Read file content
            if hasattr(file, 'read'):
                content = file.read()
                if isinstance(content, bytes):
                    content = content.decode('utf-8')
            else:
                # If file is a path string
                with open(file, 'r', encoding='utf-8') as f:
                    content = f.read()
            
            data = json.loads(content)
            
            if "saved_prompts" in data:
                state.saved_prompts.update(data["saved_prompts"])
            if "analytics" in data:
                state.analytics.extend(data["analytics"])
            
            return "✅ Data imported successfully"
        except Exception as e:
            return f"❌ Import failed: {str(e)}"

    # Wire up all the interactions
    
    # Template system
    template_category.change(
        update_template_dropdown,
        inputs=[template_category],
        outputs=[template_name]
    )
    
    template_category.change(
        update_template_info,
        inputs=[template_category, template_name],
        outputs=[template_info]
    )
    
    template_name.change(
        update_template_info,
        inputs=[template_category, template_name],
        outputs=[template_info]
    )
    
    template_name.change(
        show_template_variables,
        inputs=[template_category, template_name],
        outputs=[variables_section] + [variable_inputs[f"var_{i}"] for i in range(10)]
    )
    
    # Chat interface
    send_btn.click(
        process_chat_message,
        inputs=[
            chat_input, chat_history, selected_model, temperature,
            custom_system, cognitive_techniques, output_structure, quality_assurance, audience_adaptation,
            tone_style, length_preference, format_preference, role_persona, custom_system, banned_words, banned_patterns,
            example_1_input, example_1_output, example_2_input, example_2_output
        ],
        outputs=[chatbot, chat_history, final_prompt_preview, prompt_stats, optimization_display]
    )
    
    def clear_chat_handler():
        """Clear chat history safely"""
        return [], []
    
    clear_chat_btn.click(
        clear_chat_handler,
        outputs=[chatbot, chat_history]
    )
    
    save_prompt_btn.click(
        save_current_prompt,
        inputs=[final_prompt_preview, selected_model, prompt_stats],
        outputs=[save_status]
    )
    
    # A/B Testing
    test_btn.click(
        run_ab_test,
        inputs=[model_a, model_b, prompt_a, prompt_b],
        outputs=[response_a, response_b]
    )
    
    # Analytics
    refresh_analytics_btn.click(
        refresh_analytics,
        outputs=[analytics_display, analytics_json]
    )
    
    # Export/Import
    export_btn.click(
        export_data,
        outputs=[export_data_display]
    )
    
    # Load and delete saved prompts
    load_prompt_btn.click(
        load_saved_prompt,
        inputs=[saved_prompts_list],
        outputs=[saved_prompt_display]
    )
    
    delete_prompt_btn.click(
        delete_saved_prompt,
        inputs=[saved_prompts_list],
        outputs=[saved_prompts_list, saved_prompt_display]
    )
    
    # Import functionality
    import_file.upload(
        import_data_handler,
        inputs=[import_file],
        outputs=[import_status]
    )
    
    def refresh_all_on_load():
        """Helper function to refresh analytics and update dropdowns on load"""
        try:
            analytics_html, analytics_data = refresh_analytics()
            saved_choices = list(state.saved_prompts.keys())
            return analytics_html, analytics_data, gr.update(choices=saved_choices)
        except Exception as e:
            error_html = f"<div class='analytics-card'>❌ Load Error: {str(e)}</div>"
            return error_html, {}, gr.update(choices=[])
    
    # Auto-refresh analytics on load and update saved prompts
    demo.load(
        refresh_all_on_load,
        outputs=[analytics_display, analytics_json, saved_prompts_list]
    )

if __name__ == "__main__":
    demo.launch(share=False, debug=True, show_api=False)