import httpx
from typing import Dict, Any, Optional, List, TypeVar, Callable, Awaitable, Union, cast
import functools
import inspect
import logging
from datetime import datetime, timedelta, timezone
import json
import asyncio
import ssl
import websockets

from app.config import HA_URL, HA_TOKEN, get_ha_headers

# Set up logging
logger = logging.getLogger(__name__)

# Define a generic type for our API function return values
T = TypeVar('T')
F = TypeVar('F', bound=Callable[..., Awaitable[Any]])

# HTTP client
_client: Optional[httpx.AsyncClient] = None

# Default field sets for different verbosity levels
# Lean fields for standard requests (optimized for token efficiency)
DEFAULT_LEAN_FIELDS = ["entity_id", "state", "attr.friendly_name"]

# Common fields that are typically needed for entity operations
DEFAULT_STANDARD_FIELDS = ["entity_id", "state", "attributes", "last_updated"]

# Domain-specific important attributes to include in lean responses
DOMAIN_IMPORTANT_ATTRIBUTES = {
    "light": ["brightness", "color_temp", "rgb_color", "supported_color_modes"],
    "switch": ["device_class"],
    "binary_sensor": ["device_class"],
    "sensor": ["device_class", "unit_of_measurement", "state_class"],
    "climate": ["hvac_mode", "current_temperature", "temperature", "hvac_action"],
    "media_player": ["media_title", "media_artist", "source", "volume_level"],
    "cover": ["current_position", "current_tilt_position"],
    "fan": ["percentage", "preset_mode"],
    "camera": ["entity_picture"],
    "automation": ["last_triggered"],
    "scene": [],
    "script": ["last_triggered"],
}

# Sensitive keys that should be redacted in logs
_SENSITIVE_KEYS = frozenset([
    "token", "password", "api_key", "secret", "access_token",
    "authorization", "auth", "credential", "key", "pin", "code",
    "bearer", "refresh_token", "client_secret", "private_key",
    "passphrase", "session", "cookie"
])


def sanitize_for_logging(data: Any) -> Any:
    """
    Sanitize data for safe logging by redacting sensitive fields.

    Args:
        data: The data to sanitize (dict, list, or primitive)

    Returns:
        Sanitized copy of the data with sensitive values redacted
    """
    if isinstance(data, dict):
        return {
            k: "***REDACTED***" if k.lower() in _SENSITIVE_KEYS else sanitize_for_logging(v)
            for k, v in data.items()
        }
    elif isinstance(data, list):
        return [sanitize_for_logging(item) for item in data]
    else:
        return data

def handle_api_errors(func: F) -> F:
    """
    Decorator to handle common error cases for Home Assistant API calls
    
    Args:
        func: The async function to decorate
        
    Returns:
        Wrapped function that handles errors
    """
    @functools.wraps(func)
    async def wrapper(*args: Any, **kwargs: Any) -> Any:
        # Determine return type from function annotation
        return_type = inspect.signature(func).return_annotation
        is_dict_return = 'Dict' in str(return_type)
        is_list_return = 'List' in str(return_type)
        
        # Prepare error formatters based on return type
        def format_error(msg: str) -> Any:
            if is_dict_return:
                return {"error": msg}
            elif is_list_return:
                return [{"error": msg}]
            else:
                return msg
        
        try:
            # Check if token is available
            if not HA_TOKEN:
                return format_error("No Home Assistant token provided. Please set HA_TOKEN in .env file.")
            
            # Call the original function
            return await func(*args, **kwargs)
        except httpx.ConnectError:
            return format_error(f"Connection error: Cannot connect to Home Assistant at {HA_URL}")
        except httpx.TimeoutException:
            return format_error(f"Timeout error: Home Assistant at {HA_URL} did not respond in time")
        except httpx.HTTPStatusError as e:
            return format_error(f"HTTP error: {e.response.status_code} - {e.response.reason_phrase}")
        except httpx.RequestError as e:
            return format_error(f"Error connecting to Home Assistant: {str(e)}")
        except Exception as e:
            return format_error(f"Unexpected error: {str(e)}")
    
    return cast(F, wrapper)

# Persistent HTTP client with granular timeout configuration
_timeout_config = httpx.Timeout(
    connect=5.0,    # Connection timeout - fail fast on unreachable hosts
    read=30.0,      # Read timeout - allow time for long-running queries (stats, history)
    write=5.0,      # Write timeout - service calls should be quick
    pool=5.0        # Pool timeout - waiting for available connection
)


