"""
Atom builders for AMI SDK.

Factory functions for creating type-safe atoms: state(), tag(), event(), ulid()
"""
import math
import re
from datetime import datetime, timezone
from typing import List, Optional, Tuple, Dict, Any, Literal
from dataclasses import dataclass
from ulid import ULID


# ============================================================================
# Atom Source Type
# ============================================================================

AtomSource = Literal["system", "human", "ml"]


# ============================================================================
# State Atom
# ============================================================================

@dataclass
class StateAtom:
    """
    State atom representing a key-value fact.

    Attributes:
        key: Dotted path (2-7 parts, lowercase, underscores allowed)
        value: Primitive value (str/int/float/bool)
        source: Origin of the atom
    """
    key: str
    value: Any  # str | int | float | bool
    source: AtomSource = "system"

    def __post_init__(self):
        _validate_state_key(self.key)
        if isinstance(self.value, str) and len(self.value) > 256:
            raise ValueError("String values must be ≤ 256 characters")

    def to_dict(self) -> Dict[str, Any]:
        """Serialize to wire format."""
        return {
            "type": "state",
            "key": self.key,
            "value": self.value,
            "source": self.source
        }


def state(key: str, value: Any, *, source: AtomSource = "system") -> StateAtom:
    """
    Create a state atom.

    Args:
        key: Dotted path (e.g., "user.id", "support.ticket.status")
        value: Primitive value (str, int, float, bool)
        source: Atom source ("system", "human", "ml")

    Returns:
        StateAtom instance

    Example:
        state("user.id", "U-123")
        state("cart.total", 99.99, source="system")
    """
    return StateAtom(key=key, value=value, source=source)


def atoms_from_dict(
    data: dict,
    prefix: str,
    *,
    source: AtomSource = "system",
) -> List[StateAtom]:
    """Convert a flat or nested dict to a list of state atoms.

    Recursively flattens nested dicts using dot notation. Skips None, lists,
    fields starting with ``_``, and invalid field names. Mirrors the TypeScript
    ``atomsFromDict`` helper.

    Args:
        data: Object to convert.
        prefix: Key prefix (1-6 lowercase dot-separated parts).
        source: Atom source ("system", "human", "ml").

    Returns:
        List of StateAtom instances.

    Example::

        atoms_from_dict({"id": "U-123", "profile": {"tier": "premium"}}, "user")
        # [StateAtom(key="user.id", ...), StateAtom(key="user.profile.tier", ...)]
    """
    PREFIX_PATTERN = re.compile(r"^[a-z][a-z0-9_]*(\.[a-z][a-z0-9_]*){0,5}$")
    if not PREFIX_PATTERN.match(prefix):
        raise ValueError(
            f"Invalid prefix (must be 1-6 lowercase dot-separated parts): {prefix}"
        )

    FIELD_PATTERN = re.compile(r"^[a-z][a-z0-9_]*$")
    atoms: List[StateAtom] = []
    prefix_depth = len(prefix.split("."))

    def flatten(obj: dict, current_prefix: str, depth: int) -> None:
        for key, value in obj.items():
            if key.startswith("_"):
                continue
            if not FIELD_PATTERN.match(key):
                continue

            full_key = f"{current_prefix}.{key}"
            total_depth = depth + 1

            if total_depth > 7:
                continue

            if value is None:
                continue
            elif isinstance(value, bool):
                atoms.append(state(full_key, value, source=source))
            elif isinstance(value, str):
                if len(value) <= 256:
                    atoms.append(state(full_key, value, source=source))
            elif isinstance(value, (int, float)):
                if math.isfinite(value):
                    atoms.append(state(full_key, value, source=source))
            elif isinstance(value, dict):
                flatten(value, full_key, total_depth)
            # lists and other types are silently skipped

    flatten(data, prefix, prefix_depth)
    return atoms


# ============================================================================
# Tag Atom
# ============================================================================

