#!/usr/bin/env python3
"""
variant_manager.py
==================
Manage KiCad design variants - track which parts are fitted/unfitted in different variants.

Variants are stored in a JSON file alongside the project.
"""

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


class VariantManager:
    """Manages design variants for a KiCad project"""

    def __init__(self, project_path: str):
        """
        Initialize variant manager for a project.

        Args:
            project_path: Path to the .kicad_pro or .kicad_sch file
        """
        self.project_path = Path(project_path)
        self.project_dir = self.project_path.parent
        self.project_name = self.project_path.stem

        # Variants file stored alongside project
        self.variants_file = self.project_dir / f"{self.project_name}.variants.json"
        self.variants = self._load_variants()

    def _load_variants(self) -> Dict:
        """Load variants from JSON file"""
        if self.variants_file.exists():
            with open(self.variants_file, 'r', encoding='utf-8') as f:
                data = json.load(f)

            # Migrate from version 2 to version 3
            if data.get("meta", {}).get("version", 1) < 3:
                print(f"Migrating variant file from version {data.get('meta', {}).get('version', 1)} to 3")
                if "base_values" not in data:
                    data["base_values"] = {}
                for variant in data.get("variants", {}).values():
                    if "part_overrides" not in variant:
                        variant["part_overrides"] = {}
                data["meta"]["version"] = 3
                # Save migrated version
                with open(self.variants_file, 'w', encoding='utf-8') as f:
                    json.dump(data, f, indent=2)

            return data
        else:
            # Default structure
            return {
                "meta": {
                    "version": 3,  # Version 3 adds base_values and part_overrides
                    "active_variant": "default"
                },
                "base_values": {
                    # UUID -> {property_name: value}
                    # Stores the base/default values for parts that vary between variants
                },
                "variants": {
                    "default": {
                        "name": "default",
                        "description": "Default variant - all parts fitted",
                        "dnp_parts": [],  # List of UUIDs that are DNP (Do Not Place)
                        "part_overrides": {}  # UUID -> {property_name: override_value}
                    }
                }
            }

    def _save_variants(self):
        """Save variants to JSON file"""
        with open(self.variants_file, 'w', encoding='utf-8') as f:
            json.dump(self.variants, f, indent=2)

    def get_variants(self) -> Dict:
        """Get all variants"""
        return self.variants["variants"]

    def get_active_variant(self) -> str:
        """Get the name of the active variant"""
        return self.variants["meta"]["active_variant"]

    def create_variant(self, name: str, description: str = "", based_on: str = None) -> bool:
        """
        Create a new variant.

        Args:
            name: Variant name
            description: Variant description
            based_on: Name of variant to copy from (None = start empty)

        Returns:
            True if created, False if already exists
        """
        if name in self.variants["variants"]:
            return False

        if based_on and based_on in self.variants["variants"]:
            # Copy DNP list from base variant
            dnp_parts = self.variants["variants"][based_on]["dnp_parts"].copy()
        else:
            dnp_parts = []

        self.variants["variants"][name] = {
            "name": name,
            "description": description,
            "dnp_parts": dnp_parts
        }

        self._save_variants()
        return True

    def delete_variant(self, name: str) -> bool:
        """
        Delete a variant.

        Args:
            name: Variant name

        Returns:
            True if deleted, False if doesn't exist or is active
        """
        if name not in self.variants["variants"]:
            return False

        if name == self.variants["meta"]["active_variant"]:
            return False  # Can't delete active variant

        if name == "default":
            return False  # Can't delete default variant

        del self.variants["variants"][name]
        self._save_variants()
        return True

    def set_active_variant(self, name: str) -> bool:
        """
        Set the active variant.

        Args:
            name: Variant name

        Returns:
            True if set, False if variant doesn't exist
        """
        if name not in self.variants["variants"]:
            return False

        self.variants["meta"]["active_variant"] = name
        self._save_variants()
        return True

    def set_part_dnp(self, variant_name: str, uuid: str, is_dnp: bool) -> bool:
        """
        Set whether a part is DNP (Do Not Place) in a variant.

        Args:
            variant_name: Variant name
            uuid: Component UUID (e.g., "681abb84-6eb2-4c95-9a2f-a9fc19a34beb")
            is_dnp: True to mark as DNP, False to mark as fitted

        Returns:
            True if successful, False if variant doesn't exist
        """
        if variant_name not in self.variants["variants"]:
            return False

        dnp_list = self.variants["variants"][variant_name]["dnp_parts"]

        if is_dnp:
            if uuid not in dnp_list:
                dnp_list.append(uuid)
                dnp_list.sort()  # Keep sorted
        else:
            if uuid in dnp_list:
                dnp_list.remove(uuid)

        self._save_variants()
        return True

    def get_dnp_parts(self, variant_name: str) -> List[str]:
        """
        Get list of DNP parts for a variant.

        Args:
            variant_name: Variant name

        Returns:
            List of UUIDs, or empty list if variant doesn't exist
        """
        if variant_name not in self.variants["variants"]:
            return []

        return self.variants["variants"][variant_name]["dnp_parts"].copy()

    def is_part_dnp(self, variant_name: str, uuid: str) -> bool:
        """
        Check if a part is DNP in a variant.

        Args:
            variant_name: Variant name
            uuid: Component UUID

        Returns:
            True if DNP, False if fitted or variant doesn't exist
        """
        if variant_name not in self.variants["variants"]:
            return False

        return uuid in self.variants["variants"][variant_name]["dnp_parts"]

    def get_part_properties(self, variant_name: str, uuid: str) -> Dict[str, str]:
        """
        Get effective property values for a part in a variant.
        Merges base_values with variant-specific overrides.

        Args:
            variant_name: Variant name
            uuid: Component UUID

        Returns:
            Dict of property_name -> value (merged base + overrides)
        """
        if variant_name not in self.variants["variants"]:
            return {}

        # Start with base values
        properties = self.variants.get("base_values", {}).get(uuid, {}).copy()

        # Apply variant overrides
        overrides = self.variants["variants"][variant_name].get("part_overrides", {}).get(uuid, {})
        properties.update(overrides)

        return properties

    def set_part_property(self, variant_name: str, uuid: str, property_name: str, value: str) -> bool:
        """
        Set a property override for a part in a variant.

        Args:
            variant_name: Variant name
            uuid: Component UUID
            property_name: Property to set (e.g., "Value", "MPN")
            value: New value

        Returns:
            True if successful, False if variant doesn't exist
        """
        if variant_name not in self.variants["variants"]:
            return False

        # Ensure part_overrides exists for this variant
        if "part_overrides" not in self.variants["variants"][variant_name]:
            self.variants["variants"][variant_name]["part_overrides"] = {}

        # Get or create override dict for this UUID
        if uuid not in self.variants["variants"][variant_name]["part_overrides"]:
            self.variants["variants"][variant_name]["part_overrides"][uuid] = {}

        # Check if this matches base value - if so, remove override (keep it sparse)
        base_value = self.variants.get("base_values", {}).get(uuid, {}).get(property_name)
        if value == base_value:
            # Remove from overrides since it matches base
            if property_name in self.variants["variants"][variant_name]["part_overrides"][uuid]:
                del self.variants["variants"][variant_name]["part_overrides"][uuid][property_name]
                # Clean up empty override dicts
                if not self.variants["variants"][variant_name]["part_overrides"][uuid]:
                    del self.variants["variants"][variant_name]["part_overrides"][uuid]
        else:
            # Store override
            self.variants["variants"][variant_name]["part_overrides"][uuid][property_name] = value

        self._save_variants()
        return True

    def set_base_value(self, uuid: str, property_name: str, value: str):
        """
        Set a base property value for a part.

        Args:
            uuid: Component UUID
            property_name: Property name
            value: Value to set
        """
        if "base_values" not in self.variants:
            self.variants["base_values"] = {}

        if uuid not in self.variants["base_values"]:
            self.variants["base_values"][uuid] = {}

        self.variants["base_values"][uuid][property_name] = value
        self._save_variants()


if __name__ == "__main__":
    if len(sys.argv) < 2:
        print("Usage: python variant_manager.py <project.kicad_pro>")
        sys.exit(1)

    manager = VariantManager(sys.argv[1])

    # Print current variants
    print(f"Project: {manager.project_name}")
    print(f"Active variant: {manager.get_active_variant()}")
    print("\nVariants:")
    for name, variant in manager.get_variants().items():
        dnp_count = len(variant["dnp_parts"])
        print(f"  {name}: {variant['description']} ({dnp_count} DNP parts)")