"""
Policy-Aware Receipt Writer for Loom.

Wraps the existing writer.py with domain policy enforcement, attestation
tagging, and unified proof pipeline.

This module is the bridge between the old fragmented paths and the new
policy-driven architecture. It:
1. Enforces domain policies (from config/domain_policies.yaml)
2. Tags receipts with attestation levels
3. Validates lineage when required
4. Provides a single entry point (emit_with_policy)

The goal is to gradually migrate all emission paths to use this module,
then inline the logic into receipts/writer.py.

Usage:
    from receipts.policy_writer import emit_with_policy
    
    # Emits with clinical domain policy (L3+ required, TSA required)
    signed = emit_with_policy(receipt, domain="clinical")
    
    # Emits with test domain policy (L0 allowed, no requirements)
    signed = emit_with_policy(receipt, domain="test")
"""

from __future__ import annotations

import logging
from dataclasses import dataclass, field
from datetime import datetime, timezone
from typing import Any, Dict, List, Optional, TYPE_CHECKING, Mapping
import ulid

from loom.proof_config import (
    AttestationLevel,
    ProofConfig,
    ProofMode,
    get_proof_config,
)
from receipts.policy_registry import (
    DomainPolicy,
    LineagePolicy,
    get_domain_policy,
    get_policy_registry,
)
from receipts.lineage_governor import (
    LineageProof,
    get_lineage_governor,
)
from receipts.validate import validate_receipt
from receipts.models import BaseReceipt, load_receipt
from receipts.compat.pyd import ConfigDictLike
from receipts.proof_writer import emit_proof_receipt

if TYPE_CHECKING:
    from receipts.models import BaseReceipt

logger = logging.getLogger(__name__)


# =============================================================================
# Errors
# =============================================================================

class PolicyWriterError(Exception):
    """Base error for policy-aware writer."""
    pass


class PolicyViolationError(PolicyWriterError):
    """Raised when receipt violates domain policy."""
    pass


class AttestationError(PolicyWriterError):
    """Raised when required attestation level cannot be achieved."""
    pass


class SchemaViolationError(PolicyWriterError):
    """Raised when schema not allowed for domain."""
    pass


class LineageViolationError(PolicyWriterError):
    """Raised when lineage validation fails."""
    pass


# =============================================================================
# Attestation Result
# =============================================================================

@dataclass
class AttestationResult:
    """
    Result of attestation computation.
    
    Contains the achieved level and what proofs are present.
    """
    level: AttestationLevel
    signed: bool = False
    timestamped: bool = False
    anchored: bool = False
    witnessed: bool = False
    rekor_logged: bool = False
    
    # Key/service IDs
    signer_key_id: Optional[str] = None
    tsa_server: Optional[str] = None
    merkle_batch_id: Optional[str] = None
    
    # Watermark for dev receipts
    is_dev_watermarked: bool = False
    
    def to_proof_block(self) -> Dict[str, Any]:
        """Convert to proof block for receipt embedding."""
        return {
            "attestation_level": self.level.value,
            "signed": self.signed,
            "timestamped": self.timestamped,
            "anchored": self.anchored,
            "witnessed": self.witnessed,
            "rekor_logged": self.rekor_logged,
            "signer_key_id": self.signer_key_id,
            "tsa_server": self.tsa_server,
            "merkle_batch_id": self.merkle_batch_id,
            "is_dev_watermarked": self.is_dev_watermarked,
        }


def compute_attestation(
    signed: bool = False,
    timestamped: bool = False,
    anchored: bool = False,
    witnessed: bool = False,
    rekor_logged: bool = False,
) -> AttestationLevel:
    """Compute attestation level from proof flags."""
    if rekor_logged and witnessed:
        return AttestationLevel.L5_COURT_GRADE
    if witnessed:
        return AttestationLevel.L4_WITNESSED
    if anchored and timestamped and signed:
        return AttestationLevel.L3_ANCHORED
    if timestamped and signed:
        return AttestationLevel.L2_TIMESTAMPED
    if signed:
        return AttestationLevel.L1_SIGNED
    return AttestationLevel.L0_DEV


