"""
Helper functions for coupling prompts to receipts via anchors.

Anchors enable direct lineage tracking:
- prompt_id (e.g., "C12") → anchor_type (e.g., "program_receipt_id") → receipt_id
- Supports glimpse_id, proof_hash, gatebook_change_id, program_receipt_id

Usage:
    from receipts.anchors import resolve_anchor, validate_anchor_chain
    
    # Find receipts generated by a prompt
    receipts = resolve_anchor("C12", "program_receipt_id", inbox_dir)
    
    # Validate anchor chain integrity
    is_valid = validate_anchor_chain(receipt, prompt_card)
"""

import json
import logging
from datetime import datetime, timezone
from pathlib import Path
from typing import Any, Dict, List, Optional, Set

logger = logging.getLogger(__name__)


ALLOWED_ANCHOR_TYPES = {
    "glimpse_id",        # Links to Glimpse protocol receipts
    "proof_hash",        # Links to DIP/proof chain
    "gatebook_change_id", # Links to gatebook change receipts
    "program_receipt_id", # Links to general program execution receipts
}


def resolve_anchor(
    prompt_id: str,
    anchor_type: str,
    inbox_dir: Path,
    limit: int = 100,
) -> List[Dict[str, Any]]:
    """
    Resolve receipts matching prompt_id and anchor_type.
    
    Args:
        prompt_id: Prompt identifier (e.g., "C12", "R1")
        anchor_type: Type of anchor (must be in ALLOWED_ANCHOR_TYPES)
        inbox_dir: Receipt inbox directory to search
        limit: Max receipts to return
        
    Returns:
        List of receipt dicts matching the anchor criteria
    """
    if anchor_type not in ALLOWED_ANCHOR_TYPES:
        raise ValueError(f"Invalid anchor_type: {anchor_type}. Must be one of {ALLOWED_ANCHOR_TYPES}")
    
    if not inbox_dir.exists():
        logger.warning(f"Inbox directory does not exist: {inbox_dir}")
        return []
    
    matching_receipts = []
    json_files = list(inbox_dir.glob("**/*.json"))[:limit * 2]  # Search more, filter later
    
    for path in json_files:
        try:
            with open(path, "r") as f:
                receipt = json.load(f)
            
            # Check if receipt has prompt_id in context
            context = receipt.get("context", {})
            receipt_prompt_id = context.get("prompt_id")
            
            if receipt_prompt_id == prompt_id:
                # Match found
                matching_receipts.append({
                    "receipt_id": receipt.get("receipt_id"),
                    "receipt_type": receipt.get("receipt_type"),
                    "domain": receipt.get("domain"),
                    "ts": receipt.get("ts"),
                    "path": str(path),
                    "anchor_type": anchor_type,
                    "prompt_id": prompt_id,
                })
                
                if len(matching_receipts) >= limit:
                    break
        
        except Exception as e:
            logger.debug(f"Skipping {path}: {e}")
            continue
    
    logger.info(f"Found {len(matching_receipts)} receipts for prompt_id={prompt_id}, anchor_type={anchor_type}")
    return matching_receipts


def validate_anchor_chain(
    receipt: Dict[str, Any],
    prompt_card: Dict[str, Any],
) -> bool:
    """
    Validate that receipt properly links to prompt via anchor.
    
    Args:
        receipt: Receipt dict with context
        prompt_card: Prompt card from pack_v3.yaml
        
    Returns:
        True if anchor chain is valid
    """
    # Check if prompt declares anchors
    declared_anchors = prompt_card.get("anchors", [])
    if not declared_anchors:
        # No anchors declared, nothing to validate
        return True
    
    # Check if receipt has prompt_id in context
    context = receipt.get("context", {})
    receipt_prompt_id = context.get("prompt_id")
    
    if not receipt_prompt_id:
        logger.warning(f"Receipt {receipt.get('receipt_id')} missing prompt_id in context")
        return False
    
    # Validate anchor types are allowed
    for anchor_type in declared_anchors:
        if anchor_type not in ALLOWED_ANCHOR_TYPES:
            logger.error(f"Invalid anchor type: {anchor_type}")
            return False
    
    return True


def list_prompts_with_anchors(pack_data: Dict[str, Any]) -> Dict[str, List[str]]:
    """
    List all prompts with their declared anchor types.
    
    Args:
        pack_data: Full pack_v3.yaml data
        
    Returns:
        Dict mapping prompt_id → list of anchor types
    """
    result = {}
    prompts = pack_data.get("prompts", {})
    
    for prompt_id, card in prompts.items():
        anchors = card.get("anchors")
        if anchors:
            result[prompt_id] = anchors
    
    return result


def get_anchor_coverage_stats(pack_data: Dict[str, Any]) -> Dict[str, Any]:
    """
    Calculate anchor coverage statistics across the prompt pack.
    
    Args:
        pack_data: Full pack_v3.yaml data
        
    Returns:
        Dict with coverage stats
    """
    prompts = pack_data.get("prompts", {})
    total_prompts = len(prompts)
    
    with_anchors = sum(1 for card in prompts.values() if card.get("anchors"))
    by_domain = {}
    anchor_type_counts = {anchor_type: 0 for anchor_type in ALLOWED_ANCHOR_TYPES}
    
    for card in prompts.values():
        domain = card.get("domain", "unknown")
        if domain not in by_domain:
            by_domain[domain] = {"total": 0, "with_anchors": 0}
        
        by_domain[domain]["total"] += 1
        
        if card.get("anchors"):
            by_domain[domain]["with_anchors"] += 1
            for anchor_type in card.get("anchors", []):
                if anchor_type in anchor_type_counts:
                    anchor_type_counts[anchor_type] += 1
    
    return {
        "total_prompts": total_prompts,
        "prompts_with_anchors": with_anchors,
        "coverage_pct": round(100 * with_anchors / total_prompts, 1) if total_prompts > 0 else 0,
        "by_domain": by_domain,
        "anchor_type_usage": anchor_type_counts,
    }


def annotate_receipt_with_prompt(
    receipt_id: str,
    prompt_id: str,
    anchor_type: str,
    sink=None,
) -> Optional[str]:
    """
    Annotate a receipt in telemetry sink with prompt linkage.
    
    Args:
        receipt_id: Receipt to annotate
        prompt_id: Prompt that generated the receipt
        anchor_type: Type of anchor linking them
        sink: Optional TelemetrySink instance
        
    Returns:
        Annotation ID if successful, None otherwise
    """
    if not sink:
        logger.warning("No telemetry sink provided, skipping annotation")
        return None
    
    if anchor_type not in ALLOWED_ANCHOR_TYPES:
        raise ValueError(f"Invalid anchor_type: {anchor_type}")
    
    annotation_id = f"prompt_anchor_{receipt_id}_{prompt_id}"
    
    try:
        # This assumes receipt already exists in sink
        # We're just adding the prompt linkage annotation
        sink.annotate_metric(
            annotation_id=annotation_id,
            metric_id=None,  # Not metric-specific, receipt-level
            receipt_id=receipt_id,
            prompt_id=prompt_id,
            annotation_type=f"prompt_anchor_{anchor_type}",
            metadata={
                "anchor_type": anchor_type,
                "annotated_at": datetime.now(timezone.utc).isoformat(),
            },
        )
        logger.info(f"Annotated receipt {receipt_id} with prompt {prompt_id} ({anchor_type})")
        return annotation_id
    
    except Exception as e:
        logger.error(f"Failed to annotate receipt {receipt_id}: {e}")
        return None
