"""
helix_cdc/capsules/schema.py
============================

Three-Layer Capsule Schema for Helix-CDC.

Layers:
1. Genotype    - The "species" layer: seed, topology, quant, FibPi3D lattice
2. BloodType   - The "compatibility" layer: domain, safety, compat class, Rh flags
3. InstanceResidual - The "clone" layer: delta bytes for exact instance reconstruction

This is additive - existing capsules without these fields will still load.
"""

from dataclasses import dataclass, field, asdict
from typing import Optional, Dict, Any, List


@dataclass
class Genotype:
    """
    Layer 1: Genotype - The "species" of the model.

    This describes WHAT the model is architecturally.
    """
    model_family: str                      # e.g. 'mistral', 'llama', 'phi'
    topology_id: str                       # e.g. '7b_v0.2', '13b_v1.0'
    quant_profile: str                     # e.g. 'q4_k_m', 'q6_k', 'fp16'
    fibpi3d_lattice: Optional[str] = None  # e.g. 'golden_spiral_v2'
    seed_sha256: Optional[str] = None      # SHA256 of the DNA seed
    dna_seed: Optional[str] = None         # DNA seed hex (e.g. '1f647d91f6')
    entropy_seed: Optional[str] = None     # Entropy seed hex (e.g. 'bf690b54')
    atlas_sha256: Optional[str] = None     # SHA256 of centroid atlas

    def to_dict(self) -> Dict[str, Any]:
        """Convert to dict, omitting None values."""
        return {k: v for k, v in asdict(self).items() if v is not None}


@dataclass
class BloodType:
    """
    Layer 2: BloodType - Compatibility classification.

    Determines what this capsule can safely "transfuse" with.
    Like ABO blood types, this prevents incompatible merges.
    """
    domain: str                            # e.g. 'radiology', 'genomics', 'general-llm'
    safety_class: str                      # e.g. 'clinical', 'research-only', 'toy'
    compat_class: str                      # e.g. 'M7B-Q4KM-CLIN-v1'
    blood_type_code: Optional[str] = None  # e.g. 'O+', 'A-', 'AB+' (from blood bank)
    rh_like_flags: Dict[str, bool] = field(default_factory=dict)
    # Example rh_like_flags:
    #   "no_phi_overlay": True       - don't mix with Phi residuals
    #   "no_pseudo16_mix": True      - don't upcast to fp16 priors
    #   "requires_hospital_prior_v3": True  - needs specific reference data

    def to_dict(self) -> Dict[str, Any]:
        """Convert to dict."""
        d = asdict(self)
        if not d['rh_like_flags']:
            del d['rh_like_flags']
        if d.get('blood_type_code') is None:
            del d['blood_type_code']
        return d


@dataclass
class InstanceResidual:
    """
    Layer 3: InstanceResidual - The "clone mode" layer.

    When mode='clone', the residual bytes are REQUIRED for exact reconstruction.
    When mode='generic', residual is optional (procedural regen only).
    When mode='procedural', no residual - pure seed-based regeneration.

    v0 uses naive XOR diff. Future versions will use structured, compressed residuals.
    """
    mode: str                              # 'procedural' | 'generic' | 'clone'
    residual_sha256: Optional[str] = None  # SHA256 of the residual bytes
    residual_bytes_relpath: Optional[str] = None  # relative path to residual file
    baseline_hash: Optional[str] = None    # SHA256 of baseline (vendor/reference) model
    regen_hash: Optional[str] = None       # SHA256 of regenerated model
    residual_size_bytes: Optional[int] = None  # size of residual file
    codec: str = "xor_v0"                  # 'xor_v0' for now, future: 'tensor_diff_v1'

    def to_dict(self) -> Dict[str, Any]:
        """Convert to dict, omitting None values."""
        return {k: v for k, v in asdict(self).items() if v is not None}


