app.py

import os
import gradio as gr
from crewai import Agent, Task, Crew, Process, LLM
from crewai_tools import FileWriterTool, FileReadTool
from dotenv import load_dotenv
import traceback
import shutil
import glob
import json
from datetime import datetime

load_dotenv()

# --- WEBSITE TYPE SYSTEM DEFINITIONS ---
WEBSITE_TYPES = {
    "SaaS Platform": {
        "emoji": "🚀",
        "description": "Software as a Service platform website",
        "pages": {
            "single": ["Landing Page"],
            "multi": ["Home", "Features", "Pricing", "About", "Blog", "Contact", "Login/Signup"]
        },
        "sections": {
            "Hero": "Value proposition with clear CTA for free trial",
            "Features": "Key features with icons and descriptions",
            "How It Works": "3-step process visualization",
            "Pricing": "Tiered pricing cards with feature comparison",
            "Testimonials": "Customer success stories with metrics",
            "Integration": "Logos of integrated platforms",
            "FAQ": "Common questions addressed",
            "CTA": "Final conversion section with trial offer"
        },
        "color_scheme": {
            "primary": "#4F46E5",  # Indigo
            "secondary": "#06B6D4", # Cyan
            "accent": "#10B981",    # Emerald
            "background": "#FFFFFF",
            "text": "#1F2937",
            "muted": "#6B7280"
        },
        "design_style": "Modern, clean, professional with subtle gradients",
        "fonts": {
            "heading": "Inter",
            "body": "Inter"
        },
        "must_have_elements": [
            "Sticky navigation with CTA button",
            "Animated feature cards",
            "Pricing toggle (monthly/yearly)",
            "Trust badges and security icons",
            "Live chat widget placeholder",
            "Cookie consent banner"
        ]
    },
    
    "E-commerce Store": {
        "emoji": "🛍️",
        "description": "Online shopping platform",
        "pages": {
            "single": ["Product Showcase"],
            "multi": ["Home", "Shop", "Product Details", "Cart", "Checkout", "Account", "About", "Contact"]
        },
        "sections": {
            "Hero": "Featured products carousel with offers",
            "Categories": "Product categories grid",
            "Featured Products": "Best sellers and new arrivals",
            "Deals": "Special offers and discounts",
            "Benefits": "Free shipping, returns, security",
            "Newsletter": "Email subscription for deals",
            "Reviews": "Customer testimonials",
            "Brands": "Partner brand logos"
        },
        "color_scheme": {
            "primary": "#DC2626",    # Red
            "secondary": "#F59E0B",  # Amber
            "accent": "#10B981",     # Green for CTAs
            "background": "#FFFFFF",
            "text": "#111827",
            "muted": "#6B7280"
        },
        "design_style": "Clean, product-focused with high-quality imagery",
        "fonts": {
            "heading": "Poppins",
            "body": "Open Sans"
        },
        "must_have_elements": [
            "Shopping cart icon with counter",
            "Search bar with filters",
            "Product quick view modals",
            "Wishlist functionality",
            "Product image zoom",
            "Rating stars",
            "Size/color selectors",
            "Secure payment badges"
        ]
    },
    
    "Portfolio/Creative": {
        "emoji": "🎨",
        "description": "Creative professional or agency portfolio",
        "pages": {
            "single": ["Portfolio Showcase"],
            "multi": ["Home", "Portfolio", "About", "Services", "Process", "Contact"]
        },
        "sections": {
            "Hero": "Creative statement with animated text",
            "Work": "Portfolio grid with hover effects",
            "About": "Story and expertise",
            "Services": "What we offer",
            "Process": "Creative workflow",
            "Clients": "Client logos and testimonials",
            "Awards": "Recognition and achievements",
            "Contact": "Get in touch form"
        },
        "color_scheme": {
            "primary": "#000000",
            "secondary": "#FFFFFF",
            "accent": "#FF6B6B",
            "background": "#FAFAFA",
            "text": "#1A1A1A",
            "muted": "#666666"
        },
        "design_style": "Bold, minimalist, typography-focused with creative layouts",
        "fonts": {
            "heading": "Playfair Display",
            "body": "Roboto"
        },
        "must_have_elements": [
            "Masonry or grid portfolio layout",
            "Smooth scroll animations",
            "Before/after sliders",
            "Video backgrounds",
            "Parallax effects",
            "Custom cursor",
            "Loading animation",
            "Social media links"
        ]
    },
    
    "Corporate/Business": {
        "emoji": "🏢",
        "description": "Professional business or corporate website",
        "pages": {
            "single": ["Company Overview"],
            "multi": ["Home", "About", "Services", "Industries", "Case Studies", "Team", "Careers", "Contact"]
        },
        "sections": {
            "Hero": "Company vision with professional imagery",
            "Services": "Core service offerings",
            "About": "Company history and values",
            "Statistics": "Key metrics and achievements",
            "Industries": "Sectors served",
            "Team": "Leadership profiles",
            "Partners": "Strategic partnerships",
            "News": "Latest updates and press"
        },
        "color_scheme": {
            "primary": "#1E40AF",    # Blue
            "secondary": "#0F766E",  # Teal
            "accent": "#F59E0B",     # Amber
            "background": "#FFFFFF",
            "text": "#1F2937",
            "muted": "#6B7280"
        },
        "design_style": "Professional, trustworthy, clean with structured layouts",
        "fonts": {
            "heading": "Montserrat",
            "body": "Source Sans Pro"
        },
        "must_have_elements": [
            "Professional navigation",
            "Dropdown menus",
            "Team member cards",
            "Timeline components",
            "Statistics counters",
            "Certificate badges",
            "Newsletter signup",
            "Office locations map"
        ]
    },
    
    "Restaurant/Food": {
        "emoji": "🍕",
        "description": "Restaurant, cafe, or food service website",
        "pages": {
            "single": ["Restaurant Landing"],
            "multi": ["Home", "Menu", "About", "Gallery", "Reservations", "Events", "Contact"]
        },
        "sections": {
            "Hero": "Appetizing hero with reservation CTA",
            "Menu Preview": "Featured dishes showcase",
            "About": "Our story and chef",
            "Specials": "Daily specials and offers",
            "Gallery": "Food and ambiance photos",
            "Testimonials": "Customer reviews",
            "Hours": "Opening hours and location",
            "Reservation": "Table booking form"
        },
        "color_scheme": {
            "primary": "#DC2626",    # Red
            "secondary": "#854D0E",  # Brown
            "accent": "#16A34A",     # Green
            "background": "#FEF3C7",
            "text": "#1F2937",
            "muted": "#78716C"
        },
        "design_style": "Warm, inviting, appetizing with rich imagery",
        "fonts": {
            "heading": "Merriweather",
            "body": "Open Sans"
        },
        "must_have_elements": [
            "Menu with prices",
            "Online reservation system",
            "Photo gallery",
            "Google Maps integration",
            "Social media feed",
            "Special offers banner",
            "Delivery/takeout options",
            "Customer reviews"
        ]
    },
    
    "Healthcare/Medical": {
        "emoji": "🏥",
        "description": "Medical, healthcare, or wellness website",
        "pages": {
            "single": ["Clinic Overview"],
            "multi": ["Home", "Services", "Doctors", "Appointments", "Patient Portal", "Resources", "About", "Contact"]
        },
        "sections": {
            "Hero": "Welcome with appointment booking",
            "Services": "Medical services offered",
            "Doctors": "Medical team profiles",
            "Why Choose Us": "Credentials and approach",
            "Testimonials": "Patient stories",
            "Insurance": "Accepted providers",
            "Resources": "Health tips and articles",
            "Emergency": "Emergency contact info"
        },
        "color_scheme": {
            "primary": "#0891B2",    # Cyan
            "secondary": "#10B981",  # Green
            "accent": "#3B82F6",     # Blue
            "background": "#F0FDFA",
            "text": "#0F172A",
            "muted": "#64748B"
        },
        "design_style": "Clean, professional, trustworthy with calming colors",
        "fonts": {
            "heading": "Nunito Sans",
            "body": "Open Sans"
        },
        "must_have_elements": [
            "Appointment booking system",
            "Doctor profiles with credentials",
            "Service descriptions",
            "Patient portal login",
            "Insurance information",
            "Emergency contacts",
            "Health resources/blog",
            "HIPAA compliance notice"
        ]
    },
    
    "Education/E-learning": {
        "emoji": "🎓",
        "description": "Educational institution or online learning platform",
        "pages": {
            "single": ["Course Landing"],
            "multi": ["Home", "Courses", "Programs", "Instructors", "Resources", "Student Portal", "About", "Contact"]
        },
        "sections": {
            "Hero": "Learn and grow message with enrollment CTA",
            "Courses": "Popular courses grid",
            "Features": "Learning platform benefits",
            "Instructors": "Expert teacher profiles",
            "Testimonials": "Student success stories",
            "Stats": "Graduation and success rates",
            "Process": "How to get started",
            "FAQ": "Common questions"
        },
        "color_scheme": {
            "primary": "#7C3AED",    # Purple
            "secondary": "#2563EB",  # Blue
            "accent": "#F59E0B",     # Amber
            "background": "#FAFAF9",
            "text": "#18181B",
            "muted": "#71717A"
        },
        "design_style": "Modern, engaging, accessible with clear information hierarchy",
        "fonts": {
            "heading": "Raleway",
            "body": "Inter"
        },
        "must_have_elements": [
            "Course catalog with filters",
            "Enrollment/registration forms",
            "Student dashboard mockup",
            "Video player placeholder",
            "Progress indicators",
            "Certificate badges",
            "Calendar/schedule",
            "Resource library"
        ]
    },
    
    "Real Estate": {
        "emoji": "🏠",
        "description": "Real estate agency or property listing website",
        "pages": {
            "single": ["Property Showcase"],
            "multi": ["Home", "Properties", "Property Details", "Agents", "Services", "About", "Contact"]
        },
        "sections": {
            "Hero": "Property search with filters",
            "Featured": "Featured properties carousel",
            "Search": "Advanced search options",
            "Services": "Buying, selling, renting",
            "Agents": "Team of realtors",
            "Testimonials": "Client experiences",
            "Market": "Market trends and insights",
            "Contact": "Schedule viewing"
        },
        "color_scheme": {
            "primary": "#0F766E",    # Teal
            "secondary": "#B45309",  # Brown
            "accent": "#3B82F6",     # Blue
            "background": "#FFFFFF",
            "text": "#1F2937",
            "muted": "#6B7280"
        },
        "design_style": "Professional, clean, image-heavy with property showcases",
        "fonts": {
            "heading": "Playfair Display",
            "body": "Source Sans Pro"
        },
        "must_have_elements": [
            "Property search with filters",
            "Property cards with details",
            "Image galleries",
            "Virtual tour buttons",
            "Mortgage calculator",
            "Agent contact cards",
            "Map integration",
            "Property comparison tool"
        ]
    },
    
    "Fitness/Gym": {
        "emoji": "🏋️",
        "description": "Gym, fitness center, or personal training website",
        "pages": {
            "single": ["Fitness Landing"],
            "multi": ["Home", "Classes", "Trainers", "Membership", "Schedule", "Blog", "Contact"]
        },
        "sections": {
            "Hero": "Transform your body message with join CTA",
            "Classes": "Class types and schedules",
            "Trainers": "Expert trainer profiles",
            "Membership": "Pricing plans",
            "Transformation": "Before/after gallery",
            "Features": "Gym facilities and equipment",
            "Schedule": "Class timetable",
            "Blog": "Fitness tips and nutrition"
        },
        "color_scheme": {
            "primary": "#DC2626",    # Red
            "secondary": "#000000",  # Black
            "accent": "#22C55E",     # Green
            "background": "#FAFAFA",
            "text": "#171717",
            "muted": "#525252"
        },
        "design_style": "Energetic, bold, motivational with dynamic imagery",
        "fonts": {
            "heading": "Bebas Neue",
            "body": "Roboto"
        },
        "must_have_elements": [
            "Class schedule/timetable",
            "Membership pricing cards",
            "Trainer profiles",
            "Transformation gallery",
            "Class booking system",
            "Nutrition tips section",
            "BMI calculator",
            "Free trial CTA"
        ]
    },
    
    "Technology/Startup": {
        "emoji": "💡",
        "description": "Tech startup or innovation company website",
        "pages": {
            "single": ["Startup Landing"],
            "multi": ["Home", "Product", "Solutions", "Technology", "Company", "Blog", "Contact"]
        },
        "sections": {
            "Hero": "Innovation message with demo request",
            "Problem/Solution": "What we solve",
            "Product": "Product features showcase",
            "Technology": "Our tech stack",
            "Partners": "Integration partners",
            "Team": "Founding team",
            "Investors": "Backed by",
            "Press": "Media mentions"
        },
        "color_scheme": {
            "primary": "#8B5CF6",    # Purple
            "secondary": "#3B82F6",  # Blue
            "accent": "#10B981",     # Green
            "background": "#0F0F1E",
            "text": "#FFFFFF",
            "muted": "#9CA3AF"
        },
        "design_style": "Futuristic, gradient-heavy, dark mode with neon accents",
        "fonts": {
            "heading": "Space Grotesk",
            "body": "Inter"
        },
        "must_have_elements": [
            "Animated hero section",
            "Interactive product demo",
            "Tech stack icons",
            "Gradient backgrounds",
            "Particle effects",
            "Timeline roadmap",
            "API documentation link",
            "Developer resources"
        ]
    }
}

