from flask import Flask, request, jsonify, make_response
from flask_cors import CORS
import anthropic
import json
import os
import xml.etree.ElementTree as ET
import html
from typing import Dict, Any, List
import psycopg2
from psycopg2.extras import RealDictCursor
import uuid
from datetime import datetime
import jwt
import requests
from functools import wraps

app = Flask(__name__)
CORS(app)  # Enable CORS for frontend communication

# Initialize Claude client
client = anthropic.Anthropic(
    api_key=os.getenv("CLAUDE_API_KEY", "your-claude-api-key-here")
)

# Database connection configuration
DB_CONFIG = {
    'host': os.getenv('POSTGRES_HOST', 'localhost'),
    'database': os.getenv('POSTGRES_DB', 'dev_pipeline'),
    'user': os.getenv('POSTGRES_USER', 'pipeline_admin'),
    'password': os.getenv('POSTGRES_PASSWORD', 'secure_pipeline_2024'),
    'port': os.getenv('POSTGRES_PORT', '5432')
}

# JWT and User Auth Service Configuration
JWT_SECRET = os.getenv('JWT_ACCESS_SECRET', 'access-secret-key-2024-tech4biz-secure_pipeline_2024')  # Use same secret as user-auth service
USER_AUTH_SERVICE_URL = os.getenv('USER_AUTH_SERVICE_URL', 'http://user-auth:8011')  # Use Docker service name
JWT_ALGORITHM = 'HS256'

def get_db_connection():
    """Get database connection"""
    try:
        conn = psycopg2.connect(**DB_CONFIG)
        return conn
    except Exception as e:
        print(f"Database connection error: {e}")
        return None

def verify_jwt_token(token):
    """Verify JWT token and return user data"""
    try:
        # First try to verify with local JWT secret (same as user-auth service)
        payload = jwt.decode(token, JWT_SECRET, algorithms=[JWT_ALGORITHM])
        return payload
    except jwt.ExpiredSignatureError:
        raise Exception("Token has expired")
    except jwt.InvalidTokenError:
        # If local verification fails, try to verify with user-auth service
        try:
            print(f"Local JWT verification failed, trying user-auth service at {USER_AUTH_SERVICE_URL}")
            response = requests.get(
                f"{USER_AUTH_SERVICE_URL}/api/auth/verify",
                headers={'Authorization': f'Bearer {token}'},
                timeout=5  # Reduced timeout
            )
            print(f"User-auth service response: {response.status_code}")
            
            if response.status_code == 200:
                result = response.json()
                if result.get('success') and result.get('data', {}).get('user'):
                    return result['data']['user']
                else:
                    raise Exception("Invalid response format from auth service")
            else:
                error_text = response.text
                print(f"Auth service error response: {error_text}")
                try:
                    error_json = response.json()
                    error_msg = error_json.get('error', 'Unknown error')
                except:
                    error_msg = error_text
                raise Exception(f"Auth service error: {error_msg}")
        except requests.RequestException as req_err:
            print(f"Request to auth service failed: {req_err}")
            # Don't raise exception, try to continue with local verification
            print("Continuing with local verification only")
            raise Exception(f"Unable to verify token with auth service: {str(req_err)}")
        except Exception as e:
            print(f"Error processing auth service response: {e}")
            raise Exception(f"Token verification failed: {str(e)}")
    except Exception as e:
        raise Exception(f"Token verification failed: {str(e)}")

def extract_user_id_from_token(user_data):
    """Extract user ID from various possible token formats"""
    if not user_data:
        return None
    
    # Try different possible user ID fields
    user_id = (user_data.get('id') or 
              user_data.get('userId') or 
              user_data.get('user_id') or
              user_data.get('sub') or  # JWT standard subject field
              user_data.get('user', {}).get('id'))  # Nested user object
    
    return user_id

def require_auth(f):
    """Decorator to require authentication"""
    @wraps(f)
    def decorated_function(*args, **kwargs):
        auth_header = request.headers.get('Authorization')
        if not auth_header:
            return jsonify({"error": "Authorization header required"}), 401
        
        try:
            token = auth_header.split(' ')[1] if ' ' in auth_header else auth_header
            user_data = verify_jwt_token(token)
            user_id = extract_user_id_from_token(user_data)
            
            if not user_id:
                return jsonify({"error": "Invalid token - no user ID found"}), 401
            
            # Store both user data and user ID for easy access
            request.user = user_data
            request.user_id = user_id
            return f(*args, **kwargs)
        except Exception as e:
            return jsonify({"error": str(e)}), 401
    
    return decorated_function

def optional_auth(f):
    """Decorator for optional authentication"""
    @wraps(f)
    def decorated_function(*args, **kwargs):
        auth_header = request.headers.get('Authorization')
        if auth_header:
            try:
                token = auth_header.split(' ')[1] if ' ' in auth_header else auth_header
                user_data = verify_jwt_token(token)
                user_id = extract_user_id_from_token(user_data)
                
                if user_id:
                    request.user = user_data
                    request.user_id = user_id
            except Exception:
                # Continue without authentication
                pass
        return f(*args, **kwargs)
    
    return decorated_function

# System prompt for SVG wireframe generation
SYSTEM_PROMPT = """You are an expert UI/UX designer and wireframe architect. Your task is to analyze user prompts and generate detailed wireframe specifications that can be converted to SVG.

Generate a JSON response with the following structure:
{
  "layout": {
    "page": {"width": 1200, "height": 800},
    "header": {"enabled": true, "height": 72, "elements": ["logo", "nav", "cta"]},
    "sidebar": {"enabled": false, "width": 240, "position": "left", "elements": []},
    "hero": {"enabled": false, "height": 200, "elements": []},
    "main_content": {
      "sections": [
        {
          "type": "grid",
          "rows": 2,
          "cols": 3,
          "height": 200,
          "elements": ["card1", "card2", "card3", "card4", "card5", "card6"]
        }
      ]
    },
    "footer": {"enabled": true, "height": 64, "elements": ["links", "copyright"]}
  },
  "styling": {
    "theme": "balsamiq",
    "colors": {"primary": "#3B82F6", "secondary": "#6B7280", "background": "#FFFFFF", "card": "#F8FAFC"},
    "spacing": {"gap": 16, "padding": 20}
  },
  "annotations": {
    "title": "Generated Wireframe",
    "description": "AI-generated wireframe based on user prompt"
  }
}

Focus on creating clean, Balsamiq-style wireframes with proper headings and descriptions.
Return only valid JSON, no additional text."""

