#!/usr/bin/env python3
"""
Pencil Design Build System
Merges component library into page files
"""

import json
import sys
from pathlib import Path
from typing import Dict, Any, List

class PencilBuilder:
    def __init__(self, config_path: str = "pencil.config.json"):
        """Initialize builder with config"""
        # Get script directory and project root
        self.script_dir = Path(__file__).parent.absolute()
        self.project_root = self.script_dir.parent
        
        # Load config from project root
        config_full_path = self.project_root / config_path
        self.config = self.load_config(config_full_path)
        self.components = {}
        
    def load_config(self, path: Path) -> Dict[str, Any]:
        """Load build configuration"""
        if path.exists():
            with open(path, 'r') as f:
                return json.load(f)
        
        # Default config
        return {
            "sourceDir": "src",
            "buildDir": "build",
            "componentLibrary": "src/components/ui-kit.pen",
            "buildOptions": {
                "validateAfterBuild": True
            }
        }
    
    def _load_components_from_file(self, file_path: Path, components: Dict[str, Any], prefix: str = "") -> None:
        """Helper to load components from a single .pen file"""
        try:
            with open(file_path, 'r') as f:
                data = json.load(f)
            
            def extract_recursive(children: List[Dict], path_prefix: str):
                for child in children:
                    name = child.get('name', '')
                    
                    if name:
                        full_name = f"{path_prefix}/{name}" if path_prefix else name
                        
                        # Only store if reusable or if we want to track everything
                        # For now, let's allow referencing any named frame
                        components[full_name] = child
                        
                        # Store by simple ID/Name too if unique? 
                        # No, strict pathing is better for atomic design
                    
                    if 'children' in child:
                        extract_recursive(child['children'], full_name if name else path_prefix)

            extract_recursive(data.get('children', []), prefix)
            
            # If this is the main library file (legacy), capture variables too
            # For atomic, we handle variables separately or merge them
            if not self.variables:
                self.variables = data.get('variables', {})
                
        except Exception as e:
            print(f"⚠️  Error reading {file_path.name}: {e}")

    def load_component_library(self) -> Dict[str, Any]:
        """Load all components from ui-kit.pen OR Atomic Design structure"""
        
        components = {}
        self.variables = {}
        
        # 1. Check for Atomic Design Structure
        src_dir = self.project_root / self.config['sourceDir']
        atomic_dirs = ['foundations', 'atoms', 'molecules', 'organisms']
        found_atomic = False
        
        for atom_dir in atomic_dirs:
            dir_path = src_dir / atom_dir
            if dir_path.exists():
                found_atomic = True
                print(f"Scanning {atom_dir}...")
                for pen_file in dir_path.glob('*.pen'):
                    # Use filename or directory logic as prefix? 
                    # Atomic design usually implies explicit naming in the file, 
                    # so we don't necessarily need file prefix if component names are good.
                    self._load_components_from_file(pen_file, components)

        if found_atomic:
            # Ensure variables are loaded from foundations
            tokens_path = src_dir / 'foundations/design-tokens.pen'
            if tokens_path.exists():
                with open(tokens_path, 'r') as f:
                    data = json.load(f)
                    self.variables = data.get('variables', {})
            
            print(f"✓ Loaded {len(components)} components from Atomic Design structure")
            return components

        # 2. Fallback to Legacy Single File
        lib_path = self.project_root / self.config['componentLibrary']
        
        if not lib_path.exists():
            print(f"❌ Component library not found: {lib_path}")
            sys.exit(1)
            
        print(f"Loading legacy library: {lib_path.name}")
        self._load_components_from_file(lib_path, components)
        
        print(f"✓ Loaded {len(components)} components from library")
        return components
    
    def process_component_refs(self, children: List[Dict], injected_components: List[Dict]) -> None:
        """Process component_ref and replace with proper Pencil refs"""
        for i, child in enumerate(children):
            if child.get('type') == 'component_ref':
                # Get component name from ref
                component_path = child.get('component', '')
                
                # Extract component name (after #)
                if '#' in component_path:
                    component_name = component_path.split('#')[1]
                else:
                    component_name = component_path
                
                # Find matching component
                component = None
                for comp_path, comp_obj in self.components.items():
                    if comp_path.endswith(component_name):
                        component = comp_obj
                        break
                
                if component:
                    # Add component to injected list (if not already added)
                    comp_id = component.get('id')
                    if not any(c.get('id') == comp_id for c in injected_components):
                        injected_components.append(component.copy())
                    
                    # Replace component_ref with proper ref
                    overrides = child.get('overrides', {})
                    
                    children[i] = {
                        'type': 'ref',
                        'ref': comp_id,
                        'name': child.get('name', f'ref_{comp_id}'),
                        'descendants': overrides
                    }
                    
                    # Copy position/size if specified
                    for key in ['x', 'y', 'width', 'height']:
                        if key in child:
                            children[i][key] = child[key]
                else:
                    print(f"⚠️  Component not found: {component_name}")
            
            # Recurse into children
            if 'children' in child:
                self.process_component_refs(child['children'], injected_components)
    
    def build_page(self, page_path: Path, output_path: Path) -> None:
        """Build a single page file"""
        print(f"Building {page_path.name}...")
        
        # Load page
        with open(page_path, 'r') as f:
            page = json.load(f)
        
        # Collect injected components
        injected_components = []
        
        # Process all component_ref in the page
        self.process_component_refs(page.get('children', []), injected_components)
        
        # Inject components at the beginning of children array
        # (Pencil expects reusable components to be defined before usage)
        page['children'] = injected_components + page['children']
        
        # Ensure design variables are included
        if not page.get('variables'):
            page['variables'] = self.variables
        
        # Write output
        output_path.parent.mkdir(parents=True, exist_ok=True)
        with open(output_path, 'w') as f:
            json.dump(page, f, indent=2)
        
        print(f"  ✓ {len(injected_components)} components injected")
        print(f"  ✓ Saved to {output_path}")
    
    def validate_output(self, file_path: Path) -> bool:
        """Validate built file"""
        try:
            with open(file_path, 'r') as f:
                data = json.load(f)
            
            # Check structure
            assert 'version' in data, "Missing version field"
            assert 'children' in data, "Missing children array"
            assert 'variables' in data, "Missing variables object"
            
            return True
        except Exception as e:
            print(f"❌ Validation failed: {e}")
            return False
    
    def build_monolithic(self, output_filename: str = "tPOS.pen") -> None:
        """Build monolithic file merging all pages and component library"""
        print("🔨 Building Monolithic File")
        print("=" * 50)
        print(f"📁 Project root: {self.project_root}")
        print()
        
        # Load component library
        lib_path = self.project_root / self.config['componentLibrary']
        if not lib_path.exists():
            print(f"❌ Component library not found: {lib_path}")
            sys.exit(1)
        
        with open(lib_path, 'r') as f:
            lib_data = json.load(f)
        
        # Get variables from component library
        variables = lib_data.get('variables', {})
        
        # Find all pages in src/pages/
        src_pages = self.project_root / self.config['sourceDir'] / 'pages'
        if not src_pages.exists():
            print(f"❌ Source pages directory not found: {src_pages}")
            sys.exit(1)
        
        pages = sorted(src_pages.glob('*.pen'))
        if not pages:
            print(f"⚠️  No .pen files found in {src_pages}")
            return
        
        print(f"Found {len(pages)} page(s) to merge:")
        for page in pages:
            print(f"  - {page.name}")
        print()
        
        # Create monolithic structure
        monolithic = {
            "version": "2.6",
            "children": [],
            "variables": variables
        }
        
        # Load and add all pages
        current_x = 0  # Start position
        page_spacing = 470  # Spacing between pages
        
        for page_path in pages:
            with open(page_path, 'r') as f:
                page_data = json.load(f)
            
            # Extract the main page frame (usually first child or children array)
            page_children = page_data.get('children', [])
            
            # Add each top-level frame to monolithic children
            for child in page_children:
                # Set position for this page
                child['x'] = current_x
                child['y'] = 0
                
                # Calculate next position based on this page's width
                page_width = child.get('width', 1440)
                if isinstance(page_width, str):  # Handle "fill_container"
                    page_width = 1440  # Default width
                
                monolithic['children'].append(child)
                
                # Update position for next page
                current_x += page_width + page_spacing
            
            print(f"  ✓ Merged {page_path.name} ({len(page_children)} frame(s)) at x={child.get('x', 0)}")
        
        # Add component library page at the end
        lib_children = lib_data.get('children', [])
        for child in lib_children:
            child['x'] = current_x
            child['y'] = 0
            monolithic['children'].append(child)
            
            print(f"  ✓ Merged component library ({len(lib_children)} frame(s)) at x={current_x}")
        print()
        
        # Write output to project root
        output_path = self.project_root / output_filename
        with open(output_path, 'w', encoding='utf-8') as f:
            json.dump(monolithic, f, indent=2, ensure_ascii=False)
        
        # Validate
        if self.validate_output(output_path):
            print(f"  ✓ Validation passed")
        
        print()
        print("=" * 50)
        print(f"✅ Monolithic build complete!")
        print(f"📄 Output: {output_path.absolute()}")
        print(f"📊 Total frames: {len(monolithic['children'])}")
        print(f"🎨 Design tokens: {len(variables)}")
    
    def _extract_components_recursive(self, children: List[Dict], gathered: List[Dict]) -> None:
        """Recursively find components, unwrapping structural frames"""
        for child in children:
            # 1. If reusable, ALWAYS keep it (it's explicitly a component)
            if child.get('reusable'):
                gathered.append(child)
                continue
                
            # 2. Check if it's a structural wrapper (Showcase, Section, Examples, Row)
            name = child.get('name', '')
            is_frame = child.get('type') == 'frame'
            is_structural = any(kw in name.lower() for kw in ['showcase', 'section', 'examples', 'row', 'header', 'container'])
            
            if is_frame and is_structural:
                # Unwrap structural frames and recurse to find reusable components inside
                self._extract_components_recursive(child.get('children', []), gathered)
            elif is_frame:
                # Non-structural frame without reusable flag
                # Check if it has nested reusable components inside
                nested_reusables = []
                self._find_nested_reusables(child.get('children', []), nested_reusables)
                
                if nested_reusables:
                    # Has nested reusables, extract them instead of the wrapper
                    gathered.extend(nested_reusables)
                else:
                    # No nested reusables, keep this frame as potential component
                    gathered.append(child)
            # Primitives are ignored at extraction level
    
    def _find_nested_reusables(self, children: List[Dict], found: List[Dict]) -> None:
        """Find all nested reusable components"""
        for child in children:
            if child.get('reusable'):
                found.append(child)
            elif child.get('type') == 'frame':
                self._find_nested_reusables(child.get('children', []), found)

    def build_library(self, output_path: str = None) -> None:
        """Build merged component library from Atomic Design files"""
        print("🔨 Building Component Library")
        print("=" * 50)
        
        components = []
        variables = {}
        
        # 1. Load Variables from Foundations
        src_dir = self.project_root / self.config['sourceDir']
        tokens_path = src_dir / 'foundations/design-tokens.pen'
        
        if tokens_path.exists():
            with open(tokens_path, 'r') as f:
                data = json.load(f)
                variables = data.get('variables', {})
                print(f"✓ Loaded design tokens from {tokens_path.name}")
        else:
            print(f"⚠️  Design tokens not found at {tokens_path}")

        # 2. Collect Components from Atomic Dirs
        atomic_dirs = ['atoms', 'molecules', 'organisms']
        
        total_found = 0
        for atom_dir in atomic_dirs:
            dir_path = src_dir / atom_dir
            if not dir_path.exists():
                continue
                
            print(f"Scanning {atom_dir}...")
            # Sorting ensures deterministic order - use rglob for subdirectories
            for pen_file in sorted(dir_path.rglob('*.pen')):
                with open(pen_file, 'r') as f:
                    data = json.load(f)
                
                # Merge variables from source file (won't override existing tokens)
                file_vars = data.get('variables', {})
                for key, val in file_vars.items():
                    if key not in variables:
                        variables[key] = val
                
                # Use recursive extraction
                file_components = []
                self._extract_components_recursive(data.get('children', []), file_components)
                
                components.extend(file_components)
                total_found += len(file_components)
                print(f"  ✓ Added {len(file_components)} components from {pen_file.name}")

        # 3. Create Library File
        library_file = {
            "version": "2.6",
            "children": [
                {
                    "type": "frame",
                    "name": "aPOS Component Library",  # Main page name
                    "width": 1440,
                    "fill": "$bg-page",
                    "layout": "vertical",
                    "gap": 40,
                    "padding": 40,
                    "children": components
                }
            ],
            "variables": variables
        }
        
        # Determine output path
        if output_path:
            out_file = self.project_root / output_path
        else:
            out_file = self.project_root / self.config['componentLibrary']

        # Ensure directory exists
        out_file.parent.mkdir(parents=True, exist_ok=True)
        
        with open(out_file, 'w', encoding='utf-8') as f:
            json.dump(library_file, f, indent=2)
            
        print()
        print("=" * 50)
        print(f"✅ Library build complete!")
        print(f"📄 Output: {out_file.absolute()}")
        print(f"🧩 Components: {total_found}")
        print(f"🎨 Tokens: {len(variables)}")

    def build_all(self) -> None:
        """Build all pages"""
        print("🔨 Pencil Design Build System")
        print("=" * 50)
        print(f"📁 Project root: {self.project_root}")
        print()
        
        # Load component library
        self.components = self.load_component_library()
        
        # Find all pages in src/pages/ (relative to project root)
        src_pages = self.project_root / self.config['sourceDir'] / 'pages'
        build_dir = self.project_root / self.config['buildDir']
        
        if not src_pages.exists():
            print(f"❌ Source pages directory not found: {src_pages}")
            sys.exit(1)
        
        pages = list(src_pages.glob('*.pen'))
        
        if not pages:
            print(f"⚠️  No .pen files found in {src_pages}")
            return
        
        print(f"\nBuilding {len(pages)} page(s)...\n")
        
        # Build each page
        built_count = 0
        for page_path in pages:
            output_path = build_dir / page_path.name
            self.build_page(page_path, output_path)
            
            # Validate if enabled
            if self.config['buildOptions'].get('validateAfterBuild'):
                if self.validate_output(output_path):
                    print(f"  ✓ Validation passed")
                    built_count += 1
            else:
                built_count += 1
            
            print()
        
        print("=" * 50)
        print(f"✅ Build complete! ({built_count}/{len(pages)} pages)")
        print(f"📂 Output: {build_dir.absolute()}")