# --- ENHANCED PROMPTS ---
ENHANCED_IMAGE_PROMPT = """
Image Guidelines for {website_type}:

Use static.photos service with appropriate categories:
- {category_suggestions}

Recommended dimensions:
- Hero: 1200x630
- Features: 640x360
- Testimonials: 200x200
- Gallery: 640x360
- Icons: 200x200

Example URLs:
{example_urls}

Match images to the {design_style} style.
"""

PROJECT_CONTEXT_PROMPT = """
You are creating a {website_type} website with the following specifications:

Project Name: {project_name}
Type: {page_type}
Style: {design_style}
Primary Purpose: {description}

Color Palette:
{color_details}

Typography:
- Headings: {heading_font}
- Body: {body_font}

Required Sections:
{sections}

Must-Have Elements:
{must_have_elements}

Target Audience: {target_audience}
"""

# --- SYSTEM SETUP ---
try:
    llm = LLM(
        model="gemini/gemini-2.5-pro",
        temperature=0.7,
        api_key=os.environ.get("GOOGLE_API_KEY")
    )
except Exception as e:
    print(f"Error initializing Gemini LLM: {e}")
    llm = None

os.makedirs("outputs", exist_ok=True)
file_write_tool = FileWriterTool()
# Add FileReadTool for modification tasks
file_read_tool = FileReadTool()