@dataclass
class CryptoDNA:
    """
    Layer 4: CryptoDNA - Crypto-DNA Regenerative Substrate.

    Uses RSA-style padding and factorization for regenerative encoding.
    WO#50 addition.
    """
    dna_b64: str                               # Original Base64 seed
    padded_hex: Optional[str] = None           # Padded bigint as hex
    bigint_modulus: Optional[int] = None       # N' value
    factorization_p: Optional[int] = None      # p factor
    factorization_q: Optional[int] = None      # q factor
    regulatory_map: Dict[str, Any] = field(default_factory=dict)
    fibpi3d_offsets: Optional[Dict[str, float]] = None  # {x, y, z}
    epigenetic_flags: int = 0                  # 32-bit flag field
    crt_layout: Dict[str, Any] = field(default_factory=dict)

    def to_dict(self) -> Dict[str, Any]:
        """Convert to dict, omitting None values."""
        d = {"dna_b64": self.dna_b64}
        if self.padded_hex:
            d["padded_hex"] = self.padded_hex
        if self.bigint_modulus is not None:
            d["bigint_modulus"] = self.bigint_modulus
        if self.factorization_p is not None:
            d["factorization"] = [self.factorization_p, self.factorization_q]
        if self.regulatory_map:
            d["regulatory_map"] = self.regulatory_map
        if self.fibpi3d_offsets:
            d["fibpi3d_offsets"] = self.fibpi3d_offsets
        if self.epigenetic_flags:
            d["epigenetic_flags"] = self.epigenetic_flags
        if self.crt_layout:
            d["crt_layout"] = self.crt_layout
        return d


@dataclass
class CapsuleManifest:
    """
    Extended Capsule Manifest with three-layer schema.

    This is additive to existing .hxz format - existing fields preserved,
    new layers added as optional nested objects.
    """
    # Existing fields (keep compatibility with current .hxz format)
    format: str = "helix_hxz"
    version: str = "v0.2.0"
    kat_schema: Optional[str] = None
    se_spec_hash: Optional[str] = None
    seed_count: Optional[int] = None
    glyph_count: Optional[int] = None
    kat_hash: Optional[str] = None

    # New three-layer schema (all optional for backwards compatibility)
    genotype: Optional[Genotype] = None
    blood_type: Optional[BloodType] = None
    instance_residual: Optional[InstanceResidual] = None
    crypto_dna: Optional[CryptoDNA] = None  # WO#50: Crypto-DNA substrate

    # Manifest entries (existing)
    manifest: List[Dict[str, Any]] = field(default_factory=list)

    # Extra fields for extensibility
    extra: Dict[str, Any] = field(default_factory=dict)

    def to_dict(self) -> Dict[str, Any]:
        """Convert to dict for JSON serialization."""
        d = {
            "format": self.format,
            "version": self.version,
        }

        # Add optional existing fields
        if self.kat_schema:
            d["kat_schema"] = self.kat_schema
        if self.se_spec_hash:
            d["se_spec_hash"] = self.se_spec_hash
        if self.seed_count is not None:
            d["seed_count"] = self.seed_count
        if self.glyph_count is not None:
            d["glyph_count"] = self.glyph_count
        if self.kat_hash:
            d["kat_hash"] = self.kat_hash

        # Add new three-layer schema
        if self.genotype:
            d["genotype"] = self.genotype.to_dict()
        if self.blood_type:
            d["blood_type"] = self.blood_type.to_dict()
        if self.instance_residual:
            d["instance_residual"] = self.instance_residual.to_dict()
        if self.crypto_dna:
            d["crypto_dna"] = self.crypto_dna.to_dict()

        # Add manifest entries
        if self.manifest:
            d["manifest"] = self.manifest

        # Add extra fields
        if self.extra:
            d.update(self.extra)

        return d

    @classmethod
    def from_dict(cls, d: Dict[str, Any]) -> "CapsuleManifest":
        """
        Parse a dict (e.g., from JSON) into a CapsuleManifest.

        Gracefully handles missing fields for backwards compatibility.
        """
        genotype = None
        if "genotype" in d:
            genotype = Genotype(**d["genotype"])

        blood_type = None
        if "blood_type" in d:
            bt = d["blood_type"]
            blood_type = BloodType(
                domain=bt.get("domain", "unknown"),
                safety_class=bt.get("safety_class", "unknown"),
                compat_class=bt.get("compat_class", "UNKNOWN"),
                blood_type_code=bt.get("blood_type_code"),
                rh_like_flags=bt.get("rh_like_flags", {}),
            )

        instance_residual = None
        if "instance_residual" in d:
            ir = d["instance_residual"]
            instance_residual = InstanceResidual(
                mode=ir.get("mode", "procedural"),
                residual_sha256=ir.get("residual_sha256"),
                residual_bytes_relpath=ir.get("residual_bytes_relpath"),
                baseline_hash=ir.get("baseline_hash"),
                regen_hash=ir.get("regen_hash"),
                residual_size_bytes=ir.get("residual_size_bytes"),
                codec=ir.get("codec", "xor_v0"),
            )

        crypto_dna = None
        if "crypto_dna" in d:
            cd = d["crypto_dna"]
            factorization = cd.get("factorization", [None, None])
            crypto_dna = CryptoDNA(
                dna_b64=cd.get("dna_b64", ""),
                padded_hex=cd.get("padded_hex"),
                bigint_modulus=cd.get("bigint_modulus"),
                factorization_p=factorization[0] if len(factorization) > 0 else None,
                factorization_q=factorization[1] if len(factorization) > 1 else None,
                regulatory_map=cd.get("regulatory_map", {}),
                fibpi3d_offsets=cd.get("fibpi3d_offsets"),
                epigenetic_flags=cd.get("epigenetic_flags", 0),
                crt_layout=cd.get("crt_layout", {}),
            )

        # Extract known fields, put rest in extra
        known_keys = {
            "format", "version", "kat_schema", "se_spec_hash",
            "seed_count", "glyph_count", "kat_hash", "manifest",
            "genotype", "blood_type", "instance_residual", "crypto_dna"
        }
        extra = {k: v for k, v in d.items() if k not in known_keys}

        return cls(
            format=d.get("format", "helix_hxz"),
            version=d.get("version", "v0.1.0"),
            kat_schema=d.get("kat_schema"),
            se_spec_hash=d.get("se_spec_hash"),
            seed_count=d.get("seed_count"),
            glyph_count=d.get("glyph_count"),
            kat_hash=d.get("kat_hash"),
            genotype=genotype,
            blood_type=blood_type,
            instance_residual=instance_residual,
            crypto_dna=crypto_dna,
            manifest=d.get("manifest", []),
            extra=extra,
        )


