import os
import sys
import importlib
import inspect
from pathlib import Path
import matplotlib
import shutil
matplotlib.use('Agg')
import matplotlib.pyplot as plt

import nazca as nd
import mxpic as mx

# 2. Define your mxPIC Color Palette
# You can map by your custom layer names (if defined) or raw GDS (layer, datatype) tuples.
# Matplotlib accepts standard color names ('blue', 'cyan') or hex codes ('#FFD700').
PALETTE = {
    "WG": "blue",             # Example: Core waveguide layer
    "SLAB": "cyan",           # Example: Shallow etch / Slab
    "M1": "#FFD700",          # Example: Metal 1 (Gold)
    "M2": "silver",           # Example: Metal 2
    "DEEP_TRENCH": "black",   # Example: Trenching
    (1, 0): "darkred",        # Fallback: You can use raw GDS tuples directly
    (2, 0): "green",
    (1111, 0): "green",
    (63, 30): "#FFD700",
}

def apply_mxpic_colors():
    """Applies the custom color palette to Nazca's active layer map."""
    print("🎨 Applying mxPIC layer colors...")
    for layer_id, color in PALETTE.items():
        try:
            # Nazca's built-in command to register plot colors
            nd.set_layercolor(layer=layer_id, color=color)
        except Exception:
            # If a specific layer name doesn't exist in the registry yet, it safely skips
            pass

def generate_component_images(img_root="images/components"):
    print("📸 Starting mxPIC Component Image Generation...")

    # Define our source and target directories
    src_root = Path("mxpic/components")
    img_root = Path(img_root)

    # 1. Delete the directory and all its contents
    if img_root.exists() and img_root.is_dir():
        shutil.rmtree(img_root)

    # 2. Recreate the directory
    img_root.mkdir(parents=True, exist_ok=True)

    if not src_root.exists():
        print(f"❌ Error: Source directory '{src_root}' not found.")
        sys.exit(1)

    # Walk through all Python files in the components folder
    success_count = 0
    fail_count = 0
    
    tapeout = mx.foundries.Silterra.EOM1_2ML_CU()

    for py_file in src_root.rglob("*.py"):
        if py_file.name == "__init__.py":
            continue

        # Convert the file path to a Python module path (e.g., mxpic.components.mzm)
        rel_path = py_file.relative_to(src_root)
        module_name = "mxpic.components." + str(rel_path.with_suffix("")).replace(os.sep, ".")

        try:
            # Dynamically import the module
            module = importlib.import_module(module_name)
        except Exception as e:
            print(f"⚠️ Could not import {module_name}: {e}")
            fail_count += 1
            continue

            
        # Scan the module for all defined classes
        for name, obj in inspect.getmembers(module, inspect.isclass):
            # Only process classes actually defined IN this file (ignore imported classes)
            if obj.__module__ == module_name:
                
                # Determine where to save the image
                target_dir = img_root / rel_path.parent
                target_dir.mkdir(parents=True, exist_ok=True)
                
                # img_path = target_dir / f"{name}.png"

                try:
                    # 1. Clear the Nazca canvas so components don't overlap!
                    nd.clear_layout()

                    # 2. Instantiate the class (assuming zero arguments)
                    instance = obj()

                    # 4. Export the image using Nazca
                    nd.export_plt(path="",title=f"{name}",topcells=[instance.cell])
                    # 3. Explicitly save to disk with tight borders
                    plt.savefig(str(target_dir)+f"\\{name}.png", bbox_inches='tight', dpi=300)
                    
                    # 4. CRITICAL: Clear the figure from RAM so the next loop is clean
                    plt.close()
                    
                    print(f"✅ Generated: {str(target_dir)}\\{name}.png")
                    success_count += 1

                except Exception as e:
                    print(f"❌ Failed to generate image for {name} in {module_name}: {e}")
                    fail_count += 1

    print("\n✨ Image generation complete!")
    print(f"📊 Success: {success_count} | Failures: {fail_count}")

if __name__ == "__main__":
    generate_component_images(img_root="docs/source/mxpic/components/")