# --- ENHANCED AGENTS ---
designer = Agent(
    role='Senior Design Systems Architect',
    goal='Create comprehensive design systems tailored to specific website types',
    backstory=(
        "You are a design systems expert who specializes in creating cohesive, "
        "type-specific design languages. You understand that different website types "
        "have distinct requirements:\n\n"
        "- SaaS needs trust indicators and clear value props\n"
        "- E-commerce requires product focus and purchase flow\n"
        "- Portfolio demands creative expression\n"
        "- Corporate needs professionalism and structure\n\n"
        "You always consider:\n"
        "- Industry-specific design patterns\n"
        "- User expectations for each website type\n"
        "- Conversion optimization for the specific goal\n"
        "- Accessibility and performance\n"
        "- Mobile-first responsive design\n\n"
        "You create detailed design tokens including:\n"
        "- Precise color values with semantic naming\n"
        "- Typography scales with exact sizes\n"
        "- Spacing systems (4px/8px grid)\n"
        "- Component specifications\n"
        "- Animation timings and easings"
    ),
    verbose=True,
    llm=llm,
    allow_delegation=False
)

copywriter = Agent(
    role='Conversion Copywriter & Content Strategist',
    goal='Write compelling, type-specific copy that converts visitors into customers',
    backstory=(
        "You are a conversion-focused copywriter who understands that different "
        "website types require different voices and messaging strategies:\n\n"
        "- SaaS: Focus on benefits, ROI, and problem-solving\n"
        "- E-commerce: Product benefits, urgency, and trust\n"
        "- Portfolio: Creativity, expertise, and personality\n"
        "- Corporate: Authority, reliability, and professionalism\n"
        "- Healthcare: Empathy, trust, and expertise\n\n"
        "You craft:\n"
        "- Headlines that grab attention\n"
        "- Subheadings that maintain interest\n"
        "- CTAs that drive action\n"
        "- Microcopy that guides users\n"
        "- SEO-optimized content\n\n"
        "You always match tone to the target audience and industry expectations."
    ),
    verbose=True,
    llm=llm,
    allow_delegation=False
)