class PolicyWrappedReceipt(BaseReceipt):
    """
    Generic, extra-tolerant receipt wrapper used when we cannot map a payload
    to a typed receipt. Allows arbitrary fields to flow through while still
    satisfying the writer's BaseReceipt requirement.
    """

    model_config = ConfigDictLike(frozen=True, protected_namespaces=(), extra="allow")
    receipt_type: str = "PolicyWrappedReceipt"


# =============================================================================
# Helpers
# =============================================================================


def _coerce_receipt_dict(receipt: Any) -> Dict[str, Any]:
    """Convert a receipt-like object into a mutable dict without losing fields."""
    if isinstance(receipt, Mapping):
        return dict(receipt)
    if hasattr(receipt, "model_dump"):
        return receipt.model_dump()
    if hasattr(receipt, "dict"):
        return dict(receipt.dict())  # type: ignore[arg-type]
    return dict(receipt)


def _coerce_to_base_receipt(receipt_data: Dict[str, Any], domain: str) -> BaseReceipt:
    """
    Build a BaseReceipt (or tolerant wrapper) from untyped data so the writer
    is never handed a raw dict.
    """
    prepared = dict(receipt_data)
    prepared.setdefault("domain", domain)
    prepared.setdefault("ts", prepared.get("ts") or datetime.now(timezone.utc))
    prepared.setdefault(
        "receipt_type",
        prepared.get("receipt_type") or prepared.get("type") or "PolicyWrappedReceipt",
    )
    try:
        return load_receipt(prepared)
    except Exception:
        return PolicyWrappedReceipt.model_validate(prepared)


def _embed_proof_chain(receipt_data: Dict[str, Any], domain: str) -> None:
    """
    Attach tri-temporal proof chain metadata into the receipt context.
    Uses a deterministic stream key derived from domain and receipt_id.
    """
    try:
        stream_key = f"{domain}/{receipt_data.get('receipt_type') or 'receipt'}/{receipt_data.get('receipt_id')}"
        proof_payload = emit_proof_receipt(
            dict(receipt_data),
            stream_key=stream_key,
            timestamp=False,  # avoid overstating TSA; writer handles TSA separately
            anchor_critical=False,
        )
        proof_block = proof_payload.get("proof_chain")
        if proof_block:
            context = dict(receipt_data.get("context") or {})
            context["proof_chain"] = proof_block
            receipt_data["context"] = context
    except Exception:
        # Proof chain is best-effort; never block emission
        logger.debug("Skipping proof chain embedding", exc_info=True)


def _attach_lineage_if_needed(
    receipt_data: Dict[str, Any],
    parent_ids: List[str],
    lineage_policy: LineagePolicy,
    lineage_validated: bool,
) -> tuple[bool, Optional[LineageProof]]:
    """Run lineage governor when parents are provided."""
    if not parent_ids:
        return lineage_validated, None
    if lineage_validated or receipt_data.get("_lineage_validated"):
        return True, None

    governor = get_lineage_governor()
    try:
        proof = governor.attach_parents(receipt_data, parent_ids, policy=lineage_policy)
    except Exception as exc:
        raise LineageViolationError(str(exc)) from exc
    return True, proof


def _precheck_policy(
    policy: DomainPolicy,
    proof_config: ProofConfig,
    schema_id: Optional[str],
    lineage_validated: bool,
    receipt: Dict[str, Any],
) -> List[str]:
    """Validate policy requirements that can be checked before writing."""
    violations: List[str] = []

    achievable = proof_config.max_achievable_attestation()
    if achievable < policy.min_attestation:
        violations.append(
            f"System cannot achieve required attestation {policy.min_attestation.value} "
            f"(max achievable {achievable.value})"
        )

    if policy.fail_on_missing_key and not proof_config.signing.available:
        violations.append(
            f"Domain '{policy.name}' requires signing key but none is configured"
        )

    if policy.allowed_schemas and schema_id and not policy.allows_schema(schema_id):
        violations.append(
            f"Schema '{schema_id}' not allowed for domain '{policy.name}'. "
            f"Allowed: {policy.allowed_schemas}"
        )

    if policy.require_schema_validation:
        try:
            validate_receipt(receipt)
        except Exception as exc:
            violations.append(f"Schema validation failed: {exc}")

    if policy.require_lineage and not lineage_validated and receipt.get("parents"):
        violations.append(f"Lineage must be validated for domain '{policy.name}'")

    return violations