@dataclass
class TagAtom:
    """
    Tag atom representing categorical metadata.

    Attributes:
        kind: Tag category (≤64 chars)
        value: Tag value (≤128 chars)
        source: Origin of the atom
    """
    kind: str
    value: str
    source: AtomSource = "system"

    def __post_init__(self):
        if len(self.kind) > 64:
            raise ValueError("Tag kind must be ≤ 64 characters")
        if len(self.value) > 128:
            raise ValueError("Tag value must be ≤ 128 characters")

    def to_dict(self) -> Dict[str, Any]:
        """Serialize to wire format."""
        return {
            "type": "tag",
            "kind": self.kind,
            "value": self.value,
            "source": self.source
        }


def tag(kind: str, value: str, *, source: AtomSource = "system") -> TagAtom:
    """
    Create a tag atom.

    Args:
        kind: Tag category (e.g., "intent", "persona", "tier")
        value: Tag value
        source: Atom source ("system", "human", "ml")

    Returns:
        TagAtom instance

    Example:
        tag("intent", "escalation")
        tag("segment", "enterprise", source="ml")
    """
    return TagAtom(kind=kind, value=value, source=source)


# ============================================================================
# Event Atom
# ============================================================================

@dataclass
class AnchorRole:
    """Role anchor (type + id)."""
    type: str
    id: str


@dataclass
class EventAnchor:
    """Event anchors for subject and object."""
    subject: Optional[AnchorRole] = None
    object: Optional[AnchorRole] = None


@dataclass
class EventCorrelation:
    """Event correlation metadata."""
    trace_id: Optional[str] = None
    span_id: Optional[str] = None
    parent_span_id: Optional[str] = None
    invocation_id: Optional[str] = None
    activation_id: Optional[str] = None
    action_id: Optional[str] = None
    idempotency_id: Optional[str] = None


@dataclass
class EventAtom:
    """
    Event atom representing a timestamped action.

    Attributes:
        topic: Dot-notation topic (S.V.O or P.S.V.O)
        ts: UTC timestamp
        anchor: Subject/object role anchors
        correlation: Tracing/correlation IDs
        attributes: Additional event data
    """
    topic: str
    ts: datetime
    anchor: Optional[EventAnchor] = None
    correlation: Optional[EventCorrelation] = None
    attributes: Optional[Dict[str, Any]] = None

    def __post_init__(self):
        _validate_event_topic(self.topic)
        # Ensure timestamp is UTC
        if self.ts.tzinfo != timezone.utc:
            raise ValueError("Event timestamp must be UTC")

    def to_dict(self) -> Dict[str, Any]:
        """
        Serialize to wire format.

        Wire format splits topic into {subject, verb, object} tokens.
        """
        tokens = self.topic.split(".")

        # Handle 3-token (S.V.O) or 4-token (P.S.V.O) topics
        if len(tokens) == 3:
            subject, verb, obj = tokens
            data = {
                "type": "event",
                "subject": subject,
                "verb": verb,
                "object": obj,
                "ts": self.ts.strftime("%Y-%m-%dT%H:%M:%SZ")
            }
        else:  # 4 tokens
            preposition, subject, verb, obj = tokens
            data = {
                "type": "event",
                "preposition": preposition,
                "subject": subject,
                "verb": verb,
                "object": obj,
                "ts": self.ts.strftime("%Y-%m-%dT%H:%M:%SZ")
            }

        # Add optional fields
        if self.anchor:
            anchor_data = {}
            if self.anchor.subject:
                anchor_data["subject"] = {
                    "type": self.anchor.subject.type,
                    "id": self.anchor.subject.id
                }
            if self.anchor.object:
                anchor_data["object"] = {
                    "type": self.anchor.object.type,
                    "id": self.anchor.object.id
                }
            data["anchor"] = anchor_data

        if self.correlation:
            corr_data = {}
            if self.correlation.trace_id:
                corr_data["trace_id"] = self.correlation.trace_id
            if self.correlation.span_id:
                corr_data["span_id"] = self.correlation.span_id
            if self.correlation.parent_span_id:
                corr_data["parent_span_id"] = self.correlation.parent_span_id
            if self.correlation.invocation_id:
                corr_data["invocation_id"] = self.correlation.invocation_id
            if self.correlation.activation_id:
                corr_data["activation_id"] = self.correlation.activation_id
            if self.correlation.action_id:
                corr_data["action_id"] = self.correlation.action_id
            if self.correlation.idempotency_id:
                corr_data["idempotency_id"] = self.correlation.idempotency_id
            if corr_data:
                data["correlation"] = corr_data

        if self.attributes:
            data["attributes"] = self.attributes

        return data