developer = Agent(
    role='Full-Stack Frontend Developer',
    goal='Build pixel-perfect, performant websites tailored to specific industries',
    backstory=(
        "You are an expert frontend developer who specializes in building "
        "industry-specific websites with modern technologies:\n\n"
        "**Technical Expertise:**\n"
        "- Tailwind CSS for rapid, consistent styling\n"
        "- Custom CSS for unique requirements\n"
        "- JavaScript for interactivity\n"
        "- AOS for scroll animations\n"
        "- Performance optimization\n\n"
        "**Industry-Specific Knowledge:**\n"
        "- SaaS: Pricing toggles, feature comparisons, dashboards\n"
        "- E-commerce: Product grids, cart functionality, filters\n"
        "- Portfolio: Galleries, sliders, creative layouts\n"
        "- Corporate: Professional navigation, team sections\n"
        "- Healthcare: Appointment forms, service listings\n\n"
        "You always implement:\n"
        "- Semantic HTML5\n"
        "- ARIA labels for accessibility\n"
        "- Responsive design (mobile-first)\n"
        "- SEO best practices\n"
        "- Performance optimizations\n"
        "- Cross-browser compatibility"
    ),
    verbose=True,
    llm=llm,
    tools=[file_write_tool, file_read_tool], # Give agent both read and write tools
    allow_delegation=False
)