# =============================================================================
# Policy Enforcement
# =============================================================================

def check_policy_compliance(
    policy: DomainPolicy,
    attestation: AttestationResult,
    has_parents: bool,
    lineage_validated: bool,
    schema_id: Optional[str],
    proof_config: ProofConfig,
) -> List[str]:
    """
    Check if a receipt would comply with domain policy.
    
    Returns list of violation messages (empty = compliant).
    """
    violations = []
    
    # Check attestation level
    if attestation.level < policy.min_attestation:
        violations.append(
            f"Attestation {attestation.level.value} below minimum "
            f"{policy.min_attestation.value} for domain '{policy.name}'"
        )
    
    # Check lineage
    if has_parents and policy.require_lineage and not lineage_validated:
        violations.append(
            f"Domain '{policy.name}' requires lineage validation for receipts with parents"
        )
    
    # Check schema
    if schema_id and policy.allowed_schemas:
        if not policy.allows_schema(schema_id):
            violations.append(
                f"Schema '{schema_id}' not allowed for domain '{policy.name}'. "
                f"Allowed: {policy.allowed_schemas}"
            )
    
    # Check TSA requirement
    if policy.tsa_required and not attestation.timestamped:
        violations.append(
            f"Domain '{policy.name}' requires TSA timestamp but none present"
        )
    
    return violations


# =============================================================================
# Policy-Aware Emission
# =============================================================================

def emit_with_policy(
    receipt: Any,
    domain: str = "default",
    parents: Optional[List[str]] = None,
    lineage_validated: bool = False,
    force_attestation: Optional[AttestationLevel] = None,
    skip_policy_check: bool = False,
) -> Dict[str, Any]:
    """
    Emit a receipt with domain policy enforcement.
    
    This is the primary entry point for policy-aware receipt emission.
    It wraps the existing write_receipt function with policy checks.
    
    Args:
        receipt: Receipt data (dict or Pydantic model)
        domain: Domain name for policy lookup
        parents: Parent receipt IDs (for lineage)
        lineage_validated: Whether lineage was already validated
        force_attestation: Override attestation level (for testing)
        skip_policy_check: Skip policy enforcement (dangerous, for migration)
    
    Returns:
        Enriched receipt dict with attestation metadata
        
    Raises:
        PolicyViolationError: If policy check fails (unless in dev mode)
    """
    proof_config = get_proof_config()
    policy = get_domain_policy(domain)

    # Normalize incoming payload
    receipt_data = _coerce_receipt_dict(receipt)

    # Ensure identifiers and domain tagging
    if "receipt_id" not in receipt_data:
        receipt_data["receipt_id"] = f"rcpt-{ulid.new()}"
    receipt_data.setdefault("domain", domain)
    receipt_data["_domain"] = domain

    # Normalise parents for lineage handling
    raw_parents = parents or receipt_data.get("parents") or receipt_data.get("parent_ids") or []
    if isinstance(raw_parents, (str, bytes)):
        parent_ids = [str(raw_parents)]
    else:
        parent_ids = list(raw_parents)

    # Attach lineage when provided
    lineage_policy = policy.lineage_policy or LineagePolicy()
    lineage_validated, _ = _attach_lineage_if_needed(
        receipt_data,
        parent_ids,
        lineage_policy,
        lineage_validated,
    )
    if lineage_validated:
        receipt_data["_lineage_validated"] = True
        if parent_ids:
            receipt_data["parents"] = parent_ids

    # Embed proof-chain metadata (best-effort)
    _embed_proof_chain(receipt_data, domain)

    schema_id = receipt_data.get("$schema") or receipt_data.get("dataschema")

    if not skip_policy_check:
        violations = _precheck_policy(
            policy=policy,
            proof_config=proof_config,
            schema_id=schema_id,
            lineage_validated=lineage_validated,
            receipt=receipt_data,
        )
        if violations:
            message = "; ".join(violations)
            if proof_config.mode == ProofMode.DEV_STUB:
                logger.warning(
                    f"Policy violations for domain '{domain}' (allowed in dev mode): {message}"
                )
            else:
                raise PolicyViolationError(
                    f"Receipt violates policy for domain '{domain}': {message}"
                )

    if force_attestation:
        receipt_data["_attestation_override"] = force_attestation.value

    # Mark policy metadata
    receipt_data["_policy_domain"] = domain
    receipt_data["_policy_version"] = "v1"

    # Add timestamp if not present
    if "created_at" not in receipt_data:
        receipt_data["created_at"] = datetime.now(timezone.utc).isoformat()

    # Coerce into BaseReceipt before writing
    receipt_model = _coerce_to_base_receipt(receipt_data, domain)

    # Delegate to the existing writer
    try:
        from receipts.writer import write_receipt

        result = write_receipt(receipt_model)
        if isinstance(result, dict):
            return result
    except ImportError:
        logger.warning("receipts.writer not available, returning enriched data only")
    except Exception as exc:
        logger.error("Failed to write receipt with policy enforcement", exc_info=True)
        raise

    # If writer returns None (common path), surface the best-effort payload
    return receipt_data