def generate_svg_wireframe(layout_spec: Dict[str, Any], device_type: str = "desktop") -> str:
    """Generate SVG wireframe from layout specification"""
    
    # Device-specific dimensions
    device_dimensions = {
        "mobile": {"width": 375, "height": 812},
        "tablet": {"width": 768, "height": 1024},
        "desktop": {"width": 1440, "height": 1024}
    }
    
    # Get device dimensions or fallback to layout spec
    device_config = device_dimensions.get(device_type, device_dimensions["desktop"])
    page_width = int(device_config["width"])
    page_height = int(device_config["height"])
    
    # Create SVG root element
    svg = ET.Element('svg', {
        'width': str(page_width),
        'height': str(page_height),
        'viewBox': f'0 0 {page_width} {page_height}',
        'xmlns': 'http://www.w3.org/2000/svg',
        'style': 'font-family: "Comic Sans MS", cursive, sans-serif; background: #f8f9fa;'
    })
    
    # Add background
    ET.SubElement(svg, 'rect', {
        'x': '0', 'y': '0',
        'width': str(page_width), 'height': str(page_height),
        'fill': '#f8f9fa'
    })
    
    # Add definitions (filters, patterns)
    defs = ET.SubElement(svg, 'defs')
    
    # Hand-drawn style pattern
    pattern = ET.SubElement(defs, 'pattern', {
        'id': 'rough',
        'x': '0', 'y': '0',
        'width': '4', 'height': '4',
        'patternUnits': 'userSpaceOnUse'
    })
    ET.SubElement(pattern, 'path', {
        'd': 'M0,0 L2,2 M2,0 L4,2',
        'stroke': '#ddd',
        'stroke-width': '0.5'
    })
    
    # Create main group
    main_group = ET.SubElement(svg, 'g')
    
    # Get styling information
    styling = layout_spec.get('styling', {})
    colors = styling.get('colors', {})
    spacing = styling.get('spacing', {})
    
    # Balsamiq-style colors
    primary_color = '#2c3e50'
    secondary_color = '#7f8c8d'
    background_color = '#ffffff'
    card_color = '#ffffff'
    text_color = '#2c3e50'
    border_color = '#34495e'
    
    gap = int(spacing.get('gap', 16))
    padding = int(spacing.get('padding', 20))
    
    # Adjust spacing for mobile
    if device_type == "mobile":
        gap = max(8, gap // 2)
        padding = max(12, padding // 2)
    
    # Generate layout elements
    current_y = padding
    
    # Add device frame for mobile and tablet
    if device_type in ["mobile", "tablet"]:
        add_device_frame(main_group, page_width, page_height, device_type)
        # Adjust content area for frame
        if device_type == "mobile":
            padding = int(30)
            current_y = int(50)
        else:  # tablet
            padding = int(40)
            current_y = int(60)
    
    # Header
    header_spec = layout_spec.get('layout', {}).get('header', {})
    if header_spec.get('enabled', False):
        header_height = int(header_spec.get('height', 72))
        if device_type == "mobile":
            header_height = min(header_height, 60)
        generate_header(main_group, header_spec, page_width - (padding * 2), current_y, header_height, 
                       primary_color, background_color, text_color, border_color, padding, device_type)
        current_y += header_height + gap
    
    # Content area width calculation
    content_width = page_width - (padding * 2)
    content_x = padding
    
    # Sidebar (only for desktop and tablet)
    sidebar_spec = layout_spec.get('layout', {}).get('sidebar', {})
    if sidebar_spec.get('enabled', False) and device_type != "mobile":
        sidebar_width = int(sidebar_spec.get('width', 240))
        if device_type == "tablet":
            sidebar_width = min(sidebar_width, 200)
        footer_height = int(80 if layout_spec.get('layout', {}).get('footer', {}).get('enabled', False) else 0)
        sidebar_height = page_height - current_y - padding - footer_height
        generate_sidebar(main_group, sidebar_spec, padding, current_y, sidebar_width, sidebar_height,
                        secondary_color, background_color, text_color, border_color, device_type)
        content_x = padding + sidebar_width + gap
        content_width = page_width - content_x - padding
    
    # Hero section
    hero_spec = layout_spec.get('layout', {}).get('hero', {})
    if hero_spec.get('enabled', False):
        hero_height = int(hero_spec.get('height', 200))
        if device_type == "mobile":
            hero_height = min(hero_height, 150)
        generate_hero(main_group, hero_spec, content_x, current_y, content_width, hero_height,
                     primary_color, card_color, text_color, border_color, device_type)
        current_y += hero_height + gap
    
    # Main content sections
    main_content = layout_spec.get('layout', {}).get('main_content', {})
    sections = main_content.get('sections', [])
    
    for section in sections:
        section_height = int(section.get('height', 200))
        if device_type == "mobile":
            section_height = min(section_height, 300)
        generate_section(main_group, section, content_x, current_y, content_width, section_height,
                        primary_color, card_color, text_color, border_color, gap, device_type)
        current_y += section_height + gap
    
    # Footer
    footer_spec = layout_spec.get('layout', {}).get('footer', {})
    if footer_spec.get('enabled', False):
        footer_height = int(footer_spec.get('height', 64))
        if device_type == "mobile":
            footer_height = min(footer_height, 50)
        footer_y = page_height - footer_height - padding
        generate_footer(main_group, footer_spec, padding, footer_y, page_width - padding * 2, footer_height,
                       primary_color, background_color, text_color, border_color, device_type)
    
    # Add device label
    add_device_label(main_group, device_type, page_width, page_height)
    
    return ET.tostring(svg, encoding='unicode')

def get_wireframe_spec_from_claude(user_prompt: str, device_type: str) -> Dict[str, Any]:
    """Get wireframe specification from Claude API"""
    try:
        user_message = f"""Generate a wireframe specification for this prompt: "{user_prompt}"

Device type: {device_type}

Please analyze the requirements and create a detailed wireframe layout that follows modern UI/UX principles and Balsamiq wireframe styling. Consider the user's needs and create an intuitive, well-structured interface suitable for {device_type} devices."""

        response = client.messages.create(
            model="claude-3-5-sonnet-20241022",
            max_tokens=2000,
            system=SYSTEM_PROMPT,
            messages=[{"role": "user", "content": user_message}]
        )
        
        claude_response = response.content[0].text
        
        # Extract JSON from response
        json_start = claude_response.find('{')
        json_end = claude_response.rfind('}') + 1
        
        if json_start != -1 and json_end != -1:
            json_str = claude_response[json_start:json_end]
            json_str = json_str.replace('\n', ' ').replace('\r', ' ')
            return json.loads(json_str)
        else:
            print("No JSON found in Claude response, using fallback")
            return create_fallback_spec(user_prompt)
            
    except Exception as e:
        print(f"Claude API error: {e}")
        return create_fallback_spec(user_prompt)

def add_device_frame(group, width, height, device_type):
    """Add device frame for mobile and tablet"""
    if device_type == "mobile":
        # Mobile frame
        ET.SubElement(group, 'rect', {
            'x': '10', 'y': '10',
            'width': str(width - 20), 'height': str(height - 20),
            'fill': 'none',
            'stroke': '#2c3e50',
            'stroke-width': '3',
            'rx': '25'
        })
        # Home indicator
        ET.SubElement(group, 'rect', {
            'x': str(width // 2 - 30), 'y': str(height - 25),
            'width': '60', 'height': '4',
            'fill': '#2c3e50',
            'rx': '2'
        })
    elif device_type == "tablet":
        # Tablet frame
        ET.SubElement(group, 'rect', {
            'x': '15', 'y': '15',
            'width': str(width - 30), 'height': str(height - 30),
            'fill': 'none',
            'stroke': '#2c3e50',
            'stroke-width': '3',
            'rx': '15'
        })
        # Home button
        ET.SubElement(group, 'circle', {
            'cx': str(width // 2), 'cy': str(height - 35),
            'r': '8',
            'fill': 'none',
            'stroke': '#2c3e50',
            'stroke-width': '2'
        })

def safe_set_text(element, text):
    """Safely set text content on an XML element, escaping special characters"""
    if text:
        element.text = html.escape(str(text), quote=False)

def add_device_label(group, device_type, width, height):
    """Add device type label"""
    label_text = device_type.upper()
    text_element = ET.SubElement(group, 'text', {
        'x': str(width - 10), 'y': '25',
        'font-size': '12',
        'fill': '#7f8c8d',
        'text-anchor': 'end',
        'font-weight': 'bold'
    })
    safe_set_text(text_element, label_text)

def draw_rough_rect(group, x, y, width, height, fill, stroke, stroke_width="2"):
    """Draw a rough, hand-drawn style rectangle"""
    # Ensure all parameters are integers for mathematical operations
    x, y, width, height = int(x), int(y), int(width), int(height)
    
    # Create slightly irregular path for hand-drawn look
    path_data = f"M{x},{y} L{x+width-1},{y+1} L{x+width},{y+height-1} L{x+1},{y+height} Z"
    ET.SubElement(group, 'path', {
        'd': path_data,
        'fill': fill,
        'stroke': stroke,
        'stroke-width': stroke_width,
        'stroke-linecap': 'round',
        'stroke-linejoin': 'round'
    })

def generate_header(group, header_spec, width, y, height, primary_color, bg_color, text_color, border_color, x_offset, device_type):
    """Generate header section with Balsamiq style"""
    # Header background
    draw_rough_rect(group, x_offset, y, width, height, bg_color, border_color)
    
    # Header elements
    elements = header_spec.get('elements', [])
    if elements:
        if device_type == "mobile":
            # Mobile header layout
            # Logo on left
            logo_text = ET.SubElement(group, 'text', {
                'x': str(x_offset + 10), 'y': str(y + height//2 + 5),
                'font-size': '16',
                'fill': text_color,
                'font-weight': 'bold'
            })
            safe_set_text(logo_text, "LOGO")
            
            # Hamburger menu on right
            for i in range(3):
                ET.SubElement(group, 'line', {
                    'x1': str(x_offset + width - 30), 'y1': str(y + height//2 - 8 + i*6),
                    'x2': str(x_offset + width - 10), 'y2': str(y + height//2 - 8 + i*6),
                    'stroke': text_color,
                    'stroke-width': '2',
                    'stroke-linecap': 'round'
                })
        else:
            # Desktop/tablet header layout
            element_width = (width - 40) / len(elements)
            for i, element in enumerate(elements):
                element_x = x_offset + 20 + i * element_width
                element_text = str(element) if not isinstance(element, dict) else element.get('label', 'Element')
                
                if i == 0:  # First element is usually logo
                    draw_rough_rect(group, element_x - 5, y + 10, 80, height - 20, bg_color, border_color)
                    logo_element = ET.SubElement(group, 'text', {
                        'x': str(element_x + 35), 'y': str(y + height//2 + 5),
                        'font-size': '14',
                        'fill': text_color,
                        'font-weight': 'bold',
                        'text-anchor': 'middle'
                    })
                    safe_set_text(logo_element, element_text.upper())
                elif i == len(elements) - 1:  # Last element is usually CTA
                    draw_rough_rect(group, element_x, y + 10, 80, height - 20, primary_color, border_color)
                    cta_element = ET.SubElement(group, 'text', {
                        'x': str(element_x + 40), 'y': str(y + height//2 + 5),
                        'font-size': '12',
                        'fill': 'white',
                        'font-weight': 'bold',
                        'text-anchor': 'middle'
                    })
                    safe_set_text(cta_element, element_text)
                else:  # Navigation items
                    nav_element = ET.SubElement(group, 'text', {
                        'x': str(element_x), 'y': str(y + height//2 + 5),
                        'font-size': '14',
                        'fill': text_color,
                        'font-weight': '500'
                    })
                    safe_set_text(nav_element, element_text)

def generate_sidebar(group, sidebar_spec, x, y, width, height, primary_color, bg_color, text_color, border_color, device_type):
    """Generate sidebar section"""
    # Sidebar background
    draw_rough_rect(group, x, y, width, height, bg_color, border_color)
    
    # Title
    menu_title = ET.SubElement(group, 'text', {
        'x': str(x + 15), 'y': str(y + 25),
        'font-size': '16',
        'fill': text_color,
        'font-weight': 'bold'
    })
    safe_set_text(menu_title, "MENU")
    
    # Menu items
    elements = sidebar_spec.get('elements', ['Dashboard', 'Products', 'Orders', 'Customers', 'Settings'])
    for i, element in enumerate(elements[:8]):  # Limit to 8 items
        item_y = y + 50 + i * 35
        element_text = str(element) if not isinstance(element, dict) else element.get('label', 'Menu Item')
        
        # Menu item background
        if i == 0:  # Highlight first item as active
            draw_rough_rect(group, x + 5, item_y - 15, width - 10, 30, primary_color, border_color)
            text_color_item = 'white'
        else:
            text_color_item = text_color
        
        menu_item = ET.SubElement(group, 'text', {
            'x': str(x + 15), 'y': str(item_y),
            'font-size': '12',
            'fill': text_color_item
        })
        safe_set_text(menu_item, element_text)

def generate_hero(group, hero_spec, x, y, width, height, primary_color, bg_color, text_color, border_color, device_type):
    """Generate hero section"""
    # Hero background
    draw_rough_rect(group, x, y, width, height, bg_color, border_color)
    
    # Hero content layout
    if device_type == "mobile":
        # Mobile hero - stacked layout
        title_y = y + 40
        subtitle_y = y + 70
        cta_y = y + 110
    else:
        # Desktop/tablet hero - centered layout
        title_y = y + height//2 - 20
        subtitle_y = y + height//2
        cta_y = y + height//2 + 40
    
    # Title
    hero_title = ET.SubElement(group, 'text', {
        'x': str(x + width//2), 'y': str(title_y),
        'font-size': '24' if device_type != "mobile" else '20',
        'fill': text_color,
        'font-weight': 'bold',
        'text-anchor': 'middle'
    })
    safe_set_text(hero_title, "HERO TITLE")
    
    # Subtitle
    hero_subtitle = ET.SubElement(group, 'text', {
        'x': str(x + width//2), 'y': str(subtitle_y),
        'font-size': '14' if device_type != "mobile" else '12',
        'fill': text_color,
        'text-anchor': 'middle'
    })
    safe_set_text(hero_subtitle, "Your compelling subtitle goes here")
    
    # CTA Button
    button_width = 120 if device_type != "mobile" else 100
    button_x = x + width//2 - button_width//2
    draw_rough_rect(group, button_x, cta_y - 15, button_width, 35, primary_color, border_color)
    cta_button_text = ET.SubElement(group, 'text', {
        'x': str(x + width//2), 'y': str(cta_y + 5),
        'font-size': '14',
        'fill': 'white',
        'font-weight': 'bold',
        'text-anchor': 'middle'
    })
    safe_set_text(cta_button_text, "GET STARTED")

def generate_section(group, section, x, y, width, height, primary_color, card_color, text_color, border_color, gap, device_type):
    """Generate main content section"""
    section_type = section.get('type', 'generic')
    
    if section_type == 'grid':
        generate_grid_section(group, section, x, y, width, height, primary_color, card_color, text_color, border_color, gap, device_type)
    elif section_type == 'form':
        generate_form_section(group, section, x, y, width, height, primary_color, card_color, text_color, border_color, gap, device_type)
    elif section_type == 'two_column':
        generate_two_column_section(group, section, x, y, width, height, primary_color, card_color, text_color, border_color, gap, device_type)
    else:
        # Generic section
        draw_rough_rect(group, x, y, width, height, card_color, border_color)
        section_title = ET.SubElement(group, 'text', {
            'x': str(x + 10), 'y': str(y + 20),
            'font-size': '14',
            'fill': text_color,
            'font-weight': 'bold'
        })
        safe_set_text(section_title, section_type.upper().replace('_', ' '))

def generate_grid_section(group, section, x, y, width, height, primary_color, card_color, text_color, border_color, gap, device_type):
    """Generate grid section with Balsamiq style"""
    rows = int(section.get('rows', 2))
    cols = int(section.get('cols', 3))
    elements = section.get('elements', [])
    
    # Adjust grid for mobile
    if device_type == "mobile":
        cols = min(cols, 2)  # Max 2 columns on mobile
        if cols > 2:
            rows = (len(elements) + 1) // 2  # Recalculate rows
    
    cell_width = (width - gap * (cols - 1)) / cols
    cell_height = (height - gap * (rows - 1)) / rows
    
    for r in range(rows):
        for c in range(cols):
            cell_x = x + c * (cell_width + gap)
            cell_y = y + r * (cell_height + gap)
            element_index = r * cols + c
            
            # Cell background
            draw_rough_rect(group, cell_x, cell_y, cell_width, cell_height, card_color, border_color)
            
            # Cell content
            if element_index < len(elements):
                element_data = elements[element_index]
                element_text = str(element_data) if not isinstance(element_data, dict) else element_data.get('label', f'Card {element_index + 1}')
                
                # Card title
                card_title = ET.SubElement(group, 'text', {
                    'x': str(cell_x + 10), 'y': str(cell_y + 20),
                    'font-size': '14',
                    'fill': text_color,
                    'font-weight': 'bold'
                })
                safe_set_text(card_title, element_text)
                
                # Card description
                card_desc = ET.SubElement(group, 'text', {
                    'x': str(cell_x + 10), 'y': str(cell_y + 40),
                    'font-size': '10',
                    'fill': text_color
                })
                safe_set_text(card_desc, "Description text here")
                
                # Placeholder image area
                if cell_height > 80:
                    image_y = cell_y + 50
                    image_height = min(60, cell_height - 70)
                    draw_rough_rect(group, cell_x + 10, image_y, cell_width - 20, image_height, '#f1f2f6', border_color, "1")
                    # X marks for image placeholder
                    ET.SubElement(group, 'line', {
                        'x1': str(cell_x + 15), 'y1': str(image_y + 5),
                        'x2': str(cell_x + cell_width - 15), 'y2': str(image_y + image_height - 5),
                        'stroke': '#ddd',
                        'stroke-width': '2'
                    })
                    ET.SubElement(group, 'line', {
                        'x1': str(cell_x + cell_width - 15), 'y1': str(image_y + 5),
                        'x2': str(cell_x + 15), 'y2': str(image_y + image_height - 5),
                        'stroke': '#ddd',
                        'stroke-width': '2'
                    })

def generate_form_section(group, section, x, y, width, height, primary_color, card_color, text_color, border_color, gap, device_type):
    """Generate form section"""
    # Form background
    draw_rough_rect(group, x, y, width, height, card_color, border_color)
    
    # Form title
    form_title = ET.SubElement(group, 'text', {
        'x': str(x + 20), 'y': str(y + 30),
        'font-size': '18',
        'fill': text_color,
        'font-weight': 'bold'
    })
    safe_set_text(form_title, "CONTACT FORM")
    
    # Form fields
    fields = section.get('fields', ['name', 'email', 'message', 'submit'])
    field_height = 40
    field_gap = 15
    start_y = y + 50
    field_width = width - 40
    
    if device_type == "mobile":
        field_width = width - 20
    
    for i, field in enumerate(fields):
        field_y = start_y + i * (field_height + field_gap)
        field_text = str(field) if not isinstance(field, dict) else field.get('label', field)
        
        if field_text == 'submit':
            # Submit button
            button_width = 120
            button_x = x + 20
            draw_rough_rect(group, button_x, field_y, button_width, 40, primary_color, border_color)
            submit_text = ET.SubElement(group, 'text', {
                'x': str(button_x + 60), 'y': str(field_y + 25),
                'font-size': '14',
                'fill': 'white',
                'text-anchor': 'middle',
                'font-weight': 'bold'
            })
            safe_set_text(submit_text, 'SUBMIT')
        elif field_text == 'message':
            # Text area
            draw_rough_rect(group, x + 20, field_y, field_width, 80, 'white', border_color, "1")
            field_label = ET.SubElement(group, 'text', {
                'x': str(x + 30), 'y': str(field_y - 5),
                'font-size': '12',
                'fill': text_color,
                'font-weight': 'bold'
            })
            safe_set_text(field_label, field_text.upper())
            placeholder_text = ET.SubElement(group, 'text', {
                'x': str(x + 30), 'y': str(field_y + 20),
                'font-size': '10',
                'fill': '#7f8c8d'
            })
            safe_set_text(placeholder_text, "Your message here...")
        else:
            # Regular input field
            draw_rough_rect(group, x + 20, field_y, field_width, field_height, 'white', border_color, "1")
            field_label = ET.SubElement(group, 'text', {
                'x': str(x + 30), 'y': str(field_y - 5),
                'font-size': '12',
                'fill': text_color,
                'font-weight': 'bold'
            })
            safe_set_text(field_label, field_text.upper())
            placeholder_text = ET.SubElement(group, 'text', {
                'x': str(x + 30), 'y': str(field_y + 25),
                'font-size': '10',
                'fill': '#7f8c8d'
            })
            safe_set_text(placeholder_text, f"Enter your {field_text}...")

def generate_two_column_section(group, section, x, y, width, height, primary_color, card_color, text_color, border_color, gap, device_type):
    """Generate two-column section"""
    if device_type == "mobile":
        # Stack columns on mobile
        column_height = (height - gap) / 2
        
        # Top column
        draw_rough_rect(group, x, y, width, column_height, card_color, border_color)
        section1_text = ET.SubElement(group, 'text', {
            'x': str(x + 10), 'y': str(y + 20),
            'font-size': '14',
            'fill': text_color,
            'font-weight': 'bold'
        })
        safe_set_text(section1_text, 'SECTION 1')
        
        # Bottom column
        column2_y = y + column_height + gap
        draw_rough_rect(group, x, column2_y, width, column_height, card_color, border_color)
        section2_text = ET.SubElement(group, 'text', {
            'x': str(x + 10), 'y': str(column2_y + 20),
            'font-size': '14',
            'fill': text_color,
            'font-weight': 'bold'
        })
        safe_set_text(section2_text, 'SECTION 2')
    else:
        # Side-by-side columns
        column_width = (width - gap) / 2
        
        # Left column
        draw_rough_rect(group, x, y, column_width, height, card_color, border_color)
        left_text = ET.SubElement(group, 'text', {
            'x': str(x + 10), 'y': str(y + 20),
            'font-size': '14',
            'fill': text_color,
            'font-weight': 'bold'
        })
        safe_set_text(left_text, 'LEFT COLUMN')
        
        # Right column
        column2_x = x + column_width + gap
        draw_rough_rect(group, column2_x, y, column_width, height, card_color, border_color)
        right_text = ET.SubElement(group, 'text', {
            'x': str(column2_x + 10), 'y': str(y + 20),
            'font-size': '14',
            'fill': text_color,
            'font-weight': 'bold'
        })
        safe_set_text(right_text, 'RIGHT COLUMN')

def generate_footer(group, footer_spec, x, y, width, height, primary_color, bg_color, text_color, border_color, device_type):
    """Generate footer section"""
    # Footer background
    draw_rough_rect(group, x, y, width, height, bg_color, border_color)
    
    # Footer elements
    elements = footer_spec.get('elements', ['Links', 'About', 'Contact', 'Copyright'])
    if device_type == "mobile":
        # Stack footer elements vertically on mobile
        for i, element in enumerate(elements[:3]):  # Limit to 3 on mobile
            element_text = str(element) if not isinstance(element, dict) else element.get('label', 'Footer Item')
            footer_item = ET.SubElement(group, 'text', {
                'x': str(x + 10), 'y': str(y + 15 + i * 15),
                'font-size': '10',
                'fill': text_color
            })
            safe_set_text(footer_item, element_text)
    else:
        # Horizontal layout for desktop/tablet
        if elements:
            element_width = (width - 40) / len(elements)
            for i, element in enumerate(elements):
                element_x = x + 20 + i * element_width
                element_text = str(element) if not isinstance(element, dict) else element.get('label', 'Footer Item')
                footer_item = ET.SubElement(group, 'text', {
                    'x': str(element_x), 'y': str(y + height//2 + 5),
                    'font-size': '12',
                    'fill': text_color
                })
                safe_set_text(footer_item, element_text)

def create_fallback_spec(prompt: str) -> Dict[str, Any]:
    """Create a fallback wireframe specification when Claude fails"""
    prompt_lower = prompt.lower()
    
    # Check for specific wireframe types
    is_ecommerce = any(word in prompt_lower for word in ['e-commerce', 'ecommerce', 'product', 'shop', 'store'])
    is_dashboard = any(word in prompt_lower for word in ['dashboard', 'admin', 'analytics'])
    is_landing = any(word in prompt_lower for word in ['landing', 'hero', 'banner'])
    has_form = any(word in prompt_lower for word in ['form', 'signup', 'login', 'contact'])
    
    # Basic fallback logic
    has_header = any(word in prompt_lower for word in ['header', 'nav', 'navbar'])
    has_sidebar = any(word in prompt_lower for word in ['sidebar', 'left nav', 'drawer'])
    has_hero = any(word in prompt_lower for word in ['hero', 'banner', 'jumbotron'])
    has_footer = any(word in prompt_lower for word in ['footer', 'bottom'])
    
    # Extract grid dimensions if specified
    grid_rows, grid_cols = 2, 3  # Default
    if 'x' in prompt_lower:
        try:
            parts = prompt_lower.split('x')
            if len(parts) == 2:
                grid_cols = int(parts[0].strip().split()[-1])
                grid_rows = int(parts[1].strip().split()[0])
        except:
            pass
    
    # Special handling for e-commerce product pages
    if is_ecommerce and 'product' in prompt_lower:
        return {
            "layout": {
                "page": {"width": 1440, "height": 1200},
                "header": {"enabled": True, "height": 80, "elements": ["logo", "search", "nav", "cart", "account"]},
                "sidebar": {"enabled": False, "width": 240, "position": "left", "elements": []},
                "hero": {"enabled": False, "height": 200, "elements": []},
                "main_content": {
                    "sections": [
                        {
                            "type": "two_column",
                            "height": 600,
                            "elements": ["product_gallery", "product_details"]
                        },
                        {
                            "type": "grid",
                            "rows": 1,
                            "cols": 4,
                            "height": 200,
                            "elements": ["specifications", "shipping", "reviews", "related"]
                        }
                    ]
                },
                "footer": {"enabled": True, "height": 120, "elements": ["company_info", "customer_service", "legal", "social"]}
            },
            "styling": {
                "theme": "balsamiq",
                "colors": {"primary": "#4F46E5", "secondary": "#64748B", "accent": "#EF4444", "background": "#FFFFFF", "text": "#1F2937"},
                "spacing": {"gap": 24, "padding": 32}
            },
            "annotations": {
                "title": "E-commerce Product Page Wireframe",
                "description": "Product detail page with gallery and product information"
            }
        }
    
    # Special handling for dashboards
    elif is_dashboard:
        return {
            "layout": {
                "page": {"width": 1440, "height": 1000},
                "header": {"enabled": True, "height": 80, "elements": ["logo", "nav", "user_menu", "notifications"]},
                "sidebar": {"enabled": True, "width": 280, "position": "left", "elements": ["dashboard", "analytics", "reports", "settings"]},
                "hero": {"enabled": False, "height": 200, "elements": []},
                "main_content": {
                    "sections": [
                        {
                            "type": "grid",
                            "rows": 2,
                            "cols": 4,
                            "height": 300,
                            "elements": ["stats_1", "stats_2", "stats_3", "stats_4", "chart_1", "chart_2", "table_1", "activity"]
                        }
                    ]
                },
                "footer": {"enabled": False, "height": 64, "elements": []}
            },
            "styling": {
                "theme": "balsamiq",
                "colors": {"primary": "#3B82F6", "secondary": "#6B7280", "accent": "#10B981", "background": "#F9FAFB", "text": "#1F2937"},
                "spacing": {"gap": 24, "padding": 32}
            },
            "annotations": {
                "title": "Dashboard Wireframe",
                "description": "Analytics dashboard with stats cards and charts"
            }
        }
    
    # Special handling for landing pages
    elif is_landing:
        return {
            "layout": {
                "page": {"width": 1440, "height": 1000},
                "header": {"enabled": True, "height": 80, "elements": ["logo", "nav", "cta_button"]},
                "sidebar": {"enabled": False, "width": 240, "position": "left", "elements": []},
                "hero": {"enabled": True, "height": 400, "elements": ["title", "subtitle", "cta_button", "hero_image"]},
                "main_content": {
                    "sections": [
                        {
                            "type": "grid",
                            "rows": 2,
                            "cols": 3,
                            "height": 300,
                            "elements": ["feature_1", "feature_2", "feature_3", "testimonial_1", "testimonial_2", "testimonial_3"]
                        }
                    ]
                },
                "footer": {"enabled": True, "height": 120, "elements": ["company", "product", "support", "legal"]}
            },
            "styling": {
                "theme": "balsamiq",
                "colors": {"primary": "#3B82F6", "secondary": "#6B7280", "accent": "#F59E0B", "background": "#FFFFFF", "text": "#1F2937"},
                "spacing": {"gap": 32, "padding": 40}
            },
            "annotations": {
                "title": "Landing Page Wireframe",
                "description": "Modern landing page with hero section and features"
            }
        }
    
    # Default fallback
    return {
        "layout": {
            "page": {"width": 1440, "height": 800},
            "header": {"enabled": has_header or True, "height": 80, "elements": ["logo", "navigation", "cta"]},
            "sidebar": {"enabled": has_sidebar, "width": 280, "position": "left", "elements": ["menu", "filters", "options"]},
            "hero": {"enabled": has_hero, "height": 200, "elements": ["title", "subtitle", "cta"]},
            "main_content": {
                "sections": [
                    {
                        "type": "grid",
                        "rows": grid_rows,
                        "cols": grid_cols,
                        "height": 400,
                        "elements": [f"card_{i+1}" for i in range(grid_rows * grid_cols)]
                    }
                ] + ([{
                    "type": "form",
                    "height": 300,
                    "fields": ["name", "email", "message", "submit"]
                }] if has_form else [])
            },
            "footer": {"enabled": has_footer or True, "height": 80, "elements": ["links", "about", "contact", "copyright"]}
        },
        "styling": {
            "theme": "balsamiq",
            "colors": {"primary": "#3B82F6", "secondary": "#6B7280", "background": "#FFFFFF", "text": "#1F2937"},
            "spacing": {"gap": 20, "padding": 24}
        },
        "annotations": {
            "title": "Custom Wireframe",
            "description": f"Generated from: {prompt[:100]}..."
        }
    }

# MAIN ENDPOINTS - Solution 1: Device-specific endpoints that return pure SVG

@app.route('/generate-wireframe/<device_type>', methods=['POST'])
def generate_wireframe_by_device(device_type):
    """Generate wireframe for specific device type, returns SVG directly"""
    try:
        data = request.get_json()
        user_prompt = data.get('prompt', '')
        
        if not user_prompt:
            return jsonify({"error": "No prompt provided"}), 400
        
        if device_type not in ['desktop', 'tablet', 'mobile']:
            return jsonify({"error": "Invalid device type. Use: desktop, tablet, or mobile"}), 400
        
        # Get wireframe spec from Claude
        wireframe_spec = get_wireframe_spec_from_claude(user_prompt, device_type)
        
        # Generate SVG
        svg_content = generate_svg_wireframe(wireframe_spec, device_type)
        
        # Return SVG directly with proper headers
        response = make_response(svg_content)
        response.headers['Content-Type'] = 'image/svg+xml'
        response.headers['Cache-Control'] = 'no-cache'
        response.headers['Access-Control-Allow-Origin'] = '*'
        return response
        
    except Exception as e:
        print(f"Error generating {device_type} wireframe: {e}")
        return jsonify({"error": str(e)}), 500

# Keep original endpoint for backward compatibility
@app.route('/generate-wireframe', methods=['POST'])
def generate_wireframe():
    """Original endpoint - defaults to desktop"""
    try:
        data = request.get_json()
        user_prompt = data.get('prompt', '')
        device_type = data.get('device', 'desktop')
        
        if not user_prompt:
            return jsonify({"error": "No prompt provided"}), 400
        
        if device_type not in ['desktop', 'tablet', 'mobile']:
            device_type = 'desktop'
        
        # Get wireframe spec from Claude
        wireframe_spec = get_wireframe_spec_from_claude(user_prompt, device_type)
        
        # Generate SVG
        svg_content = generate_svg_wireframe(wireframe_spec, device_type)
        
        # Return SVG directly
        response = make_response(svg_content)
        response.headers['Content-Type'] = 'image/svg+xml'
        response.headers['Cache-Control'] = 'no-cache'
        return response
        
    except Exception as e:
        print(f"Error generating wireframe: {e}")
        return jsonify({"error": str(e)}), 500

# Optional: Endpoint to get all device wireframes (if you still need this)
@app.route('/generate-all-devices', methods=['POST'])
def generate_all_devices():
    """Generate wireframes for all devices and return URLs to fetch them"""
    try:
        data = request.get_json()
        user_prompt = data.get('prompt', '')
        
        if not user_prompt:
            return jsonify({"error": "No prompt provided"}), 400
        
        # Return metadata about where to fetch each device wireframe
        devices = ['desktop', 'tablet', 'mobile']
        device_urls = {}
        
        for device in devices:
            device_urls[device] = f"/generate-wireframe/{device}"
        
        return jsonify({
            "success": True,
            "message": "Use the provided URLs to fetch SVG wireframes for each device",
            "device_endpoints": device_urls,
            "usage": "POST to each endpoint with {'prompt': 'your prompt'} to get SVG",
            "original_prompt": user_prompt
        })
        
    except Exception as e:
        return jsonify({"error": str(e)}), 500

# Wireframe Persistence Endpoints
@app.route('/api/wireframes', methods=['POST'])
@require_auth
def save_wireframe():
    """Save wireframe data to database"""
    try:
        data = request.get_json()
        
        # Extract wireframe data
        wireframe_data = data.get('wireframe', {})
        elements_data = data.get('elements', [])
        
        # Get user_id from JWT token - handle both local and remote verification
        user_id = None
        if hasattr(request, 'user_id'):
            user_id = request.user_id
        
        if not user_id:
            return jsonify({"error": "User authentication required - no valid user ID found"}), 401
        
        project_id = data.get('project_id')
        name = wireframe_data.get('name', 'Untitled Wireframe')
        description = wireframe_data.get('description', '')
        device_type = wireframe_data.get('device_type', 'desktop')
        dimensions = wireframe_data.get('dimensions', {'width': 1440, 'height': 1024})
        metadata = wireframe_data.get('metadata', {})
        
        if not user_id:
            return jsonify({"error": "User authentication required"}), 401
        
        conn = get_db_connection()
        if not conn:
            return jsonify({"error": "Database connection failed"}), 500
        
        try:
            with conn.cursor(cursor_factory=RealDictCursor) as cur:
                # Check if this is an update (wireframe has an ID) or new creation
                wireframe_id = wireframe_data.get('id')
                
                if wireframe_id:
                    # Update existing wireframe
                    cur.execute("""
                        UPDATE wireframes 
                        SET name = %s, description = %s, device_type = %s, dimensions = %s, metadata = %s, updated_at = NOW()
                        WHERE id = %s AND user_id = %s
                        RETURNING id
                    """, (name, description, device_type, json.dumps(dimensions), json.dumps(metadata), wireframe_id, user_id))
                    
                    if cur.rowcount == 0:
                        return jsonify({"error": "Wireframe not found or access denied"}), 404
                    
                    # Clear existing elements
                    cur.execute("DELETE FROM wireframe_elements WHERE wireframe_id = %s", (wireframe_id,))
                    
                else:
                    # Create new wireframe
                    cur.execute("""
                        INSERT INTO wireframes (user_id, project_id, name, description, device_type, dimensions, metadata)
                        VALUES (%s, %s, %s, %s, %s, %s, %s)
                        RETURNING id
                    """, (user_id, project_id, name, description, device_type, json.dumps(dimensions), json.dumps(metadata)))
                    
                    wireframe_id = cur.fetchone()['id']
                
                # Insert elements
                for element in elements_data:
                    # Extract TLDraw ID from the element data
                    tldraw_id = element.get('id') or element.get('data', {}).get('id')
                    
                    cur.execute("""
                        INSERT INTO wireframe_elements (wireframe_id, element_type, element_data, position, size, style, parent_id, tldraw_id, z_index)
                        VALUES (%s, %s, %s, %s, %s, %s, %s, %s, %s)
                    """, (
                        wireframe_id,
                        element.get('type', 'shape'),
                        json.dumps(element.get('data', {})),
                        json.dumps(element.get('position', {'x': 0, 'y': 0})),
                        json.dumps(element.get('size', {'width': 100, 'height': 100})),
                        json.dumps(element.get('style', {})),
                        element.get('parent_id'),  # This can be a string ID from TLDraw
                        tldraw_id,  # Store the original TLDraw ID
                        element.get('z_index', 0)
                    ))
                
                # Create version
                version_name = 'Updated Version' if wireframe_data.get('id') else 'Initial Version'
                version_description = 'Wireframe updated' if wireframe_data.get('id') else 'Initial wireframe creation'
                
                cur.execute("""
                    SELECT create_wireframe_version(%s, %s, %s, %s, %s)
                """, (
                    wireframe_id,
                    version_name,
                    version_description,
                    json.dumps({'wireframe_id': str(wireframe_id), 'elements': elements_data}),
                    user_id
                ))
                
                conn.commit()
                
                return jsonify({
                    "success": True,
                    "wireframe_id": str(wireframe_id),
                    "message": "Wireframe saved successfully"
                }), 201
                
        except Exception as e:
            conn.rollback()
            raise e
        finally:
            conn.close()
            
    except Exception as e:
        print(f"Error saving wireframe: {e}")
        return jsonify({"error": str(e)}), 500

@app.route('/api/wireframes/<wireframe_id>', methods=['GET'])
@require_auth
def get_wireframe(wireframe_id):
    """Get wireframe data from database"""
    try:
        # Get user_id from JWT token - handle both local and remote verification
        user_id = None
        if hasattr(request, 'user_id'):
            user_id = request.user_id
        
        if not user_id:
            return jsonify({"error": "User authentication required - no valid user ID found"}), 401
        
        conn = get_db_connection()
        if not conn:
            return jsonify({"error": "Database connection failed"}), 500
        
        try:
            with conn.cursor(cursor_factory=RealDictCursor) as cur:
                # First check if user owns this wireframe
                cur.execute("""
                    SELECT user_id FROM wireframes WHERE id = %s
                """, (wireframe_id,))
                
                wireframe_owner = cur.fetchone()
                if not wireframe_owner:
                    return jsonify({"error": "Wireframe not found"}), 404
                
                if wireframe_owner['user_id'] != user_id:
                    return jsonify({"error": "Access denied"}), 403
                
                # Get wireframe with elements using the function
                cur.execute("SELECT * FROM get_wireframe_with_elements(%s)", (wireframe_id,))
                result = cur.fetchone()
                
                if not result:
                    return jsonify({"error": "Wireframe not found"}), 404
                
                return jsonify({
                    "success": True,
                    "wireframe": result['wireframe_data'],
                    "elements": result['elements_data']
                }), 200
                
        finally:
            conn.close()
            
    except Exception as e:
        print(f"Error getting wireframe: {e}")
        return jsonify({"error": str(e)}), 500

@app.route('/api/wireframes/user/<user_id>', methods=['GET'])
@require_auth
def get_user_wireframes(user_id):
    """Get all wireframes for a user"""
    try:
        # Get user_id from JWT token - handle both local and remote verification
        authenticated_user_id = None
        if hasattr(request, 'user_id'):
            authenticated_user_id = request.user_id
        
        if not authenticated_user_id:
            return jsonify({"error": "User authentication required - no valid user ID found"}), 401
        
        # Users can only access their own wireframes
        if authenticated_user_id != user_id:
            return jsonify({"error": "Access denied"}), 403
        
        conn = get_db_connection()
        if not conn:
            return jsonify({"error": "Database connection failed"}), 500
        
        try:
            with conn.cursor(cursor_factory=RealDictCursor) as cur:
                # Get wireframes for the authenticated user
                cur.execute("""
                    SELECT id, name, description, device_type, dimensions, metadata, created_at, updated_at
                    FROM wireframes 
                    WHERE user_id = %s AND is_active = true
                    ORDER BY updated_at DESC
                """, (authenticated_user_id,))
                
                wireframes = cur.fetchall()
                
                return jsonify({
                    "success": True,
                    "wireframes": wireframes,
                    "count": len(wireframes)
                }), 200
                
        finally:
            conn.close()
            
    except Exception as e:
        print(f"Error getting user wireframes: {e}")
        return jsonify({"error": str(e)}), 500

@app.route('/api/wireframes/<wireframe_id>', methods=['PUT'])
@require_auth
def update_wireframe(wireframe_id):
    """Update existing wireframe data"""
    try:
        # Get user_id from JWT token - handle both local and remote verification
        user_id = None
        if hasattr(request, 'user_id'):
            user_id = request.user_id
        
        if not user_id:
            return jsonify({"error": "User authentication required - no valid user ID found"}), 401
        
        data = request.get_json()
        
        conn = get_db_connection()
        if not conn:
            return jsonify({"error": "Database connection failed"}), 500
        
        try:
            with conn.cursor(cursor_factory=RealDictCursor) as cur:
                # First check if user owns this wireframe
                cur.execute("""
                    SELECT user_id FROM wireframes WHERE id = %s
                """, (wireframe_id,))
                
                wireframe_owner = cur.fetchone()
                if not wireframe_owner:
                    return jsonify({"error": "Wireframe not found"}), 404
                
                if wireframe_owner['user_id'] != user_id:
                    return jsonify({"error": "Access denied"}), 403
                
                # Update wireframe metadata
                name = data.get('name', 'Untitled Wireframe')
                description = data.get('description', '')
                device_type = data.get('device_type', 'desktop')
                dimensions = data.get('dimensions', {'width': 1440, 'height': 1024})
                metadata = data.get('metadata', {})
                
                cur.execute("""
                    UPDATE wireframes 
                    SET name = %s, description = %s, device_type = %s, dimensions = %s, metadata = %s, updated_at = NOW()
                    WHERE id = %s
                """, (name, description, device_type, json.dumps(dimensions), json.dumps(metadata), wireframe_id))
                
                # Delete existing elements
                cur.execute("DELETE FROM wireframe_elements WHERE wireframe_id = %s", (wireframe_id,))
                
                # Insert new elements
                elements_data = data.get('elements', [])
                for element in elements_data:
                    cur.execute("""
                        INSERT INTO wireframe_elements (wireframe_id, element_type, element_data, position, size, style, parent_id, z_index)
                        VALUES (%s, %s, %s, %s, %s, %s, %s, %s)
                    """, (
                        wireframe_id,
                        element.get('type', 'shape'),
                        json.dumps(element.get('data', {})),
                        json.dumps(element.get('position', {'x': 0, 'y': 0})),
                        json.dumps(element.get('size', {'width': 100, 'height': 100})),
                        json.dumps(element.get('style', {})),
                        element.get('parent_id'),
                        element.get('z_index', 0)
                    ))
                
                # Create new version
                cur.execute("""
                    SELECT create_wireframe_version(%s, %s, %s, %s, %s)
                """, (
                    wireframe_id,
                    'Updated Version',
                    'Wireframe updated',
                    json.dumps({'wireframe_id': str(wireframe_id), 'elements': elements_data}),
                    user_id
                ))
                
                conn.commit()
                
                return jsonify({
                    "success": True,
                    "wireframe_id": str(wireframe_id),
                    "message": "Wireframe updated successfully"
                }), 200
                
        except Exception as e:
            conn.rollback()
            raise e
        finally:
            conn.close()
            
    except Exception as e:
        print(f"Error updating wireframe: {e}")
        return jsonify({"error": str(e)}), 500

@app.route('/api/wireframes/<wireframe_id>', methods=['DELETE'])
@require_auth
def delete_wireframe(wireframe_id):
    """Delete a wireframe (soft delete)"""
    try:
        user_id = request.user_id
        
        conn = get_db_connection()
        if not conn:
            return jsonify({"error": "Database connection failed"}), 500
        
        try:
            with conn.cursor(cursor_factory=RealDictCursor) as cur:
                # First check if user owns this wireframe
                cur.execute("""
                    SELECT user_id FROM wireframes WHERE id = %s
                """, (wireframe_id,))
                
                wireframe_owner = cur.fetchone()
                if not wireframe_owner:
                    return jsonify({"error": "Wireframe not found"}), 404
                
                if wireframe_owner['user_id'] != user_id:
                    return jsonify({"error": "Access denied"}), 403
                
                # Soft delete wireframe
                cur.execute("""
                    UPDATE wireframes SET is_active = false, updated_at = NOW() WHERE id = %s
                """, (wireframe_id,))
                
                conn.commit()
                
                return jsonify({
                    "success": True,
                    "message": "Wireframe deleted successfully"
                }), 200
                
        finally:
            conn.close()
            
    except Exception as e:
        print(f"Error deleting wireframe: {e}")
        return jsonify({"error": str(e)}), 500

# ========================================
# MAIN APPLICATION
# ========================================

@app.route('/health', methods=['GET'])
def health_check():
    """Health check endpoint"""
    try:
        # Test database connection
        conn = get_db_connection()
        db_status = 'connected' if conn else 'disconnected'
        if conn:
            conn.close()
        
        # Test user-auth service connection
        auth_status = 'connected'
        try:
            response = requests.get(f"{USER_AUTH_SERVICE_URL}/health", timeout=5)
            if response.status_code != 200:
                auth_status = 'error'
        except:
            auth_status = 'disconnected'
        
        return jsonify({
            "status": "healthy",
            "service": "ai-mockup-service",
            "version": "1.0.0",
            "timestamp": datetime.now().isoformat(),
            "features": {
                "wireframe_generation": True,
                "authentication": True,
                "real_time_updates": True,
                "user_isolation": True
            },
            "services": {
                "database": db_status,
                "user_auth": auth_status
            },
            "environment": os.getenv('FLASK_ENV', 'development')
        }), 200
        
    except Exception as e:
        return jsonify({
            "status": "unhealthy",
            "error": str(e)
        }), 500

@app.route('/', methods=['GET'])
def root():
    """Root endpoint with API documentation"""
    return jsonify({
        "message": "AI Mockup Service - Wireframe Generation with Authentication",
        "version": "1.0.0",
        "authentication": "JWT Bearer Token Required",
        "endpoints": {
            "health": "GET /health",
            "generate_wireframe": "POST /generate-wireframe",
            "generate_wireframe_desktop": "POST /generate-wireframe/desktop",
            "generate_wireframe_tablet": "POST /generate-wireframe/tablet", 
            "generate_wireframe_mobile": "POST /generate-wireframe/mobile",
            "generate_all_devices": "POST /generate-all-devices",
            "save_wireframe": "POST /api/wireframes (Auth Required)",
            "get_wireframe": "GET /api/wireframes/<id> (Auth Required)",
            "update_wireframe": "PUT /api/wireframes/<id> (Auth Required)",
            "delete_wireframe": "DELETE /api/wireframes/<id> (Auth Required)",
            "get_user_wireframes": "GET /api/wireframes/user/<user_id> (Auth Required)"
        },
        "authentication": {
            "type": "JWT Bearer Token",
            "header": "Authorization: Bearer <token>",
            "service": USER_AUTH_SERVICE_URL
        }
    }), 200

if __name__ == '__main__':
    # Load environment variables
    try:
        from dotenv import load_dotenv
        load_dotenv()
    except ImportError:
        print("python-dotenv not installed, skipping .env file loading")
    
    port = int(os.environ.get('PORT', 8021))
    app.run(debug=True, host='0.0.0.0', port=port)