class RateLimiter:
    """
    Simple async rate limiter using token bucket algorithm.
    Limits requests to max_rate per second.
    """

    def __init__(self, max_rate: float = 10.0):
        """
        Initialize rate limiter.

        Args:
            max_rate: Maximum requests per second (default: 10)
        """
        self.max_rate = max_rate
        self.tokens = max_rate
        self.last_update: float = 0.0  # Will be set on first acquire()
        self._lock = asyncio.Lock()

    async def acquire(self) -> None:
        """Wait until a request token is available."""
        async with self._lock:
            loop = asyncio.get_running_loop()
            now = loop.time()

            # On first call, just set last_update and use existing tokens
            if self.last_update == 0.0:
                self.last_update = now
            else:
                # Refill tokens based on time passed
                time_passed = now - self.last_update
                self.tokens = min(self.max_rate, self.tokens + time_passed * self.max_rate)
                self.last_update = now

            if self.tokens < 1:
                # Wait for token to become available
                wait_time = (1 - self.tokens) / self.max_rate
                await asyncio.sleep(wait_time)
                # Token is now available after waiting - deduct it
                self.tokens -= 1
            else:
                self.tokens -= 1


# Global rate limiter - 10 requests per second
_rate_limiter = RateLimiter(max_rate=10.0)


async def get_client() -> httpx.AsyncClient:
    """Get a persistent httpx client for Home Assistant API calls"""
    global _client
    if _client is None:
        logger.debug("Creating new HTTP client")
        _client = httpx.AsyncClient(timeout=_timeout_config)
    return _client

async def cleanup_client() -> None:
    """Close the HTTP client when shutting down"""
    global _client
    if _client:
        logger.debug("Closing HTTP client")
        await _client.aclose()
        _client = None

async def call_websocket_api(message_type: str, **kwargs) -> Dict[str, Any]:
    """
    Call Home Assistant WebSocket API.

    Args:
        message_type: The WebSocket message type (e.g., 'recorder/statistics_during_period')
        **kwargs: Additional parameters for the message

    Returns:
        The response from Home Assistant
    """
    # Convert HTTP URL to WebSocket URL
    ws_url = HA_URL.replace("http://", "ws://").replace("https://", "wss://")
    ws_url = f"{ws_url}/api/websocket"

    # Use SSL context for secure WebSocket connections
    ssl_context = None
    if ws_url.startswith("wss://"):
        ssl_context = ssl.create_default_context()
        # Validates certificates by default - no need for explicit verify_mode

    try:
        # Apply rate limiting before making connection
        await _rate_limiter.acquire()
        async with websockets.connect(ws_url, ssl=ssl_context) as websocket:
            # Wait for auth required message
            auth_msg = await websocket.recv()
            auth_data = json.loads(auth_msg)

            if auth_data.get("type") == "auth_required":
                # Send authentication
                await websocket.send(json.dumps({
                    "type": "auth",
                    "access_token": HA_TOKEN
                }))

                # Wait for auth result
                auth_result = await websocket.recv()
                auth_result_data = json.loads(auth_result)

                if auth_result_data.get("type") != "auth_ok":
                    raise Exception(f"Authentication failed: {auth_result_data}")

                # Now send the actual request
                message_id = 1
                message = {
                    "id": message_id,
                    "type": message_type,
                    **kwargs
                }

                await websocket.send(json.dumps(message))

                # Wait for response
                response = await websocket.recv()
                response_data = json.loads(response)

                # Check for success
                if response_data.get("success") is False:
                    error = response_data.get("error", {})
                    raise Exception(f"WebSocket API error: {error.get('message', 'Unknown error')}")

                return response_data.get("result", {})
            else:
                raise Exception(f"Unexpected auth message: {auth_data}")

    except ssl.SSLError as e:
        logger.error(f"SSL/TLS error connecting to Home Assistant: {str(e)}")
        raise Exception(f"SSL certificate error: {str(e)}")
    except websockets.exceptions.InvalidURI as e:
        logger.error(f"Invalid WebSocket URI: {ws_url}")
        raise Exception(f"Invalid WebSocket URI: {str(e)}")
    except websockets.exceptions.ConnectionClosed as e:
        logger.error(f"WebSocket connection closed unexpectedly: {e.code} - {e.reason}")
        raise Exception(f"WebSocket connection closed: {e.reason or 'Unknown reason'}")
    except websockets.exceptions.WebSocketException as e:
        logger.error(f"WebSocket connection error: {str(e)}")
        raise Exception(f"WebSocket connection failed: {str(e)}")
    except json.JSONDecodeError as e:
        logger.error(f"Failed to parse WebSocket response: {str(e)}")
        raise Exception(f"Invalid response from Home Assistant: {str(e)}")
    except asyncio.TimeoutError:
        logger.error(f"WebSocket connection to {ws_url} timed out")
        raise Exception(f"WebSocket connection timed out")
    except ConnectionRefusedError:
        logger.error(f"Connection refused to {ws_url}")
        raise Exception(f"Connection refused: Home Assistant not reachable at {HA_URL}")
    except OSError as e:
        logger.error(f"Network error connecting to {ws_url}: {str(e)}")
        raise Exception(f"Network error: {str(e)}")
    except Exception as e:
        logger.error(f"WebSocket API error: {str(e)}")
        raise

