#!/usr/bin/env python3
"""
Atomic Design Extraction Script
Extracts components from monolithic UI Kit into Atomic Design structure.
"""

import json
import os
import sys
from pathlib import Path

def extract_atomic():
    # Configuration
    source_kit = 'src/components/tPOS-ui-kit.pen'
    output_base = 'src'
    
    if not os.path.exists(source_kit):
        # Try finding in root if not in src/components
        if os.path.exists('tPOS-ui-kit.pen'):
            source_kit = 'tPOS-ui-kit.pen'
        else:
            print(f"❌ Source file not found: {source_kit}")
            sys.exit(1)

    print(f"🔨 Extracting from {source_kit}...")

    # Read monolithic UI Kit
    with open(source_kit, 'r') as f:
        data = json.load(f)

    variables = data.get('variables', {})
    
    # Find Component Library page (usually the first child)
    children = data.get('children', [])
    if not children:
        print("❌ No pages found in source file")
        sys.exit(1)
        
    # Assume first page is component library, or find it by name
    main_page = children[0]
    for child in children:
        if "library" in child.get("name", "").lower() or "kit" in child.get("name", "").lower():
            main_page = child
            break
            
    sections = main_page.get('children', [])
    print(f"✓ Found library page: {main_page.get('name')}")

    # 1. Create directory structure
    dirs = ['foundations', 'atoms', 'molecules', 'organisms']
    for d in dirs:
        os.makedirs(os.path.join(output_base, d), exist_ok=True)

    # 2. Extract design tokens
    tokens_file = {
        'version': '2.6',
        'children': [{
            'type': 'frame',
            'name': 'Design Tokens Reference',
            'layout': 'vertical',
            'gap': 20,
            'children': [
                {
                    'type': 'text',
                    'id': 'tokens_title',
                    'name': 'Title',
                    'content': 'Design System Tokens',
                    'fontSize': 32,
                    'fill': '$text-primary'
                }
            ]
        }],
        'variables': variables
    }

    with open(os.path.join(output_base, 'foundations', 'design-tokens.pen'), 'w') as f:
        json.dump(tokens_file, f, indent=2)
    print(f"✅ Created foundations/design-tokens.pen")

    # 3. Categorize components by atomic level
    atomic_mapping = {
        'atoms': {
            'patterns': ['button', 'badge', 'input', 'icon', 'typography', 'chip', 'avatar', 'divider'],
            'components': []
        },
        'molecules': {
            'patterns': ['formfield', 'searchbar', 'cardheader', 'listitem', 'row', 'field', 'dropdown'],
            'components': []
        },
        'organisms': {
            'patterns': ['card', 'navigation', 'header', 'footer', 'form', 'modal', 'dialog', 'sidebar', 'table'],
            'components': []
        }
    }

    # Helper to check if name matches patterns
    def matches_pattern(name, patterns):
        normalize_name = name.lower()
        for p in patterns:
            # Check for exact word matches or clear indicators
            if p in normalize_name:
                return True
        return False

    # Categorize each section/component
    uncategorized = []
    
    for section in sections:
        section_name = section.get('name', '').lower()
        
        # Try to guess based on name
        matched = False
        
        # Check explicit naming first (e.g. "Atom/...")
        if "atom" in section_name:
            atomic_mapping['atoms']['components'].append(section)
            continue
        if "molecule" in section_name:
            atomic_mapping['molecules']['components'].append(section)
            continue
        if "organism" in section_name:
            atomic_mapping['organisms']['components'].append(section)
            continue

        # Check patterns
        if matches_pattern(section_name, atomic_mapping['atoms']['patterns']):
            atomic_mapping['atoms']['components'].append(section)
            matched = True
        elif matches_pattern(section_name, atomic_mapping['molecules']['patterns']):
            atomic_mapping['molecules']['components'].append(section)
            matched = True
        elif matches_pattern(section_name, atomic_mapping['organisms']['patterns']):
            atomic_mapping['organisms']['components'].append(section)
            matched = True
            
        if not matched:
            # Default to organisms for complex unknown things, or store as uncategorized
            uncategorized.append(section)

    # 4. Generate atomic files
    for level, config in atomic_mapping.items():
        if not config['components']:
            continue
        
        # Group by component type
        grouped = {}
        for component in config['components']:
            name = component.get('name', 'Unknown')
            # Extract type (e.g., "Buttons & Actions Section" -> "buttons")
            # Simple heuristic: use the first word or the matched pattern
            comp_type = "misc"
            
            # Try to find which pattern matched
            name_lower = name.lower()
            for p in config['patterns']:
                if p in name_lower:
                    comp_type = p
                    break
            
            # Fallback to first word if no pattern matched (for manual assignments)
            if comp_type == "misc":
                comp_type = name.split()[0].lower().replace('&', '').strip()

            if comp_type not in grouped:
                grouped[comp_type] = []
            grouped[comp_type].append(component)
        
        # Create file for each type
        for comp_type, components in grouped.items():
            # Add a showcase wrapper
            showcase_wrapper = {
                'type': 'frame',
                'name': f'{comp_type.title()} Showcase',
                'width': 1440,
                'fill': '$bg-page',
                'layout': 'vertical',
                'gap': 80,
                'padding': [80, 120],
                'children': components
            }

            output_file = {
                'version': '2.6',
                'children': [showcase_wrapper],
                'variables': variables
            }
            
            # Pluralize filename
            filename = f"{comp_type}s.pen" if not comp_type.endswith('s') else f"{comp_type}.pen"
            filepath = os.path.join(output_base, level, filename)
            
            with open(filepath, 'w') as f:
                json.dump(output_file, f, indent=2)
            
            print(f'✅ Created {filepath} ({len(components)} components)')

    # Handle uncategorized
    if uncategorized:
        print(f"\n⚠️  {len(uncategorized)} components could not be automatically categorized:")
        for u in uncategorized:
            print(f"  - {u.get('name')}")
        
        # Save them to a 'misc' organism file
        filepath = os.path.join(output_base, 'organisms', 'misc.pen')
        with open(filepath, 'w') as f:
            json.dump({
                'version': '2.6',
                'children': uncategorized,
                'variables': variables
            }, f, indent=2)
        print(f"✅ Saved uncategorized to {filepath}")

    print(f'\n🎉 Atomic Design extraction complete!')

if __name__ == "__main__":
    extract_atomic()