"""
edge_types.py
-------------
Canonical edge type constants and builders for conversational memory in HIEROMEM.

Enables:
- Temporal queries: "What happened after X?" via TURN_FOLLOWS edges
- Speaker filtering: "What did Alice say?" via SPOKEN_BY edges
- Session scoping: "In session 2, what was discussed?" via IN_SESSION edges

Edge weights are included to enable retrieval ranking boost for:
- Temporally adjacent turns
- Same speaker context
- Same session context

Author: Antigravity
"""

from __future__ import annotations
from typing import Any, Optional
from datetime import datetime
import math

# ==============================================================================
# Canonical Edge Type Constants
# ==============================================================================

# Temporal ordering: prev_episode → curr_episode (forward in time)
TURN_FOLLOWS = "turn_follows"

# Speaker attribution: episode → speaker_node
SPOKEN_BY = "spoken_by"

# Session membership: episode → session_node
IN_SESSION = "in_session"

# Special node prefixes for typed nodes
SPEAKER_NODE_PREFIX = "SPEAKER:"
SESSION_NODE_PREFIX = "SESSION:"


# ==============================================================================
# Edge Builder Class
# ==============================================================================

class ConversationalEdgeBuilder:
    """
    Builds structural edges for conversational memory.

    These edges create the scaffolding needed for temporal, speaker-based,
    and session-based queries in HIEROMEM.

    All edges include weights for retrieval ranking boost.
    """

    @staticmethod
    def build_turn_follows_edge(
        prev_episode_id: str,
        curr_episode_id: str,
        temporal_distance: int = 1,
        weight_decay: float = 0.9,
        namespace: Optional[str] = None,
        **kwargs: Any
    ) -> tuple[str, str, dict[str, Any]]:
        """
        Build a TURN_FOLLOWS edge from previous to current episode.

        Args:
            prev_episode_id: ID of the earlier episode
            curr_episode_id: ID of the later episode
            temporal_distance: Number of turns apart (default 1 for consecutive)
            weight_decay: Decay factor per turn distance (default 0.9)
            namespace: Optional namespace for multi-tenant
            **kwargs: Additional edge attributes

        Returns:
            Tuple of (source, target, attributes) for graph insertion

        Example:
            >>> edge = builder.build_turn_follows_edge("ep_1", "ep_2")
            >>> # edge[2]["weight"] = 1.0 for adjacent turns
            >>> edge = builder.build_turn_follows_edge("ep_1", "ep_3", temporal_distance=2)
            >>> # edge[2]["weight"] = 0.9 (decayed)
        """
        # Weight decays with temporal distance
        weight = weight_decay ** (temporal_distance - 1)

        attrs = {
            "relation_type": TURN_FOLLOWS,
            "relation": TURN_FOLLOWS,
            "weight": weight,
            "temporal_distance": temporal_distance,
            "trust": 1.0,  # Structural edges are fully trusted
            **kwargs
        }

        if namespace:
            attrs["namespace"] = namespace

        return (prev_episode_id, curr_episode_id, attrs)

    @staticmethod
    def build_spoken_by_edge(
        episode_id: str,
        speaker: str,
        weight: float = 1.0,
        namespace: Optional[str] = None,
        **kwargs: Any
    ) -> tuple[str, str, dict[str, Any]]:
        """
        Build a SPOKEN_BY edge from episode to speaker node.

        Args:
            episode_id: ID of the episode
            speaker: Name/ID of the speaker
            weight: Edge weight for retrieval boost (default 1.0)
            namespace: Optional namespace
            **kwargs: Additional edge attributes

        Returns:
            Tuple of (source, target, attributes) for graph insertion
        """
        # Create canonical speaker node ID
        speaker_node = f"{SPEAKER_NODE_PREFIX}{speaker.strip()}"

        attrs = {
            "relation_type": SPOKEN_BY,
            "relation": SPOKEN_BY,
            "weight": weight,
            "speaker": speaker.strip(),
            "trust": 1.0,
            **kwargs
        }

        if namespace:
            attrs["namespace"] = namespace

        return (episode_id, speaker_node, attrs)

    @staticmethod
    def build_in_session_edge(
        episode_id: str,
        session_id: str,
        turn_index: Optional[int] = None,
        weight: float = 1.0,
        session_timestamp: Optional[str] = None,
        namespace: Optional[str] = None,
        **kwargs: Any
    ) -> tuple[str, str, dict[str, Any]]:
        """
        Build an IN_SESSION edge from episode to session node.

        Args:
            episode_id: ID of the episode
            session_id: Session identifier (e.g., "session_1", "2023_05_07")
            turn_index: Optional position within session for ordering
            weight: Edge weight for retrieval boost (default 1.0)
            session_timestamp: Optional timestamp string for the session
            namespace: Optional namespace
            **kwargs: Additional edge attributes

        Returns:
            Tuple of (source, target, attributes) for graph insertion
        """
        # Create canonical session node ID
        session_node = f"{SESSION_NODE_PREFIX}{session_id.strip()}"

        attrs = {
            "relation_type": IN_SESSION,
            "relation": IN_SESSION,
            "weight": weight,
            "session_id": session_id.strip(),
            "trust": 1.0,
            **kwargs
        }

        if turn_index is not None:
            attrs["turn_index"] = turn_index

        if session_timestamp:
            attrs["session_timestamp"] = session_timestamp

        if namespace:
            attrs["namespace"] = namespace

        return (episode_id, session_node, attrs)


