"""
Temporal query builders for WalkFrame.

Handles time-less, single-timestamp, and dwell-time data scenarios.
Provides edge-level and node-level aggregations with temporal filtering.
"""

from typing import TYPE_CHECKING, Optional, Tuple, Union
from datetime import datetime, time, date

if TYPE_CHECKING:
    from ..core.walk import WalkFrame


def edge_frequencies(walkframe: 'WalkFrame') -> 'WalkFrame':
    """Count distinct walks per edge.

    Aggregates by ``exit_edge`` to count how many distinct walks
    pass through each edge.

    Args:
        walkframe: WalkFrame instance.

    Returns:
        WalkFrame: WalkFrame with aggregation configured.

    Examples:
        >>> walks = WalkFrame("data/walks/*.parquet", network=network)
        >>> frequencies = walks.edge_frequencies().compute()
    """
    walkframe._query_state['select_cols'] = "exit_edge AS edge_id, COUNT(DISTINCT walk_id) AS traffic_count"
    walkframe._query_state['group_by'] = "exit_edge"
    walkframe._query_state['filters'].append("exit_edge IS NOT NULL")
    
    return walkframe


def node_frequencies(walkframe: 'WalkFrame') -> 'WalkFrame':
    """Count distinct walks per node.

    Aggregates by ``node_id`` to count how many distinct walks
    visit each node.

    Args:
        walkframe: WalkFrame instance.

    Returns:
        WalkFrame: WalkFrame with aggregation configured.

    Examples:
        >>> walks = WalkFrame("data/walks/*.parquet", network=network)
        >>> frequencies = walks.node_frequencies().compute()
    """
    walkframe._query_state['select_cols'] = "node_id, COUNT(DISTINCT walk_id) AS visit_count"
    walkframe._query_state['group_by'] = "node_id"
    
    return walkframe


def walk_lengths(walkframe: 'WalkFrame') -> 'WalkFrame':
    """Calculate length (number of steps) for each walk.

    Args:
        walkframe: WalkFrame instance.

    Returns:
        WalkFrame: WalkFrame with aggregation configured.

    Examples:
        >>> walks = WalkFrame("data/walks/*.parquet", network=network)
        >>> lengths = walks.walk_lengths().compute()
    """
    walkframe._query_state['select_cols'] = "walk_id, COUNT(*) AS length"
    walkframe._query_state['group_by'] = "walk_id"
    
    return walkframe


def edge_traffic_at(walkframe: 'WalkFrame', query_time: Union[datetime, time, str]) -> 'WalkFrame':
    """Get traffic on edges at a specific time.

    Handles three data scenarios:

    1. **Time-less data** -- returns empty result or warning.
    2. **Single-timestamp data** -- uses ``LEAD()`` to infer edge traversal time.
    3. **Dwell-time data** -- uses ``timestamp + dwell_time`` to next timestamp.

    Args:
        walkframe: WalkFrame instance.
        query_time: Time to query traffic at.

    Returns:
        WalkFrame: WalkFrame with temporal aggregation configured.

    Examples:
        >>> from datetime import datetime
        >>> walks = WalkFrame("data/walks/*.parquet", network=network)
        >>> traffic = walks.edge_traffic_at(datetime(2024, 1, 1, 8, 30)).compute()
    """
    # Convert to string format for SQL (time type = time-of-day only, no date)
    use_time_of_day = isinstance(query_time, time) and not isinstance(query_time, datetime)
    if isinstance(query_time, datetime):
        time_str = query_time.strftime('%Y-%m-%d %H:%M:%S')
    elif isinstance(query_time, time):
        # Time-only: use HH:MM:SS for time-of-day comparison (every day)
        time_str = query_time.strftime('%H:%M:%S')
    else:
        time_str = str(query_time)
        use_time_of_day = False
    
    # Complex query: use timestamp and LEAD(timestamp) for time window
    # (DuckDB may store dwell_time as BIGINT, so we avoid timestamp + dwell_time)
    if use_time_of_day:
        sql = f"""
        WITH edge_times AS (
            SELECT 
                walk_id,
                sequence_idx,
                exit_edge,
                timestamp AS edge_start_time,
                LEAD(timestamp) OVER (PARTITION BY walk_id ORDER BY sequence_idx) AS edge_end_time
            FROM {walkframe._view_name}
            WHERE exit_edge IS NOT NULL
        )
        SELECT 
            exit_edge AS edge_id,
            COUNT(DISTINCT walk_id) AS traffic_count
        FROM edge_times
        WHERE 
            edge_start_time IS NOT NULL
            AND CAST(edge_start_time AS TIME) <= '{time_str}'
            AND (edge_end_time IS NULL OR CAST(edge_end_time AS TIME) >= '{time_str}')
        GROUP BY exit_edge
        """
    else:
        sql = f"""
        WITH edge_times AS (
            SELECT 
                walk_id,
                sequence_idx,
                exit_edge,
                timestamp AS edge_start_time,
                LEAD(timestamp) OVER (PARTITION BY walk_id ORDER BY sequence_idx) AS edge_end_time
            FROM {walkframe._view_name}
            WHERE exit_edge IS NOT NULL
        )
        SELECT 
            exit_edge AS edge_id,
            COUNT(DISTINCT walk_id) AS traffic_count
        FROM edge_times
        WHERE 
            edge_start_time IS NOT NULL
            AND edge_start_time <= '{time_str}'
            AND (edge_end_time IS NULL OR edge_end_time >= '{time_str}')
        GROUP BY exit_edge
        """
    
    # Store as a custom query (will override normal SQL building)
    walkframe._query_state['custom_sql'] = sql
    
    return walkframe


def filter_time_range(walkframe: 'WalkFrame', time_range: Tuple[Union[datetime, time, str], Union[datetime, time, str]]) -> 'WalkFrame':
    """Filter walks by time range.

    Handles time-less data gracefully by checking if the timestamp
    column exists.

    Args:
        walkframe: WalkFrame instance.
        time_range: ``(start_time, end_time)`` tuple.

    Returns:
        WalkFrame: Filtered WalkFrame.

    Examples:
        >>> from datetime import time, datetime
        >>> walks = WalkFrame("data/walks/*.parquet", network=network)
        >>> filtered = walks.filter_time_range((time(8, 0), time(9, 0))).compute()
        >>> filtered = walks.filter_time_range(
        ...     (datetime(2024, 1, 1, 8, 0), datetime(2024, 1, 1, 9, 0))
        ... ).compute()
    """
    start, end = time_range
    
    # Pure time (no date): filter by time-of-day on every day
    both_time_only = isinstance(start, time) and not isinstance(start, datetime) and isinstance(end, time) and not isinstance(end, datetime)
    if both_time_only:
        start_str = start.strftime('%H:%M:%S')
        end_str = end.strftime('%H:%M:%S')
        walkframe._query_state['filters'].append(
            f"CAST(timestamp AS TIME) >= '{start_str}' AND CAST(timestamp AS TIME) <= '{end_str}'"
        )
    else:
        # Full timestamp or mixed: use timestamp comparison
        def format_time(t):
            if isinstance(t, datetime):
                return t.strftime('%Y-%m-%d %H:%M:%S')
            elif isinstance(t, time):
                from datetime import date
                dt = datetime.combine(date.today(), t)
                return dt.strftime('%Y-%m-%d %H:%M:%S')
            else:
                return str(t)
        start_str = format_time(start)
        end_str = format_time(end)
        walkframe._query_state['filters'].append(
            f"timestamp >= '{start_str}' AND timestamp <= '{end_str}'"
        )
    
    return walkframe