# --- ENHANCED TASK CREATION ---
def create_design_task(website_type, project_name, custom_requirements, page_type):
    """Create a design task specific to the website type"""
    
    type_config = WEBSITE_TYPES[website_type]
    
    # Build context
    context = PROJECT_CONTEXT_PROMPT.format(
        website_type=website_type,
        project_name=project_name,
        page_type=page_type,
        design_style=type_config["design_style"],
        description=type_config["description"],
        color_details=json.dumps(type_config["color_scheme"], indent=2),
        heading_font=type_config["fonts"]["heading"],
        body_font=type_config["fonts"]["body"],
        sections=json.dumps(type_config["sections"], indent=2),
        must_have_elements="\n".join(f"- {elem}" for elem in type_config["must_have_elements"]),
        target_audience=f"Target audience for {website_type}"
    )
    
    # Build image guidelines
    category_map = {
        "SaaS Platform": "technology, office, abstract, gradient",
        "E-commerce Store": "retail, product, lifestyle, fashion",
        "Portfolio/Creative": "abstract, minimal, artistic, workspace",
        "Corporate/Business": "office, business, professional, cityscape",
        "Restaurant/Food": "food, restaurant, culinary, ambiance",
        "Healthcare/Medical": "medical, wellness, healthcare, people",
        "Education/E-learning": "education, study, classroom, books",
        "Real Estate": "estate, architecture, interior, exterior",
        "Fitness/Gym": "fitness, sport, gym, wellness",
        "Technology/Startup": "technology, innovation, abstract, futuristic"
    }
    
    image_guidelines = ENHANCED_IMAGE_PROMPT.format(
        website_type=website_type,
        category_suggestions=category_map.get(website_type, "abstract, gradient"),
        design_style=type_config["design_style"],
        example_urls="\n".join([
            f"- Hero: http://static.photos/{category_map.get(website_type, 'abstract').split(',')[0].strip()}/1200x630/100",
            f"- Features: http://static.photos/{category_map.get(website_type, 'abstract').split(',')[0].strip()}/640x360/200",
            f"- Team: http://static.photos/people/200x200/300"
        ])
    )
    
    return Task(
        description=(
            f"{context}\n\n"
            f"Additional Requirements from User:\n{custom_requirements}\n\n"
            f"{image_guidelines}\n\n"
            "Create a comprehensive Design System that includes:\n"
            "1. Complete color palette with exact hex codes\n"
            "2. Typography system with sizes and weights\n"
            "3. Spacing and grid system\n"
            "4. Component specifications\n"
            "5. Animation guidelines\n"
            "6. Image URLs using static.photos\n"
            "7. Layout structure for all sections\n"
            "8. Responsive breakpoints\n"
            "9. Interaction patterns\n"
            "10. Accessibility requirements"
        ),
        expected_output=(
            "A detailed Design System document in markdown format with:\n"
            "- All specifications clearly defined\n"
            "- Exact measurements and values\n"
            "- Static.photos URLs for all images\n"
            "- Component states and variations\n"
            "- Clear implementation guidelines"
        ),
        agent=designer
    )

def create_copywriting_task(website_type, project_name):
    """Create a copywriting task specific to the website type"""
    
    type_config = WEBSITE_TYPES[website_type]
    
    return Task(
        description=(
            f"Write compelling copy for a {website_type} called '{project_name}'.\n\n"
            f"Website Sections:\n{json.dumps(type_config['sections'], indent=2)}\n\n"
            "For each section, provide:\n"
            "- Headlines (powerful, benefit-focused)\n"
            "- Subheadlines (supporting, clarifying)\n"
            "- Body copy (clear, engaging)\n"
            "- CTAs (action-oriented)\n"
            "- Microcopy (helpful, guiding)\n\n"
            f"Match the tone to {website_type} expectations.\n"
            "Include industry-specific terminology and benefits."
        ),
        expected_output=(
            "Complete website copy in markdown format with:\n"
            "- All sections populated with relevant content\n"
            "- Consistent brand voice throughout\n"
            "- SEO-optimized headlines\n"
            "- Clear CTAs for each section\n"
            "- Professional, polished copy"
        ),
        agent=copywriter
    )

def create_development_task(website_type, page_type):
    """Create a development task specific to the website type"""
    
    type_config = WEBSITE_TYPES[website_type]
    pages = type_config["pages"]["multi" if page_type == "Multi-Page" else "single"]
    
    return Task(
        description=(
            f"Build a complete {website_type} website.\n\n"
            f"Pages to create: {', '.join(pages)}\n"
            f"Must-have elements: {', '.join(type_config['must_have_elements'])}\n\n"
            "Technical Requirements:\n"
            "- Use Tailwind CSS via CDN\n"
            "- Implement AOS animations\n"
            "- Use static.photos for all images\n"
            "- Create responsive layouts\n"
            "- Add interactive elements\n"
            "- Implement smooth scrolling\n"
            "- Add hover effects\n"
            "- Include mobile menu\n\n"
            f"Save files to 'outputs/' directory using the FileWriterTool:\n"
            f"- {'index.html' if page_type == 'Single-Page' else 'Multiple HTML files'}"
        ),
        expected_output=(
            "Complete, production-ready HTML files that:\n"
            "- Implement the exact design system\n"
            "- Use all specified copy\n"
            "- Include all required elements\n"
            "- Are fully responsive\n"
            "- Have smooth animations\n"
            "- Work in all browsers"
        ),
        agent=developer,
        tools=[file_write_tool]
    )

# --- CHAT HISTORY MANAGEMENT ---
class ChatHistory:
    def __init__(self):
        self.history = []
        self.current_project = None
        self.website_type = None
        self.project_name = None
        
    def add_message(self, role, content):
        self.history.append({
            "role": role,
            "content": content,
            "timestamp": datetime.now().isoformat()
        })
    
    def get_context(self):
        if not self.history:
            return ""
        return "\n".join([f"{msg['role']}: {msg['content']}" for msg in self.history[-5:]])
    
    def reset(self):
        self.history = []
        self.current_project = None
        self.website_type = None
        self.project_name = None

chat_history = ChatHistory()