# Direct entity retrieval function
async def get_all_entity_states() -> Dict[str, Dict[str, Any]]:
    """Fetch all entity states from Home Assistant"""
    await _rate_limiter.acquire()
    client = await get_client()
    response = await client.get(f"{HA_URL}/api/states", headers=get_ha_headers())
    response.raise_for_status()
    entities = response.json()
    
    # Create a mapping for easier access
    return {entity["entity_id"]: entity for entity in entities}

def filter_fields(data: Dict[str, Any], fields: List[str]) -> Dict[str, Any]:
    """
    Filter entity data to only include requested fields
    
    This function helps reduce token usage by returning only requested fields.
    
    Args:
        data: The complete entity data dictionary
        fields: List of fields to include in the result
               - "state": Include the entity state
               - "attributes": Include all attributes
               - "attr.X": Include only attribute X (e.g. "attr.brightness")
               - "context": Include context data
               - "last_updated"/"last_changed": Include timestamp fields
    
    Returns:
        A filtered dictionary with only the requested fields
    """
    if not fields:
        return data
        
    result = {"entity_id": data["entity_id"]}
    
    for field in fields:
        if field == "state":
            result["state"] = data.get("state")
        elif field == "attributes":
            result["attributes"] = data.get("attributes", {})
        elif field.startswith("attr.") and len(field) > 5:
            attr_name = field[5:]
            attributes = data.get("attributes", {})
            if attr_name in attributes:
                if "attributes" not in result:
                    result["attributes"] = {}
                result["attributes"][attr_name] = attributes[attr_name]
        elif field == "context":
            if "context" in data:
                result["context"] = data["context"]
        elif field in ["last_updated", "last_changed"]:
            if field in data:
                result[field] = data[field]
    
    return result

# API Functions
@handle_api_errors
async def get_hass_version() -> str:
    """Get the Home Assistant version from the API"""
    await _rate_limiter.acquire()
    client = await get_client()
    response = await client.get(f"{HA_URL}/api/config", headers=get_ha_headers())
    response.raise_for_status()
    data = response.json()
    return data.get("version", "unknown")

@handle_api_errors
async def get_entity_state(
    entity_id: str,
    fields: Optional[List[str]] = None,
    lean: bool = False,
    use_cache: bool = False  # Accept but ignore for now - TODO: implement caching
) -> Dict[str, Any]:
    """
    Get the state of a Home Assistant entity

    TODO: Implement caching logic when use_cache=True
    This would reduce API calls and improve performance for frequently accessed entities
    
    Args:
        entity_id: The entity ID to get
        fields: Optional list of specific fields to include in the response
        lean: If True, returns a token-efficient version with minimal fields
              (overridden by fields parameter if provided)
    
    Returns:
        Entity state dictionary, optionally filtered to include only specified fields
    """
    # Fetch directly
    await _rate_limiter.acquire()
    client = await get_client()
    response = await client.get(
        f"{HA_URL}/api/states/{entity_id}",
        headers=get_ha_headers()
    )
    response.raise_for_status()
    entity_data = response.json()
    
    # Apply field filtering if requested
    if fields:
        # User-specified fields take precedence
        return filter_fields(entity_data, fields)
    elif lean:
        # Build domain-specific lean fields
        lean_fields = DEFAULT_LEAN_FIELDS.copy()
        
        # Add domain-specific important attributes
        domain = entity_id.split('.')[0]
        if domain in DOMAIN_IMPORTANT_ATTRIBUTES:
            for attr in DOMAIN_IMPORTANT_ATTRIBUTES[domain]:
                lean_fields.append(f"attr.{attr}")
        
        return filter_fields(entity_data, lean_fields)
    else:
        # Return full entity data
        return entity_data

