"""
Lineage Propagation - Task ID generation and receipt chain management.

This module provides utilities for:
- Canonical task_id generation from case_id (the Task ID Charter format)
- LineageContext for tracking parent receipts through a request flow
- Helper functions for wiring lineage into receipts

Task ID Charter Format:
    task_id = f"{domain}:{case_id}"

    Examples:
    - clinical:PA-2024-001
    - governance:freeze-proposal-abc
    - arena:behavior-test-xyz

Usage:
    from receipts.lineage import make_task_id, LineageContext

    # Generate canonical task_id
    task_id = make_task_id(case_id="PA-2024-001", domain="clinical")
    # -> "clinical:PA-2024-001"

    # Track lineage through a flow
    ctx = LineageContext(task_id=task_id)

    # After emitting a receipt, register it
    ctx.register_receipt(receipt_id="gov_abc", receipt_hash="sha256...")

    # Get parent_ids for next receipt in chain
    parent_ids = ctx.get_parent_ids()
"""

from __future__ import annotations

import logging
import threading
from dataclasses import dataclass, field
from typing import Any, Dict, List, Optional

logger = logging.getLogger(__name__)

# Thread-local storage for lineage context propagation
_lineage_context = threading.local()


def make_task_id(case_id: str, domain: str = "clinical") -> str:
    """
    Generate canonical task_id from case_id following Task ID Charter.

    Format: {domain}:{case_id}

    Args:
        case_id: The case/PA/request identifier
        domain: Domain namespace (clinical, governance, arena, etc.)

    Returns:
        Canonical task_id string

    Examples:
        >>> make_task_id("PA-2024-001")
        'clinical:PA-2024-001'
        >>> make_task_id("freeze-abc", domain="governance")
        'governance:freeze-abc'
    """
    if not case_id:
        return ""

    # Clean the case_id (remove any existing domain prefix)
    if ":" in case_id:
        # Already has a domain prefix, return as-is
        return case_id

    return f"{domain}:{case_id}"


def parse_task_id(task_id: str) -> tuple[str, str]:
    """
    Parse a task_id into (domain, case_id).

    Args:
        task_id: Canonical task_id string

    Returns:
        Tuple of (domain, case_id)

    Examples:
        >>> parse_task_id("clinical:PA-2024-001")
        ('clinical', 'PA-2024-001')
    """
    if not task_id or ":" not in task_id:
        return ("unknown", task_id or "")

    parts = task_id.split(":", 1)
    return (parts[0], parts[1])


@dataclass
class LineageContext:
    """
    Context for tracking receipt lineage through a request flow.

    Usage:
        ctx = LineageContext(task_id="clinical:PA-001")

        # Register receipts as they're emitted
        ctx.register_receipt("gov_abc", "sha256...")
        ctx.register_receipt("guard_def", "sha256...")

        # Get parent_ids for the next receipt
        parent_ids = ctx.get_parent_ids()  # ["gov_abc", "guard_def"]

        # Get thread_id for chain grouping
        thread_id = ctx.task_id  # "clinical:PA-001"
    """

    task_id: str
    case_id: Optional[str] = None
    domain: str = "clinical"

    # Receipt chain tracking
    receipt_ids: List[str] = field(default_factory=list)
    receipt_hashes: Dict[str, str] = field(default_factory=dict)

    # Optional metadata
    request_id: Optional[str] = None
    tenant: Optional[str] = None

    def __post_init__(self):
        # Auto-extract case_id from task_id if not provided
        if not self.case_id and self.task_id:
            domain, case_id = parse_task_id(self.task_id)
            self.case_id = case_id
            self.domain = domain

    def register_receipt(
        self,
        receipt_id: str,
        receipt_hash: Optional[str] = None,
    ) -> None:
        """
        Register a receipt that was emitted in this flow.

        Args:
            receipt_id: The receipt_id of the emitted receipt
            receipt_hash: Optional hash for parent_hashes linkage
        """
        if receipt_id and receipt_id not in self.receipt_ids:
            self.receipt_ids.append(receipt_id)
            if receipt_hash:
                self.receipt_hashes[receipt_id] = receipt_hash
            logger.debug(
                "Lineage registered: %s in task %s (chain: %d)",
                receipt_id,
                self.task_id,
                len(self.receipt_ids),
            )

    def get_parent_ids(self) -> List[str]:
        """
        Get parent_ids for the next receipt in chain.

        Returns all previously registered receipt_ids.
        """
        return list(self.receipt_ids)

    def get_parent_hashes(self) -> List[str]:
        """
        Get parent_hashes for the next receipt in chain.
        """
        return [
            self.receipt_hashes[rid]
            for rid in self.receipt_ids
            if rid in self.receipt_hashes
        ]

    def get_latest_parent_id(self) -> Optional[str]:
        """Get the most recently registered receipt_id."""
        return self.receipt_ids[-1] if self.receipt_ids else None

    def get_latest_parent_hash(self) -> Optional[str]:
        """Get the hash of the most recently registered receipt."""
        latest_id = self.get_latest_parent_id()
        return self.receipt_hashes.get(latest_id) if latest_id else None

    def to_receipt_kwargs(self) -> Dict[str, Any]:
        """
        Get kwargs to pass to a receipt constructor for lineage wiring.

        Returns:
            Dict with thread_id, parent_ids, parent_hashes, parent_hash
        """
        kwargs: Dict[str, Any] = {
            "thread_id": self.task_id,
            "parent_ids": self.get_parent_ids(),
        }

        parent_hashes = self.get_parent_hashes()
        if parent_hashes:
            kwargs["parent_hashes"] = parent_hashes
            kwargs["parent_hash"] = parent_hashes[0]

        return kwargs