# --- HELPER FUNCTIONS ---
def clean_outputs_directory():
    """Clean the outputs directory"""
    if os.path.exists("outputs"):
        if os.path.exists("outputs/project.zip"):
            os.remove("outputs/project.zip")
        for file in glob.glob("outputs/*.html"):
            try:
                os.remove(file)
            except Exception as e:
                print(f"Error removing {file}: {e}")

def create_zip_archive():
    """Create a ZIP archive of output files"""
    try:
        # We need to archive the contents of 'outputs', not the directory itself
        # This makes the zip file cleaner when extracted.
        base_name = "outputs/project"
        root_dir = "outputs"
        # We'll temporarily remove the zip file if it exists to avoid adding it to itself
        if os.path.exists(base_name + '.zip'):
            os.remove(base_name + '.zip')
            
        shutil.make_archive(base_name, 'zip', root_dir, base_dir=None)
        return "outputs/project.zip"
    except Exception as e:
        print(f"Error creating ZIP: {e}")
        return None

# --- MAIN GENERATION FUNCTION ---
def generate_website(website_type, project_name_input, name_method, page_type, custom_requirements):
    try:
        if llm is None:
            return (
                "❌ **Error:** Gemini LLM not initialized. Please check your GOOGLE_API_KEY.",
                None, gr.update(visible=False), gr.update(visible=False), gr.update(visible=False)
            )
        
        # Handle project name
        if name_method == "Enter manually":
            project_name = project_name_input if project_name_input else "My Project"
        else:
            # Generate project name with AI
            project_name = f"{WEBSITE_TYPES[website_type]['emoji']} {website_type.split('/')[0]} Pro"
        
        # Reset chat history for a new project
        chat_history.reset()
        chat_history.website_type = website_type
        chat_history.project_name = project_name
        chat_history.add_message("User", f"Generate a new {website_type} website named '{project_name}' ({page_type}).")
        
        # Clean outputs
        clean_outputs_directory()
        
        # Create tasks
        design_task = create_design_task(website_type, project_name, custom_requirements, page_type)
        copywriting_task = create_copywriting_task(website_type, project_name)
        copywriting_task.context = [design_task]
        development_task = create_development_task(website_type, page_type)
        development_task.context = [design_task, copywriting_task]
        
        # Create and run crew
        crew = Crew(
            agents=[designer, copywriter, developer],
            tasks=[design_task, copywriting_task, development_task],
            process=Process.sequential,
            verbose=True
        )
        
        result = crew.kickoff()
        
        # Check output files
        html_files = glob.glob("outputs/*.html")
        if not html_files:
            return (
                "❌ No HTML files were generated. The AI crew might have failed. Please check the console logs and try again.",
                None, gr.update(visible=False), gr.update(visible=False), gr.update(visible=False)
            )
        
        # Create preview
        preview_file = "outputs/index.html" if os.path.exists("outputs/index.html") else html_files[0]
        with open(preview_file, "r", encoding="utf-8") as file:
            html_content = file.read()
        
        preview_html = f'<iframe srcdoc="{html_content.replace(chr(34), "&quot;")}" width="100%" height="800px" style="border: 2px solid #e0e0e0; border-radius: 8px;"></iframe>'
        
        # Prepare response
        is_multipage = len(html_files) > 1
        success_message = f"✅ **{project_name}** generated successfully!\n"
        success_message += f"Type: {website_type} | Pages: {len(html_files)}"
        
        # **** FIX: Set current_project to enable modification ****
        chat_history.current_project = project_name
        chat_history.add_message("System", f"Successfully generated {project_name}. Files are available.")
        
        if is_multipage:
            zip_path = create_zip_archive()
            return (
                success_message,
                preview_html,
                gr.update(value=preview_file, visible=True),
                gr.update(value=zip_path, visible=True) if zip_path else gr.update(visible=False),
                gr.update(visible=True)
            )
        else:
            return (
                success_message,
                preview_html,
                gr.update(value=preview_file, visible=True),
                gr.update(visible=False),
                gr.update(visible=True)
            )
            
    except Exception as e:
        error_trace = traceback.format_exc()
        print(error_trace)
        return (
            f"❌ An unexpected error occurred during generation: {str(e)}",
            None, gr.update(visible=False), gr.update(visible=False), gr.update(visible=False)
        )