@handle_api_errors
async def get_entities(
    domain: Optional[str] = None, 
    search_query: Optional[str] = None, 
    limit: int = 100,
    fields: Optional[List[str]] = None,
    lean: bool = True
) -> List[Dict[str, Any]]:
    """
    Get a list of all entities from Home Assistant with optional filtering and search
    
    Args:
        domain: Optional domain to filter entities by (e.g., 'light', 'switch')
        search_query: Optional case-insensitive search term to filter by entity_id, friendly_name or other attributes
        limit: Maximum number of entities to return (default: 100)
        fields: Optional list of specific fields to include in each entity
        lean: If True (default), returns token-efficient versions with minimal fields
    
    Returns:
        List of entity dictionaries, optionally filtered by domain and search terms,
        and optionally limited to specific fields
    """
    # Get all entities directly
    await _rate_limiter.acquire()
    client = await get_client()
    response = await client.get(f"{HA_URL}/api/states", headers=get_ha_headers())
    response.raise_for_status()
    entities = response.json()
    
    # Filter by domain if specified
    if domain:
        entities = [entity for entity in entities if entity["entity_id"].startswith(f"{domain}.")]
    
    # Search if query is provided
    if search_query and search_query.strip():
        search_term = search_query.lower().strip()
        filtered_entities = []
        
        for entity in entities:
            # Search in entity_id
            if search_term in entity["entity_id"].lower():
                filtered_entities.append(entity)
                continue
                
            # Search in friendly_name
            friendly_name = entity.get("attributes", {}).get("friendly_name", "").lower()
            if friendly_name and search_term in friendly_name:
                filtered_entities.append(entity)
                continue
                
            # Search in other common attributes (state, area_id, etc.)
            if search_term in entity.get("state", "").lower():
                filtered_entities.append(entity)
                continue
                
            # Search in other attributes
            for attr_name, attr_value in entity.get("attributes", {}).items():
                # Check if attribute value can be converted to string
                if isinstance(attr_value, (str, int, float, bool)):
                    if search_term in str(attr_value).lower():
                        filtered_entities.append(entity)
                        break
        
        entities = filtered_entities
    
    # Apply the limit
    if limit > 0 and len(entities) > limit:
        entities = entities[:limit]
    
    # Apply field filtering if requested
    if fields:
        # Use explicit field list when provided
        return [filter_fields(entity, fields) for entity in entities]
    elif lean:
        # Apply domain-specific lean fields to each entity
        result = []
        for entity in entities:
            # Get the entity's domain
            entity_domain = entity["entity_id"].split('.')[0]
            
            # Start with basic lean fields
            lean_fields = DEFAULT_LEAN_FIELDS.copy()
            
            # Add domain-specific important attributes
            if entity_domain in DOMAIN_IMPORTANT_ATTRIBUTES:
                for attr in DOMAIN_IMPORTANT_ATTRIBUTES[entity_domain]:
                    lean_fields.append(f"attr.{attr}")
            
            # Filter and add to result
            result.append(filter_fields(entity, lean_fields))
        
        return result
    else:
        # Return full entities
        return entities

@handle_api_errors
async def call_service(domain: str, service: str, data: Optional[Dict[str, Any]] = None) -> Dict[str, Any]:
    """Call a Home Assistant service"""
    if data is None:
        data = {}

    await _rate_limiter.acquire()
    client = await get_client()
    response = await client.post(
        f"{HA_URL}/api/services/{domain}/{service}", 
        headers=get_ha_headers(),
        json=data
    )
    response.raise_for_status()
    
    # Invalidate cache after service calls as they might change entity states
    global _entities_timestamp
    _entities_timestamp = 0
    
    return response.json()

@handle_api_errors
async def summarize_domain(domain: str, example_limit: int = 3) -> Dict[str, Any]:
    """
    Generate a summary of entities in a domain
    
    Args:
        domain: The domain to summarize (e.g., 'light', 'switch')
        example_limit: Maximum number of examples to include for each state
        
    Returns:
        Dictionary with summary information
    """
    entities = await get_entities(domain=domain)
    
    # Check if we got an error response
    if isinstance(entities, dict) and "error" in entities:
        return entities  # Just pass through the error
    
    try:
        # Initialize summary data
        total_count = len(entities)
        state_counts = {}
        state_examples = {}
        attributes_summary = {}
        
        # Process entities to build the summary
        for entity in entities:
            state = entity.get("state", "unknown")
            
            # Count states
            if state not in state_counts:
                state_counts[state] = 0
                state_examples[state] = []
            state_counts[state] += 1
            
            # Add examples (up to the limit)
            if len(state_examples[state]) < example_limit:
                example = {
                    "entity_id": entity["entity_id"],
                    "friendly_name": entity.get("attributes", {}).get("friendly_name", entity["entity_id"])
                }
                state_examples[state].append(example)
            
            # Collect attribute keys for summary
            for attr_key in entity.get("attributes", {}):
                if attr_key not in attributes_summary:
                    attributes_summary[attr_key] = 0
                attributes_summary[attr_key] += 1
        
        # Create the summary
        summary = {
            "domain": domain,
            "total_count": total_count,
            "state_distribution": state_counts,
            "examples": state_examples,
            "common_attributes": sorted(
                [(k, v) for k, v in attributes_summary.items()], 
                key=lambda x: x[1], 
                reverse=True
            )[:10]  # Top 10 most common attributes
        }
        
        return summary
    except Exception as e:
        return {"error": f"Error generating domain summary: {str(e)}"}