# =============================================================================
# Decorator for Policy Enforcement
# =============================================================================

def require_attestation(min_level: AttestationLevel):
    """
    Decorator to require minimum attestation level for a function.
    
    Usage:
        @require_attestation(AttestationLevel.L2_TIMESTAMPED)
        def process_clinical_receipt(receipt: Dict[str, Any]):
            ...
    """
    def decorator(func):
        def wrapper(*args, **kwargs):
            # Find receipt in args
            receipt = None
            for arg in args:
                if isinstance(arg, dict) and "_attestation" in arg:
                    receipt = arg
                    break
            
            if receipt is None:
                for v in kwargs.values():
                    if isinstance(v, dict) and "_attestation" in v:
                        receipt = v
                        break
            
            if receipt:
                attestation = receipt.get("_attestation", {})
                level_str = attestation.get("attestation_level", "L0_DEV")
                try:
                    level = AttestationLevel(level_str)
                except ValueError:
                    level = AttestationLevel.L0_DEV
                
                if level < min_level:
                    raise AttestationError(
                        f"Function {func.__name__} requires {min_level.value} "
                        f"but receipt has {level.value}"
                    )
            
            return func(*args, **kwargs)
        
        wrapper.__name__ = func.__name__
        wrapper.__doc__ = func.__doc__
        return wrapper
    
    return decorator


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

def validate_receipt_for_domain(
    receipt: Dict[str, Any],
    domain: str,
) -> List[str]:
    """
    Validate a receipt against domain policy without emitting.
    
    Returns list of issues (empty = valid).
    """
    proof_config = get_proof_config()
    policy = get_domain_policy(domain)
    
    # Check attestation in receipt
    attestation_data = receipt.get("_attestation", {})
    level_str = attestation_data.get("attestation_level", "L0_DEV")
    try:
        level = AttestationLevel(level_str)
    except ValueError:
        level = AttestationLevel.L0_DEV
    
    attestation = AttestationResult(
        level=level,
        signed=attestation_data.get("signed", False),
        timestamped=attestation_data.get("timestamped", False),
        anchored=attestation_data.get("anchored", False),
    )
    
    return check_policy_compliance(
        policy=policy,
        attestation=attestation,
        has_parents=bool(receipt.get("parents")),
        lineage_validated=receipt.get("_lineage_validated", False),
        schema_id=receipt.get("$schema") or receipt.get("dataschema"),
        proof_config=proof_config,
    )


def is_receipt_court_grade(receipt: Dict[str, Any]) -> bool:
    """Check if receipt meets court-grade (L5) attestation."""
    attestation = receipt.get("_attestation", {})
    level_str = attestation.get("attestation_level", "L0_DEV")
    try:
        level = AttestationLevel(level_str)
        return level >= AttestationLevel.L5_COURT_GRADE
    except ValueError:
        return False


def is_receipt_production_ready(receipt: Dict[str, Any]) -> bool:
    """Check if receipt is at least signed (L1+)."""
    attestation = receipt.get("_attestation", {})
    level_str = attestation.get("attestation_level", "L0_DEV")
    try:
        level = AttestationLevel(level_str)
        return level >= AttestationLevel.L1_SIGNED
    except ValueError:
        return False