# **** FIX: Complete rewrite of modify_website function ****
def modify_website(modification_request):
    """Modify the existing website based on user request and update the UI."""
    # This function now returns a 5-element tuple to update all relevant UI components
    error_response = (
        "❌ Please generate a website first before requesting modifications.",
        None, gr.update(visible=False), gr.update(visible=False), gr.update(visible=False)
    )

    if not chat_history.current_project:
        return error_response

    if not modification_request or not modification_request.strip():
        return ("⚠️ Please enter a modification request.", None, gr.update(), gr.update(), gr.update())

    chat_history.add_message("User", modification_request)

    try:
        modification_task = Task(
            description=(
                f"You are modifying an existing '{chat_history.website_type}' website named '{chat_history.project_name}'.\n\n"
                "First, you MUST use the FileReadTool to read the existing HTML file(s) in the 'outputs' directory to understand the current code structure. The main file is likely 'index.html'.\n\n"
                f"The user's new modification request is: '{modification_request}'\n\n"
                "Based on this request, you must update the HTML file(s) in the 'outputs' directory. "
                "Use the FileWriterTool to overwrite the existing file(s) with the necessary changes. "
                "Do NOT create new files unless the request explicitly asks for a new page. Overwrite the existing ones.\n"
                "Maintain the existing design system, colors, fonts, Tailwind CSS classes, and overall structure unless specifically asked to change them."
            ),
            expected_output="The complete, updated HTML code written back to the original file(s) in the 'outputs' directory using the FileWriterTool.",
            agent=developer,
            tools=[file_read_tool, file_write_tool]
        )

        crew = Crew(
            agents=[developer],
            tasks=[modification_task],
            process=Process.sequential,
            verbose=True
        )
        
        result = crew.kickoff()
        chat_history.add_message("System", f"Modification applied: {modification_request[:50]}...")

        # After modification, refresh the UI with the new files
        html_files = glob.glob("outputs/*.html")
        if not html_files:
            return "❌ Modification failed: No HTML files found after update.", None, gr.update(), gr.update(), gr.update(visible=False)
        
        preview_file = "outputs/index.html" if os.path.exists("outputs/index.html") else html_files[0]
        with open(preview_file, "r", encoding="utf-8") as file:
            html_content = file.read()

        preview_html = f'<iframe srcdoc="{html_content.replace(chr(34), "&quot;")}" width="100%" height="800px" style="border: 2px solid #e0e0e0; border-radius: 8px;"></iframe>'
        
        is_multipage = len(html_files) > 1
        success_message = f"✅ Modifications applied successfully to **{chat_history.project_name}**!"
        
        if is_multipage:
            zip_path = create_zip_archive()
            return (
                success_message,
                preview_html,
                gr.update(value=preview_file, visible=True),
                gr.update(value=zip_path, visible=True) if zip_path else gr.update(visible=False),
                gr.update(visible=True)
            )
        else:
            return (
                success_message,
                preview_html,
                gr.update(value=preview_file, visible=True),
                gr.update(visible=False),
                gr.update(visible=True)
            )

    except Exception as e:
        error_trace = traceback.format_exc()
        print(error_trace)
        return f"❌ Error applying modifications: {str(e)}", None, gr.update(), gr.update(), gr.update(visible=False)