@handle_api_errors
async def get_automations() -> List[Dict[str, Any]]:
    """Get a list of all automations from Home Assistant"""
    # Reuse the get_entities function with domain filtering
    automation_entities = await get_entities(domain="automation")
    
    # Check if we got an error response
    if isinstance(automation_entities, dict) and "error" in automation_entities:
        return automation_entities  # Just pass through the error
    
    # Process automation entities
    result = []
    try:
        for entity in automation_entities:
            # Extract relevant information
            automation_info = {
                "id": entity["entity_id"].split(".")[1],
                "entity_id": entity["entity_id"],
                "state": entity["state"],
                "alias": entity["attributes"].get("friendly_name", entity["entity_id"]),
            }
            
            # Add any additional attributes that might be useful
            if "last_triggered" in entity["attributes"]:
                automation_info["last_triggered"] = entity["attributes"]["last_triggered"]
            
            result.append(automation_info)
    except (TypeError, KeyError) as e:
        # Handle errors in processing the entities
        return {"error": f"Error processing automation entities: {str(e)}"}
        
    return result

@handle_api_errors
async def reload_automations() -> Dict[str, Any]:
    """Reload all automations in Home Assistant"""
    return await call_service("automation", "reload", {})

@handle_api_errors
async def restart_home_assistant() -> Dict[str, Any]:
    """Restart Home Assistant"""
    return await call_service("homeassistant", "restart", {})

@handle_api_errors
async def get_hass_error_log() -> Dict[str, Any]:
    """
    Get the Home Assistant error log for troubleshooting using WebSocket API

    Returns:
        A dictionary containing:
        - records: List of error/warning log records
        - error_count: Number of ERROR entries found
        - warning_count: Number of WARNING entries found
        - integration_mentions: Map of integration names to mention counts
        - error: Error message if retrieval failed
    """
    try:
        # Use WebSocket API to retrieve system_log records
        # call_websocket_api returns result["result"] directly
        records = await call_websocket_api("system_log/list")

        if not records or not isinstance(records, list):
            return {
                "error": "Failed to retrieve system log records",
                "records": [],
                "error_count": 0,
                "warning_count": 0,
                "integration_mentions": {}
            }

        # Count errors and warnings
        error_count = sum(1 for r in records if r.get("level") == "ERROR")
        warning_count = sum(1 for r in records if r.get("level") == "WARNING")

        # Extract integration mentions from logger names
        integration_mentions = {}
        for record in records:
            logger_name = record.get("name", "")
            # Extract integration name from logger like "homeassistant.components.mqtt"
            if "homeassistant.components." in logger_name:
                integration = logger_name.split("homeassistant.components.")[1].split(".")[0]
                integration_mentions[integration] = integration_mentions.get(integration, 0) + 1

        return {
            "records": records,
            "error_count": error_count,
            "warning_count": warning_count,
            "integration_mentions": integration_mentions
        }
    except Exception as e:
        logger.error(f"Error retrieving Home Assistant error log: {str(e)}")
        return {
            "error": f"Error retrieving error log: {str(e)}",
            "records": [],
            "error_count": 0,
            "warning_count": 0,
            "integration_mentions": {}
        }


@handle_api_errors
async def list_automation_traces(
    automation_id: str,
    domain: str = "automation",
    limit: int = 10
) -> Dict[str, Any]:
    """
    List recent execution traces for a specific automation.

    Token-efficient: Returns only essential fields, limited results.

    Args:
        automation_id: REQUIRED - The automation ID (e.g., 'motion_light' or 'automation.motion_light')
        domain: Domain to query ('automation' or 'script'). Default: 'automation'
        limit: Maximum traces to return (default: 10, max: 50)

    Returns:
        A dictionary containing:
        - traces: List of lean summaries (run_id, timestamp, trigger, state, outcome)
        - automation_id: The automation queried
        - domain: The domain queried
        - count: Number of traces returned
    """
    try:
        # Enforce limit bounds
        limit = min(max(1, limit), 50)

        # Strip domain prefix if present
        if automation_id.startswith(f"{domain}."):
            automation_id = automation_id.split(".", 1)[1]

        # Call trace/list WebSocket API with specific item_id
        traces = await call_websocket_api(
            "trace/list",
            domain=domain,
            item_id=automation_id
        )

        if not traces:
            return {
                "automation_id": automation_id,
                "domain": domain,
                "traces": [],
                "count": 0
            }

        # Extract traces for this automation
        raw_traces = []
        if isinstance(traces, dict) and automation_id in traces:
            raw_traces = traces[automation_id]
        elif isinstance(traces, list):
            raw_traces = traces

        # Build lean trace summaries - only essential fields
        lean_traces = []
        for trace in raw_traces[:limit]:
            lean_trace = {
                "run_id": trace.get("run_id"),
                "timestamp": trace.get("timestamp", {}).get("start"),
                "trigger": trace.get("trigger"),
                "state": trace.get("state"),
            }
            # Include outcome if present
            if trace.get("script_execution"):
                lean_trace["outcome"] = trace["script_execution"]
            lean_traces.append(lean_trace)

        return {
            "automation_id": automation_id,
            "domain": domain,
            "traces": lean_traces,
            "count": len(lean_traces)
        }

    except Exception as e:
        logger.error(f"Error listing automation traces: {str(e)}")
        return {
            "automation_id": automation_id,
            "domain": domain,
            "traces": [],
            "count": 0,
            "error": f"Error listing traces: {str(e)}"
        }


