"""Distributed tracing support via Correlation IDs."""

import threading
import time
import uuid
from typing import Dict, Any, Optional


class CorrelationTracker:
    """Tracks correlation IDs across agent delegation for distributed tracing."""
    
    def __init__(self):
        self.correlations: Dict[str, Dict[str, Any]] = {}
        self.lock = threading.Lock()
    
    def create_correlation(self, root_id: Optional[str] = None) -> str:
        """Create new correlation chain or extend existing one."""
        corr_id = root_id or str(uuid.uuid4())
        with self.lock:
            if corr_id not in self.correlations:
                self.correlations[corr_id] = {
                    "created_at": time.time(),
                    "chain": [corr_id],
                    "hop_count": 0
                }
            else:
                self.correlations[corr_id]["chain"].append(str(uuid.uuid4()))
                self.correlations[corr_id]["hop_count"] += 1
        return corr_id
    
    def get_chain(self, corr_id: str) -> list:
        """Get full delegation chain for this correlation ID."""
        with self.lock:
            if corr_id in self.correlations:
                return self.correlations[corr_id]["chain"].copy()
        return []
    
    def get_metadata(self, corr_id: str) -> Dict[str, Any]:
        """Get correlation metadata including timing and hop count."""
        with self.lock:
            if corr_id in self.correlations:
                meta = self.correlations[corr_id].copy()
                meta["lifetime_seconds"] = time.time() - meta["created_at"]
                return meta
        return {}
