dss/dss-mvp1/dss/translations/merger.py

"""
Theme Merger

Merges base DSS theme with translation mappings and custom props.
"""

from datetime import datetime, timezone
from typing import Any, Dict, Optional, Union

from dss.models.theme import DesignToken, Theme, TokenCategory
from dss.themes.default_themes import get_default_dark_theme, get_default_light_theme

from .models import ResolvedTheme, ResolvedToken, TranslationRegistry
from .resolver import TokenResolver


class ThemeMerger:
    """
    Merges base DSS theme with project-specific customizations.

    The merge hierarchy:
    1. Base Theme (DSS Light or Dark)
    2. Translation Mappings (external tokens -> DSS)
    3. Custom Props (project-specific extensions)

    Usage:
        merger = ThemeMerger(registry)
        resolved = await merger.merge(base_theme="light")
    """

    def __init__(self, registry: TranslationRegistry):
        """
        Initialize merger with translation registry.

        Args:
            registry: TranslationRegistry with loaded dictionaries
        """
        self.registry = registry
        self.resolver = TokenResolver(registry)

    async def merge(
        self, base_theme: str = "light", project_name: Optional[str] = None
    ) -> ResolvedTheme:
        """
        Merge base theme with translations and custom props.

        Args:
            base_theme: Base theme name ("light" or "dark")
            project_name: Project name for resolved theme

        Returns:
            ResolvedTheme with all tokens resolved
        """
        # Get base theme
        if base_theme == "light":
            theme = get_default_light_theme()
        elif base_theme == "dark":
            theme = get_default_dark_theme()
        else:
            raise ValueError(f"Unknown base theme: {base_theme}")

        # Convert theme tokens to dict for resolution
        base_tokens = self._theme_to_dict(theme)

        # Resolve all mapped tokens
        resolved_tokens = self.resolver.resolve_all_mappings(base_tokens)

        # Separate core tokens from custom props
        core_tokens = {}
        custom_props = {}

        for dss_path, resolved in resolved_tokens.items():
            if resolved.is_custom:
                custom_props[dss_path] = resolved
            else:
                core_tokens[dss_path] = resolved

        # Add base theme tokens that aren't in mappings
        for token_name, token in theme.tokens.items():
            # Normalize token name to DSS path
            dss_path = self._normalize_to_dss_path(token_name)
            if dss_path not in core_tokens:
                core_tokens[dss_path] = ResolvedToken(
                    dss_path=dss_path,
                    value=token.value,
                    is_custom=False,
                    provenance=[f"base_theme: {base_theme}"],
                )

        return ResolvedTheme(
            name=project_name or f"resolved-{base_theme}",
            version="1.0.0",
            base_theme=base_theme,
            tokens=core_tokens,
            custom_props=custom_props,
            translations_applied=[dict_name for dict_name in self.registry.dictionaries.keys()],
            resolved_at=datetime.now(timezone.utc),
        )

    def _theme_to_dict(self, theme: Theme) -> Dict[str, Any]:
        """Convert Theme object to nested dict for resolution."""
        result = {}
        for token_name, token in theme.tokens.items():
            # Convert flat token names to nested structure
            parts = self._normalize_to_dss_path(token_name).split(".")
            current = result
            for part in parts[:-1]:
                if part not in current:
                    current[part] = {}
                elif not isinstance(current[part], dict):
                    # Skip if this path is already set to a value
                    continue
                current = current[part]
            current[parts[-1]] = token.value
        return result

    def _normalize_to_dss_path(self, token_name: str) -> str:
        """Normalize token name to DSS canonical path."""
        # Handle various formats
        normalized = token_name.replace("-", ".").replace("_", ".")

        # Map common prefixes
        prefix_map = {
            "space.": "spacing.",
            "radius.": "border.radius.",
            "text.": "typography.size.",
        }

        for old, new in prefix_map.items():
            if normalized.startswith(old):
                normalized = new + normalized[len(old) :]
                break

        return normalized

    async def merge_custom_props(
        self, resolved_theme: ResolvedTheme, additional_props: Dict[str, Any]
    ) -> ResolvedTheme:
        """
        Add additional custom props to a resolved theme.

        Args:
            resolved_theme: Existing resolved theme
            additional_props: Additional custom props to merge

        Returns:
            Updated ResolvedTheme
        """
        for prop_name, prop_value in additional_props.items():
            resolved_theme.custom_props[prop_name] = ResolvedToken(
                dss_path=prop_name,
                value=prop_value,
                is_custom=True,
                provenance=["additional_custom_prop"],
            )

        resolved_theme.resolved_at = datetime.now(timezone.utc)
        return resolved_theme

    def export_as_theme(self, resolved: ResolvedTheme) -> Theme:
        """
        Convert ResolvedTheme back to Theme model.

        Args:
            resolved: ResolvedTheme to convert

        Returns:
            Theme model instance
        """
        tokens = {}

        # Add core tokens
        for dss_path, resolved_token in resolved.tokens.items():
            token_name = dss_path.replace(".", "-")
            tokens[token_name] = DesignToken(
                name=token_name,
                value=resolved_token.value,
                type=self._infer_type(dss_path, resolved_token.value),
                category=self._infer_category(dss_path),
                source=f"resolved:{resolved.base_theme}",
            )

        # Add custom props
        for dss_path, resolved_token in resolved.custom_props.items():
            token_name = dss_path.replace(".", "-")
            tokens[token_name] = DesignToken(
                name=token_name,
                value=resolved_token.value,
                type=self._infer_type(dss_path, resolved_token.value),
                category=TokenCategory.OTHER,
                source="custom_prop",
            )

        return Theme(name=resolved.name, version=resolved.version, tokens=tokens)

    def _infer_type(self, path: str, value: Any) -> str:
        """Infer token type from path and value."""
        if "color" in path:
            return "color"
        if "spacing" in path or "size" in path or "radius" in path:
            return "dimension"
        if "font" in path:
            return "typography"
        if "shadow" in path:
            return "shadow"
        return "string"

    def _infer_category(self, path: str) -> TokenCategory:
        """Infer token category from DSS path."""
        if path.startswith("color"):
            return TokenCategory.COLOR
        if path.startswith("spacing"):
            return TokenCategory.SPACING
        if path.startswith("typography") or path.startswith("font"):
            return TokenCategory.TYPOGRAPHY
        if path.startswith("border") or path.startswith("radius"):
            return TokenCategory.RADIUS
        if path.startswith("shadow"):
            return TokenCategory.SHADOW
        return TokenCategory.OTHER