@handle_api_errors
async def get_automation_trace(
    automation_id: str,
    run_id: str,
    domain: str = "automation"
) -> Dict[str, Any]:
    """
    Get detailed trace information for a specific automation run.

    Uses the WebSocket API to retrieve full execution details including
    trigger info, condition results, action execution, and any errors.

    Args:
        automation_id: The automation ID (e.g., 'motion_light' or 'automation.motion_light')
        run_id: The specific run/trace ID to retrieve
        domain: The domain ('automation' or 'script'). Default: 'automation'

    Returns:
        A dictionary containing the full trace with:
        - trace: Complete trace data including trigger, conditions, actions
        - automation_id: The automation ID
        - run_id: The run ID
        - domain: The domain
        - error: Error message if retrieval failed
    """
    try:
        # Strip domain prefix if present
        if automation_id.startswith(f"{domain}."):
            automation_id = automation_id.split(".", 1)[1]

        # Call trace/get WebSocket API
        trace = await call_websocket_api(
            "trace/get",
            domain=domain,
            item_id=automation_id,
            run_id=run_id
        )

        return {
            "automation_id": automation_id,
            "run_id": run_id,
            "domain": domain,
            "trace": trace
        }

    except Exception as e:
        logger.error(f"Error getting automation trace: {str(e)}")
        return {
            "automation_id": automation_id,
            "run_id": run_id,
            "domain": domain,
            "trace": None,
            "error": f"Error getting trace: {str(e)}"
        }


def parse_datetime(dt_input: Union[str, datetime]) -> datetime:
    """
    Parse datetime input to timezone-aware datetime object.
    Simple implementation supporting ISO 8601 and basic keywords.

    Supported formats:
    - ISO 8601: "2025-10-28T10:00:00Z", "2025-10-28T10:00:00+00:00"
    - Date only: "2025-10-28" (assumes start of day in UTC)
    - Keywords: "now", "today", "yesterday"

    Returns:
        datetime: Timezone-aware datetime in UTC
    """
    if isinstance(dt_input, datetime):
        return dt_input if dt_input.tzinfo else dt_input.replace(tzinfo=timezone.utc)

    dt_str = str(dt_input).strip()

    # Handle special keywords
    if dt_str.lower() == "now":
        return datetime.now(timezone.utc)
    if dt_str.lower() == "today":
        return datetime.now(timezone.utc).replace(hour=0, minute=0, second=0, microsecond=0)
    if dt_str.lower() == "yesterday":
        today = datetime.now(timezone.utc).replace(hour=0, minute=0, second=0, microsecond=0)
        return today - timedelta(days=1)

    # Handle date-only format (YYYY-MM-DD)
    if len(dt_str) == 10 and dt_str[4] == '-' and dt_str[7] == '-':
        dt = datetime.strptime(dt_str, "%Y-%m-%d")
        return dt.replace(tzinfo=timezone.utc)

    # Handle ISO 8601
    if dt_str.endswith('Z'):
        dt_str = dt_str[:-1] + '+00:00'

    try:
        dt = datetime.fromisoformat(dt_str)
        return dt if dt.tzinfo else dt.replace(tzinfo=timezone.utc)
    except ValueError:
        raise ValueError(
            f"Invalid datetime format: {dt_input}. "
            f"Use ISO 8601 (e.g., '2025-10-28T10:00:00Z'), "
            f"date only (e.g., '2025-10-28'), "
            f"or keywords: 'now', 'today', 'yesterday'"
        )