def attach_instance_residual(
    manifest_dict: Dict[str, Any],
    mode: str,
    residual_sha256: str,
    residual_relpath: str,
    baseline_hash: Optional[str] = None,
    regen_hash: Optional[str] = None,
    residual_size_bytes: Optional[int] = None,
    codec: str = "xor_v0",
) -> Dict[str, Any]:
    """
    Attach instance_residual info to a capsule manifest dict.

    This is a helper for updating existing manifests with clone-mode info.

    Args:
        manifest_dict: Existing manifest dict
        mode: 'procedural', 'generic', or 'clone'
        residual_sha256: SHA256 hash of residual bytes
        residual_relpath: Relative path to residual file
        baseline_hash: SHA256 of baseline GGUF (optional)
        regen_hash: SHA256 of regenerated GGUF (optional)
        residual_size_bytes: Size of residual file (optional)
        codec: Residual codec version (default 'xor_v0')

    Returns:
        Updated manifest_dict with instance_residual section
    """
    ir = {
        "mode": mode,
        "residual_sha256": residual_sha256,
        "residual_bytes_relpath": residual_relpath,
        "codec": codec,
    }

    if baseline_hash is not None:
        ir["baseline_hash"] = baseline_hash
    if regen_hash is not None:
        ir["regen_hash"] = regen_hash
    if residual_size_bytes is not None:
        ir["residual_size_bytes"] = residual_size_bytes

    manifest_dict["instance_residual"] = ir

    # Update version to indicate three-layer support
    if manifest_dict.get("version", "").startswith("v0.1"):
        manifest_dict["version"] = "v0.2.0"

    return manifest_dict