def event(
    topic: str,
    *,
    ts: datetime,
    subject_anchor: Optional[Tuple[str, str]] = None,
    object_anchor: Optional[Tuple[str, str]] = None,
    attributes: Optional[Dict[str, Any]] = None,
    correlation_trace_id: Optional[str] = None
) -> EventAtom:
    """
    Create an event atom.

    Args:
        topic: Dot-notation topic (S.V.O or P.S.V.O)
        ts: UTC timestamp
        subject_anchor: Subject anchor as (type, id) tuple
        object_anchor: Object anchor as (type, id) tuple
        attributes: Additional event data
        correlation_trace_id: Trace ID for correlation

    Returns:
        EventAtom instance

    Example:
        event("user.completed.onboarding_step", ts=datetime.now(timezone.utc),
              subject_anchor=("user", "U-123"))
    """
    # Build anchor
    anchor = None
    if subject_anchor or object_anchor:
        anchor = EventAnchor()
        if subject_anchor:
            if not isinstance(subject_anchor, tuple) or len(subject_anchor) != 2:
                raise ValueError("subject_anchor must be a tuple of (type, id)")
            anchor.subject = AnchorRole(type=subject_anchor[0], id=subject_anchor[1])
        if object_anchor:
            if not isinstance(object_anchor, tuple) or len(object_anchor) != 2:
                raise ValueError("object_anchor must be a tuple of (type, id)")
            anchor.object = AnchorRole(type=object_anchor[0], id=object_anchor[1])

    # Build correlation
    correlation = None
    if correlation_trace_id:
        correlation = EventCorrelation(trace_id=correlation_trace_id)

    return EventAtom(
        topic=topic,
        ts=ts,
        anchor=anchor,
        correlation=correlation,
        attributes=attributes
    )


# ============================================================================
# ULID Generator
# ============================================================================

def ulid() -> str:
    """
    Generate a ULID (Universally Unique Lexicographically Sortable Identifier).

    Returns:
        26-character ULID string (timestamp + randomness)

    Example:
        id = ulid()  # "01ARZ3NDEKTSV4RRFFQ69G5FAV"
    """
    return str(ULID())


# ============================================================================
# Validation Helpers
# ============================================================================

# State key pattern: ^[a-z][a-z0-9_]*(\.[a-z][a-z0-9_]*){1,6}$
STATE_KEY_PATTERN = re.compile(r"^[a-z][a-z0-9_]*(\.[a-z][a-z0-9_]*){1,6}$")

# Event token pattern: ^[a-z][a-z0-9_]*$
TOKEN_PATTERN = re.compile(r"^[a-z][a-z0-9_]*$")


def _validate_state_key(key: str) -> None:
    """Validate state atom key format."""
    if not STATE_KEY_PATTERN.match(key):
        parts = key.split(".")
        if len(parts) < 2:
            raise ValueError(
                f"State key must have 2-7 parts (got {len(parts)}): {key}"
            )
        if len(parts) > 7:
            raise ValueError(
                f"State key must have 2-7 parts (got {len(parts)}): {key}"
            )
        raise ValueError(
            f"State key pattern invalid (must be lowercase, start with letter, use underscores): {key}"
        )


def _validate_event_topic(topic: str) -> None:
    """Validate event topic format."""
    tokens = topic.split(".")

    # Must be 3 or 4 tokens
    if len(tokens) not in (3, 4):
        raise ValueError(
            f"Event topic must have 3 or 4 tokens (S.V.O or P.S.V.O), got {len(tokens)}: {topic}"
        )

    # Each token must match pattern
    for token in tokens:
        if not TOKEN_PATTERN.match(token):
            raise ValueError(
                f"Event topic token pattern invalid (must be lowercase, start with letter): {token}"
            )