# --- GRADIO INTERFACE ---
with gr.Blocks(theme=gr.themes.Soft(), css="""
    .container { max-width: 1400px; margin: auto; }
    .type-card {
        border: 2px solid #e2e8f0;
        border-radius: 12px;
        padding: 15px;
        margin: 10px 0;
        background: linear-gradient(135deg, #f7fafc 0%, #edf2f7 100%);
        cursor: pointer;
        transition: all 0.3s ease;
    }
    .type-card:hover {
        transform: translateY(-2px);
        box-shadow: 0 8px 20px rgba(0,0,0,0.1);
        border-color: #4F46E5;
    }
    .type-title {
        font-size: 1.2rem;
        font-weight: 600;
        color: #1a202c;
        margin-bottom: 8px;
    }
    .type-description {
        color: #4a5568;
        font-size: 0.9rem;
    }
    .status-box { 
        padding: 20px; 
        border-radius: 12px; 
        margin: 20px 0; 
        background: linear-gradient(135deg, #667eea 0%, #764ba2 100%);
        color: white;
        font-weight: 500;
    }
    .chat-box {
        border: 1px solid #e2e8f0;
        border-radius: 8px;
        padding: 15px;
        margin: 10px 0;
        background: #f8fafc;
        min-height: 100px;
    }
    .feature-badge {
        display: inline-block;
        padding: 4px 12px;
        border-radius: 20px;
        font-size: 0.85rem;
        font-weight: 600;
        margin: 4px;
        background: linear-gradient(135deg, #667eea 0%, #764ba2 100%);
        color: white;
    }
    .details-box {
        border: 1px solid #e2e8f0;
        border-radius: 8px;
        padding: 15px;
        background: #f8fafc;
    }
""") as iface:
    
    gr.Markdown(
        """
        # 🚀 Professional Website Builder AI
        ### Industry-Specific Website Generation with Tailwind CSS & Modern Animations
        
        Generate complete, production-ready websites tailored to your industry with our specialized AI team.
        Each website type comes with industry-specific components, design patterns, and content.
        """
    )
    
    with gr.Row():
        with gr.Column(scale=1):
            gr.Markdown("### 🎯 Step 1: Choose Website Type")
            
            website_type = gr.Radio(
                choices=list(WEBSITE_TYPES.keys()),
                value="SaaS Platform",
                label="Select Industry/Type",
                interactive=True
            )
            
            # Display type details in a Group instead of Box
            with gr.Group():
                type_details = gr.Markdown(
                    value=f"**{WEBSITE_TYPES['SaaS Platform']['description']}**\n\n"
                    f"Design Style: {WEBSITE_TYPES['SaaS Platform']['design_style']}\n\n"
                    f"Key Elements: {', '.join(WEBSITE_TYPES['SaaS Platform']['must_have_elements'][:3])}...",
                    elem_classes="details-box"
                )
            
        with gr.Column(scale=1):
            gr.Markdown("### 📝 Step 2: Project Details")
            
            name_method = gr.Radio(
                choices=["Generate with AI", "Enter manually"],
                value="Generate with AI",
                label="Project Name Method"
            )
            
            project_name_input = gr.Textbox(
                placeholder="Enter your project name (e.g., TechFlow Pro)",
                label="Project Name",
                visible=False
            )
            
            page_type = gr.Radio(
                choices=["Single-Page", "Multi-Page"],
                value="Single-Page",
                label="Website Structure"
            )
            
            custom_requirements = gr.Textbox(
                lines=3,
                placeholder="Any specific requirements? (e.g., 'Add a video hero section', 'Include testimonial carousel', 'Dark mode design')",
                label="Custom Requirements (Optional)"
            )
    
    with gr.Row():
        generate_btn = gr.Button(
            "🚀 Generate Professional Website", 
            variant="primary", 
            size="lg"
        )
    
    # Status and Preview
    status_output = gr.Markdown(elem_classes="status-box")
    
    with gr.Tabs():
        with gr.TabItem("🖥️ Live Preview"):
            preview_output = gr.HTML()
        
        with gr.TabItem("💬 Modify Website"):
            gr.Markdown(
                """
                ### Make Changes to Your Generated Website
                After generating a website, you can request modifications here. Your changes will update the live preview.
                """
            )
            modification_input = gr.Textbox(
                lines=3,
                placeholder="E.g., 'Change the hero background to a blue gradient', 'Add more testimonials', 'Make the navigation sticky'",
                label="What would you like to change?"
            )
            modify_btn = gr.Button("🔧 Apply Modifications", variant="secondary")
            modification_status = gr.Markdown() # This will now be the output for modification status
    
    # Download Section
    with gr.Row():
        download_single = gr.File(label="📥 Download HTML", visible=False)
        download_zip = gr.File(label="📦 Download Project (ZIP)", visible=False)
        view_code_btn = gr.Button("👁️ View Source Code", visible=False)
    
    # Code Display
    with gr.Row(visible=False) as code_row:
        code_output = gr.Code(
            label="Generated HTML Source Code",
            language="html",
            lines=20
        )
    
    # Feature showcase
    gr.Markdown(
        """
        ### ✨ Features by Website Type
        """
    )
    
    with gr.Row():
        for i, (type_name, config) in enumerate(list(WEBSITE_TYPES.items())[:3]):
            with gr.Column():
                gr.Markdown(
                    f"""
                    **{config['emoji']} {type_name}**
                    
                    {' '.join([f'<span class="feature-badge">{elem.split()[0]}</span>' for elem in config['must_have_elements'][:3]])}
                    """
                )
    
    # Event Handlers
    def update_type_details(selected_type):
        config = WEBSITE_TYPES[selected_type]
        return (
            f"**{config['description']}**\n\n"
            f"Design Style: {config['design_style']}\n\n"
            f"Fonts: {config['fonts']['heading']} / {config['fonts']['body']}\n\n"
            f"Key Elements: {', '.join(config['must_have_elements'][:3])}..."
        )
    
    def toggle_name_input(method):
        return gr.update(visible=(method == "Enter manually"))
    
    # Connect events
    website_type.change(
        fn=update_type_details,
        inputs=[website_type],
        outputs=[type_details]
    )
    
    name_method.change(
        fn=toggle_name_input,
        inputs=[name_method],
        outputs=[project_name_input]
    )
    
    generate_btn.click(
        fn=generate_website,
        inputs=[website_type, project_name_input, name_method, page_type, custom_requirements],
        outputs=[status_output, preview_output, download_single, download_zip, view_code_btn]
    )
    
    # **** FIX: Update click handler for the modify button ****
    modify_btn.click(
        fn=modify_website,
        inputs=[modification_input],
        # The outputs now include the preview and download components to be refreshed
        outputs=[modification_status, preview_output, download_single, download_zip, view_code_btn]
    )
    
    view_code_btn.click(
        fn=lambda: (
            gr.update(visible=True),
            open("outputs/index.html", "r", encoding="utf-8").read() if os.path.exists("outputs/index.html") else ""
        ),
        outputs=[code_row, code_output]
    )

if __name__ == "__main__":
    print("\n" + "="*50)
    print("🚀 Professional Website Builder AI")
    print("="*50)
    print(f"\n✅ {len(WEBSITE_TYPES)} website types loaded")
    print("🎨 Industry-specific design systems ready")
    print("💬 Chat modification system enabled")
    print("\n🌐 Launching interface...\n")
    
    iface.launch(share=False)