def main():
    """Main entry point"""
    import argparse
    
    parser = argparse.ArgumentParser(
        description='Pencil Design Build System',
        formatter_class=argparse.RawDescriptionHelpFormatter,
        epilog="""
Examples:
  # 1. Build Component Library (from Atomic Source)
  python build.py --library
  
  # 2. Build Monolithic File (Pages + Library)
  python build.py --monolithic
        """
    )
    
    parser.add_argument(
        '--monolithic', '--mono', '-m',
        action='store_true',
        help='Build monolithic file (merge all pages and component library)'
    )
    
    parser.add_argument(
        '--library', '--lib', '-l',
        action='store_true',
        help='Build component library file from Atomic Design sources'
    )
    
    parser.add_argument(
        '--mode',
        choices=['standard', 'monolithic', 'library'],
        default='standard',
        help='Build mode'
    )
    
    parser.add_argument(
        '--output', '-o',
        help='Output filename'
    )
    
    args = parser.parse_args()
    
    builder = PencilBuilder()
    
    # Determine build mode
    if args.library or args.mode == 'library':
        builder.build_library(output_path=args.output)
    elif args.monolithic or args.mode == 'monolithic':
        # Default output for monolithic is tPOS.pen, but allow override
        output = args.output if args.output else "tPOS.pen"
        builder.build_monolithic(output_filename=output)
    else:
        builder.build_all()

if __name__ == '__main__':
    main()