@handle_api_errors
async def get_entity_history_range(
    entity_id: str,
    start_time: Union[str, datetime],
    end_time: Optional[Union[str, datetime]] = None,
    minimal_response: bool = True
) -> List[Dict[str, Any]]:
    """
    Get entity history for a specific date/time range.

    Args:
        entity_id: The entity ID to get history for
        start_time: ISO 8601 string or datetime object
        end_time: ISO 8601 string or datetime object (defaults to now)
        minimal_response: Reduce response size (default: True)

    Returns:
        A list of state change objects, or an error dictionary
    """
    await _rate_limiter.acquire()
    client = await get_client()

    # Parse start_time
    start_dt = parse_datetime(start_time)

    # Parse end_time (default to now)
    if end_time is None:
        end_dt = datetime.now(timezone.utc)
    else:
        end_dt = parse_datetime(end_time)

    # Validate time range
    if start_dt >= end_dt:
        raise ValueError(f"start_time must be before end_time")

    # Format for API
    start_time_iso = start_dt.strftime("%Y-%m-%dT%H:%M:%SZ")
    end_time_iso = end_dt.strftime("%Y-%m-%dT%H:%M:%SZ")

    # Construct the API URL
    url = f"{HA_URL}/api/history/period/{start_time_iso}"

    # Set query parameters
    params = {
        "filter_entity_id": entity_id,
        "minimal_response": str(minimal_response).lower(),
        "end_time": end_time_iso,
    }

    # Make the API call
    response = await client.get(url, headers=get_ha_headers(), params=params)
    response.raise_for_status()

    # Return the JSON response
    return response.json()

@handle_api_errors
async def get_entity_history(entity_id: str, hours: int) -> List[Dict[str, Any]]:
    """
    Get the history of an entity's state changes from Home Assistant.

    Args:
        entity_id: The entity ID to get history for.
        hours: Number of hours of history to retrieve.

    Returns:
        A list of state change objects, or an error dictionary.
    """
    # Calculate time range
    end_time = datetime.now(timezone.utc)
    start_time = end_time - timedelta(hours=hours)

    # Delegate to the range function
    return await get_entity_history_range(
        entity_id=entity_id,
        start_time=start_time,
        end_time=end_time,
        minimal_response=True
    )

@handle_api_errors
async def get_entity_statistics(
    entity_id: str,
    hours: int,
    period: str = "5minute"
) -> Dict[str, Any]:
    """
    Get statistical data for an entity for recent time period.

    Args:
        entity_id: The entity ID to get statistics for
        hours: Number of hours of statistics to retrieve
        period: Statistics period: "5minute" or "hour"

    Returns:
        A dictionary containing statistical data with aggregated values
    """
    # Calculate time range
    end_time = datetime.now(timezone.utc)
    start_time = end_time - timedelta(hours=hours)

    # Delegate to the range function
    return await get_entity_statistics_range(
        entity_id=entity_id,
        start_time=start_time,
        end_time=end_time,
        period=period
    )

@handle_api_errors
async def get_entity_statistics_range(
    entity_id: str,
    start_time: Union[str, datetime],
    end_time: Optional[Union[str, datetime]] = None,
    period: str = "hour"
) -> Dict[str, Any]:
    """
    Get statistical data for an entity for a specific date/time range.

    Args:
        entity_id: The entity ID to get statistics for
        start_time: ISO 8601 string or datetime object
        end_time: ISO 8601 string or datetime object (defaults to now)
        period: Statistics period: "5minute", "hour", "day", "week", or "month"

    Returns:
        A dictionary containing:
        - entity_id: The entity ID requested
        - period: The period requested
        - start_time: The actual start time used
        - end_time: The actual end time used
        - statistics: List of statistical data points with mean, min, max values
    """
    # Parse start_time
    start_dt = parse_datetime(start_time)

    # Parse end_time (default to now)
    if end_time is None:
        end_dt = datetime.now(timezone.utc)
    else:
        end_dt = parse_datetime(end_time)

    # Validate time range
    if start_dt >= end_dt:
        raise ValueError(f"start_time must be before end_time")

    # Format for API
    start_time_iso = start_dt.strftime("%Y-%m-%dT%H:%M:%SZ")
    end_time_iso = end_dt.strftime("%Y-%m-%dT%H:%M:%SZ")

    # Map our period names to HA's expected values
    period_map = {
        "5minute": "5minute",
        "hour": "hour",
        "day": "day",
        "week": "week",
        "month": "month"
    }

    if period not in period_map:
        raise ValueError(f"Invalid period: {period}. Must be one of: {list(period_map.keys())}")

    try:
        # Use WebSocket API to get statistics
        result = await call_websocket_api(
            "recorder/statistics_during_period",
            start_time=start_time_iso,
            end_time=end_time_iso,
            statistic_ids=[entity_id],
            period=period_map[period],
            types=["mean", "min", "max", "state", "sum"]
        )

        # Extract statistics from the response
        statistics = []
        if entity_id in result:
            statistics = result[entity_id]
        elif isinstance(result, dict) and len(result) > 0:
            # Sometimes returns with different key format
            first_key = list(result.keys())[0]
            statistics = result[first_key]

        return {
            "entity_id": entity_id,
            "period": period,
            "start_time": start_time_iso,
            "end_time": end_time_iso,
            "statistics": statistics
        }

    except Exception as e:
        logger.error(f"Error getting statistics for {entity_id}: {str(e)}")
        # Return empty statistics with error message
        return {
            "entity_id": entity_id,
            "period": period,
            "start_time": start_time_iso,
            "end_time": end_time_iso,
            "statistics": [],
            "error": f"Failed to retrieve statistics: {str(e)}"
        }