# ==============================================================================
# Weight Calculation Utilities
# ==============================================================================

def compute_temporal_weight(
    distance: int,
    decay_rate: float = 0.9,
    min_weight: float = 0.1
) -> float:
    """
    Compute weight for temporal edges based on distance.

    Uses exponential decay: weight = decay_rate ^ (distance - 1)

    Args:
        distance: Number of turns apart (1 = consecutive)
        decay_rate: Decay factor per turn (default 0.9)
        min_weight: Minimum weight floor (default 0.1)

    Returns:
        Weight value in [min_weight, 1.0]
    """
    if distance <= 0:
        return 1.0
    weight = decay_rate ** (distance - 1)
    return max(min_weight, weight)


def compute_recency_weight(
    episode_timestamp: datetime,
    reference_time: Optional[datetime] = None,
    half_life_hours: float = 24.0
) -> float:
    """
    Compute weight based on recency using exponential decay.

    Args:
        episode_timestamp: When the episode occurred
        reference_time: Reference point (default: now)
        half_life_hours: Hours until weight decays to 0.5 (default 24)

    Returns:
        Weight value in (0, 1.0]
    """
    if reference_time is None:
        reference_time = datetime.now()

    age_seconds = (reference_time - episode_timestamp).total_seconds()
    age_hours = max(0, age_seconds / 3600)

    # Exponential decay with half-life
    decay_constant = math.log(2) / half_life_hours
    weight = math.exp(-decay_constant * age_hours)

    return weight


# ==============================================================================
# Edge Type Utilities
# ==============================================================================

def is_structural_edge(relation_type: str) -> bool:
    """Check if an edge type is a structural (non-semantic) edge."""
    return relation_type in {TURN_FOLLOWS, SPOKEN_BY, IN_SESSION}


def get_speaker_from_node(node_id: str) -> Optional[str]:
    """Extract speaker name from a speaker node ID."""
    if node_id.startswith(SPEAKER_NODE_PREFIX):
        return node_id[len(SPEAKER_NODE_PREFIX):]
    return None


def get_session_from_node(node_id: str) -> Optional[str]:
    """Extract session ID from a session node ID."""
    if node_id.startswith(SESSION_NODE_PREFIX):
        return node_id[len(SESSION_NODE_PREFIX):]
    return None
