"""Parse Home Assistant YAML configs to extract entity and device references."""

import re
from pathlib import Path
from typing import Any

import yaml


def extract_entity_ids(data: Any, entity_ids: set[str] | None = None) -> set[str]:
    """
    Recursively extract entity_id values from YAML data structure.
    
    Handles:
    - Direct entity_id: sensor.temperature
    - List of entity_ids: [sensor.temp1, sensor.temp2]
    - Nested in target.entity_id
    - Jinja2 templates: {{ states('sensor.temperature') }}
    
    Args:
        data: Parsed YAML data (dict, list, or primitive)
        entity_ids: Set to accumulate entity_ids (created if None)
    
    Returns:
        Set of entity_id strings
    """
    if entity_ids is None:
        entity_ids = set()
    
    if isinstance(data, dict):
        for key, value in data.items():
            # Direct entity_id field
            if key == 'entity_id':
                if isinstance(value, str):
                    # Could be entity_id or device_id (UUID)
                    # Filter out UUIDs (device_ids are 32-char hex without dots)
                    if '.' in value:
                        entity_ids.add(value)
                elif isinstance(value, list):
                    for item in value:
                        if isinstance(item, str) and '.' in item:
                            entity_ids.add(item)
            # Recursively process nested structures
            elif isinstance(value, (dict, list)):
                extract_entity_ids(value, entity_ids)
            # Extract from Jinja2 templates
            elif isinstance(value, str):
                entity_ids.update(extract_entities_from_template(value))
    
    elif isinstance(data, list):
        for item in data:
            extract_entity_ids(item, entity_ids)
    
    return entity_ids


def extract_device_ids(data: Any, device_ids: set[str] | None = None) -> set[str]:
    """
    Recursively extract device_id values from YAML data structure.
    
    Device IDs are UUIDs (32 hex characters without hyphens or dots).
    
    Args:
        data: Parsed YAML data (dict, list, or primitive)
        device_ids: Set to accumulate device_ids (created if None)
    
    Returns:
        Set of device_id strings (UUIDs)
    """
    if device_ids is None:
        device_ids = set()
    
    # UUID pattern: 32 hex characters
    uuid_pattern = re.compile(r'^[a-f0-9]{32}$')
    
    if isinstance(data, dict):
        for key, value in data.items():
            # Direct device_id field
            if key == 'device_id' and isinstance(value, str):
                if uuid_pattern.match(value):
                    device_ids.add(value)
            # Also check entity_id field for UUIDs (some automations use this)
            elif key == 'entity_id' and isinstance(value, str):
                if uuid_pattern.match(value):
                    device_ids.add(value)
            # Recursively process nested structures
            elif isinstance(value, (dict, list)):
                extract_device_ids(value, device_ids)
    
    elif isinstance(data, list):
        for item in data:
            extract_device_ids(item, device_ids)
    
    return device_ids


def extract_entities_from_template(template: str) -> set[str]:
    """
    Extract entity_ids from Jinja2 template strings.
    
    Patterns matched:
    - states('sensor.temperature')
    - state_attr('sensor.temperature', 'attribute')
    - is_state('sensor.temperature', 'on')
    
    Args:
        template: Jinja2 template string
    
    Returns:
        Set of entity_id strings found in template
    """
    entity_ids = set()
    
    # Pattern to match entity_ids in common Jinja2 functions
    patterns = [
        r"states\(['\"]([a-z_]+\.[a-z0-9_]+)['\"]",
        r"state_attr\(['\"]([a-z_]+\.[a-z0-9_]+)['\"]",
        r"is_state\(['\"]([a-z_]+\.[a-z0-9_]+)['\"]",
    ]
    
    for pattern in patterns:
        matches = re.findall(pattern, template, re.IGNORECASE)
        entity_ids.update(matches)
    
    return entity_ids


def parse_yaml_file(file_path: Path) -> dict[str, Any]:
    """
    Parse a YAML file with Home Assistant-specific handling.
    
    Args:
        file_path: Path to YAML file
    
    Returns:
        Parsed YAML data as dict
    
    Raises:
        yaml.YAMLError: If file cannot be parsed
    """
    # Add custom constructor for !secret tag (replace with placeholder)
    def secret_constructor(loader, node):
        return f"SECRET_{loader.construct_scalar(node)}"
    
    # Add custom constructor for !include (skip)
    def include_constructor(loader, node):
        return f"INCLUDE_{loader.construct_scalar(node)}"
    
    yaml.SafeLoader.add_constructor('!secret', secret_constructor)
    yaml.SafeLoader.add_constructor('!include', include_constructor)
    
    with open(file_path, 'r') as f:
        return yaml.safe_load(f) or {}


def parse_config_files(config_dir: Path) -> tuple[set[str], set[str]]:
    """
    Parse all Home Assistant config files and extract entity/device IDs.
    
    Args:
        config_dir: Path to Home Assistant config directory
    
    Returns:
        Tuple of (entity_ids, device_ids)
    """
    entity_ids = set()
    device_ids = set()
    
    # Files to parse (exclude templates.yaml per requirements)
    files_to_parse = [
        'automations.yaml',
        'scripts.yaml',
    ]
    
    for filename in files_to_parse:
        file_path = config_dir / filename
        if file_path.exists():
            try:
                data = parse_yaml_file(file_path)
                entity_ids.update(extract_entity_ids(data))
                device_ids.update(extract_device_ids(data))
            except yaml.YAMLError as e:
                print(f"Warning: Could not parse {filename}: {e}")
    
    return entity_ids, device_ids


if __name__ == '__main__':
    import sys
    
    if len(sys.argv) < 2:
        print("Usage: python parse_config.py <config_directory>")
        sys.exit(1)
    
    config_path = Path(sys.argv[1])
    entities, devices = parse_config_files(config_path)
    
    print(f"Found {len(entities)} entity IDs:")
    for entity_id in sorted(entities):
        print(f"  - {entity_id}")
    
    print(f"\nFound {len(devices)} device IDs:")
    for device_id in sorted(devices):
        print(f"  - {device_id}")