@handle_api_errors
async def get_system_overview() -> Dict[str, Any]:
    """
    Get a comprehensive overview of the entire Home Assistant system
    
    Returns:
        A dictionary containing:
        - total_entities: Total count of all entities
        - domains: Dictionary of domains with their entity counts and state distributions
        - domain_samples: Representative sample entities for each domain (2-3 per domain)
        - domain_attributes: Common attributes for each domain
        - area_distribution: Entities grouped by area (if available)
    """
    try:
        # Get ALL entities with minimal fields for efficiency
        # We retrieve all entities since API calls don't consume tokens, only responses do
        await _rate_limiter.acquire()
        client = await get_client()
        response = await client.get(f"{HA_URL}/api/states", headers=get_ha_headers())
        response.raise_for_status()
        all_entities_raw = response.json()
        
        # Apply lean formatting to reduce token usage in the response
        all_entities = []
        for entity in all_entities_raw:
            domain = entity["entity_id"].split(".")[0]
            
            # Start with basic lean fields
            lean_fields = ["entity_id", "state", "attr.friendly_name"]
            
            # Add domain-specific important attributes
            if domain in DOMAIN_IMPORTANT_ATTRIBUTES:
                for attr in DOMAIN_IMPORTANT_ATTRIBUTES[domain]:
                    lean_fields.append(f"attr.{attr}")
            
            # Filter and add to result
            all_entities.append(filter_fields(entity, lean_fields))
        
        # Initialize overview structure
        overview = {
            "total_entities": len(all_entities),
            "domains": {},
            "domain_samples": {},
            "domain_attributes": {},
            "area_distribution": {}
        }
        
        # Group entities by domain
        domain_entities = {}
        for entity in all_entities:
            domain = entity["entity_id"].split(".")[0]
            if domain not in domain_entities:
                domain_entities[domain] = []
            domain_entities[domain].append(entity)
        
        # Process each domain
        for domain, entities in domain_entities.items():
            # Count entities in this domain
            count = len(entities)
            
            # Collect state distribution
            state_distribution = {}
            for entity in entities:
                state = entity.get("state", "unknown")
                if state not in state_distribution:
                    state_distribution[state] = 0
                state_distribution[state] += 1
            
            # Store domain information
            overview["domains"][domain] = {
                "count": count,
                "states": state_distribution
            }
            
            # Select representative samples (2-3 per domain)
            sample_limit = min(3, count)
            samples = []
            for i in range(sample_limit):
                entity = entities[i]
                samples.append({
                    "entity_id": entity["entity_id"],
                    "state": entity.get("state", "unknown"),
                    "friendly_name": entity.get("attributes", {}).get("friendly_name", entity["entity_id"])
                })
            overview["domain_samples"][domain] = samples
            
            # Collect common attributes for this domain
            attribute_counts = {}
            for entity in entities:
                for attr in entity.get("attributes", {}):
                    if attr not in attribute_counts:
                        attribute_counts[attr] = 0
                    attribute_counts[attr] += 1
            
            # Get top 5 most common attributes for this domain
            common_attributes = sorted(attribute_counts.items(), key=lambda x: x[1], reverse=True)[:5]
            overview["domain_attributes"][domain] = [attr for attr, count in common_attributes]
            
            # Group by area if available
            for entity in entities:
                area_id = entity.get("attributes", {}).get("area_id", "Unknown")
                area_name = entity.get("attributes", {}).get("area_name", area_id)
                
                if area_name not in overview["area_distribution"]:
                    overview["area_distribution"][area_name] = {}
                
                if domain not in overview["area_distribution"][area_name]:
                    overview["area_distribution"][area_name][domain] = 0
                    
                overview["area_distribution"][area_name][domain] += 1
        
        # Add summary information
        overview["domain_count"] = len(domain_entities)
        overview["most_common_domains"] = sorted(
            [(domain, len(entities)) for domain, entities in domain_entities.items()],
            key=lambda x: x[1],
            reverse=True
        )[:5]
        
        return overview
    except Exception as e:
        logger.error(f"Error generating system overview: {str(e)}")
        return {"error": f"Error generating system overview: {str(e)}"}