# ============================================================================
# Thread-local context management
# ============================================================================

def get_current_lineage() -> Optional[LineageContext]:
    """
    Get the current lineage context from thread-local storage.

    Returns:
        LineageContext if one is active, None otherwise
    """
    return getattr(_lineage_context, "current", None)


def set_current_lineage(ctx: Optional[LineageContext]) -> None:
    """
    Set the current lineage context in thread-local storage.

    Args:
        ctx: LineageContext to set, or None to clear
    """
    _lineage_context.current = ctx


def clear_current_lineage() -> None:
    """Clear the current lineage context."""
    _lineage_context.current = None


class lineage_scope:
    """
    Context manager for scoped lineage tracking.

    Usage:
        with lineage_scope(task_id="clinical:PA-001") as ctx:
            # Emit receipts...
            ctx.register_receipt("receipt_1")

            # Nested code can access via get_current_lineage()
            lineage = get_current_lineage()
    """

    def __init__(
        self,
        task_id: Optional[str] = None,
        case_id: Optional[str] = None,
        domain: str = "clinical",
        **kwargs,
    ):
        if task_id:
            self.ctx = LineageContext(task_id=task_id, **kwargs)
        elif case_id:
            self.ctx = LineageContext(
                task_id=make_task_id(case_id, domain),
                case_id=case_id,
                domain=domain,
                **kwargs,
            )
        else:
            self.ctx = None

        self._previous: Optional[LineageContext] = None

    def __enter__(self) -> Optional[LineageContext]:
        self._previous = get_current_lineage()
        if self.ctx:
            set_current_lineage(self.ctx)
        return self.ctx

    def __exit__(self, exc_type, exc_val, exc_tb):
        set_current_lineage(self._previous)
        return False


# ============================================================================
# Receipt wiring helpers
# ============================================================================

def wire_lineage_to_receipt(
    receipt: Any,
    ctx: Optional[LineageContext] = None,
) -> None:
    """
    Wire lineage fields into a receipt model.

    Uses the provided context or falls back to thread-local context.

    Args:
        receipt: Receipt model (must have thread_id, parent_ids fields)
        ctx: Optional LineageContext (uses thread-local if not provided)
    """
    if ctx is None:
        ctx = get_current_lineage()

    if ctx is None:
        return

    # Use object.__setattr__ for frozen models
    try:
        if not getattr(receipt, "thread_id", None):
            object.__setattr__(receipt, "thread_id", ctx.task_id)

        existing_parents = getattr(receipt, "parent_ids", None) or []
        if not existing_parents:
            object.__setattr__(receipt, "parent_ids", ctx.get_parent_ids())

            parent_hashes = ctx.get_parent_hashes()
            if parent_hashes:
                object.__setattr__(receipt, "parent_hashes", parent_hashes)
                if not getattr(receipt, "parent_hash", None):
                    object.__setattr__(receipt, "parent_hash", parent_hashes[0])
    except (AttributeError, TypeError) as exc:
        logger.debug("Could not wire lineage to receipt: %s", exc)


def auto_register_receipt(
    receipt_id: str,
    receipt_hash: Optional[str] = None,
    ctx: Optional[LineageContext] = None,
) -> None:
    """
    Auto-register a receipt in the current lineage context.

    Called after emitting a receipt to track it for downstream parent linkage.

    Args:
        receipt_id: The receipt_id that was emitted
        receipt_hash: Optional hash for parent_hashes linkage
        ctx: Optional context (uses thread-local if not provided)
    """
    if ctx is None:
        ctx = get_current_lineage()

    if ctx is not None:
        ctx.register_receipt(receipt_id, receipt_hash)


__all__ = [
    "make_task_id",
    "parse_task_id",
    "LineageContext",
    "get_current_lineage",
    "set_current_lineage",
    "clear_current_lineage",
    "lineage_scope",
    "wire_lineage_to_receipt",
    "auto_register_receipt",
]
