#!/usr/bin/env python3
"""
Helix 14D Object Envelope

Provides a universal wrapper for any object in the Helix/CDC ecosystem.
Maps the original 5D/14D JSON concept to concrete Python objects.

Usage:
    from helix_cdc.object_14d import HelixObject14D

    # Wrap a crystal-encoded tensor
    obj = HelixObject14D.from_crystal_encoding(encoding, tensor_name="blk.0.attn_q")

    # Wrap a raw file
    obj = HelixObject14D.from_file("/path/to/file.bin", object_type="binary")

    # Serialize to JSON
    json_str = obj.to_json()

    # Validate
    obj.validate()
"""

import json
import hashlib
import os
from dataclasses import dataclass, field, asdict
from datetime import datetime
from pathlib import Path
from typing import Optional, List, Dict, Any
import base64


@dataclass
class DataDimension:
    """D1: DATA - The actual payload reference"""
    payload_type: str  # inline_base64, file_ref, crystal_encoded, vault_ref, external_url
    payload_ref: Optional[str] = None
    size_bytes: Optional[int] = None
    mime_type: Optional[str] = None


@dataclass
class FoldingDimension:
    """D2: FOLDING MAP - How data is encoded/decoded"""
    codec: str  # crystal_fibpi3d, q4_k_m, int8, zlib, lz4, none
    num_bases: int = 64
    adapter_mode: str = "int8"
    chunk_layout: Optional[List[int]] = None


@dataclass
class GrowthDimension:
    """D3: GROWTH ORDER - Regeneration sequence"""
    depends_on: List[str] = field(default_factory=list)
    decode_steps: List[str] = field(default_factory=list)
    regrow_function: Optional[str] = None


@dataclass
class EntropyDimension:
    """D4: ENTROPY SHELL - Integrity verification"""
    sha256: str
    crc32: Optional[str] = None
    entropy_bits: Optional[float] = None
    hash_chain: List[str] = field(default_factory=list)


@dataclass
class AuditDimension:
    """D5: AUDIT TRAIL - History and provenance"""
    created_at: str = field(default_factory=lambda: datetime.utcnow().isoformat() + "Z")
    created_by: Optional[str] = None
    version_history: List[Dict] = field(default_factory=list)
    source: Optional[str] = None


@dataclass
class CrossHelixDimension:
    """D6: CROSS-HELIX ANCHORS - References to other objects"""
    parent_id: Optional[str] = None
    siblings: List[str] = field(default_factory=list)
    children: List[str] = field(default_factory=list)
    vault_refs: List[str] = field(default_factory=list)


@dataclass
class TimeDriftDimension:
    """D7: TIME DRIFT - Temporal snapshots and deltas"""
    snapshot_epoch: Optional[int] = None
    delta_from_base: Optional[str] = None
    mutation_count: int = 0
    ttl_seconds: Optional[int] = None


@dataclass
class EntropyLayersDimension:
    """D8: ENTROPY LAYERS - Multiple hash chains"""
    merkle_root: Optional[str] = None
    chunk_hashes: List[str] = field(default_factory=list)
    hmac_signature: Optional[str] = None


@dataclass
class SemanticLinksDimension:
    """D9: SEMANTIC LINKS - Meaning graphs"""
    tags: List[str] = field(default_factory=list)
    biopoetica_class: Optional[str] = None
    meaning_graph_ref: Optional[str] = None


@dataclass
class NetworkCoordsDimension:
    """D10: NETWORK COORDINATES - Position in distributed lattice"""
    node_id: Optional[str] = None
    shard_index: Optional[int] = None
    replication_factor: int = 1
    lattice_address: Optional[str] = None


@dataclass
class EnergyTokensDimension:
    """D11: ENERGY TOKENS - Compute/storage cost tracking"""
    encode_cost_ms: Optional[float] = None
    decode_cost_ms: Optional[float] = None
    storage_cost_bytes: Optional[int] = None
    compute_tokens: Optional[int] = None


@dataclass
class QuantumHintsDimension:
    """D12: QUANTUM HINTS - Regrowth pointers"""
    regrow_seeds: List[str] = field(default_factory=list)
    fallback_paths: List[str] = field(default_factory=list)
    speculative_data: Optional[str] = None


@dataclass
class SpatialGeometryDimension:
    """D13: SPATIAL GEOMETRY - FibPi3D coordinates"""
    fibpi3d_coords: Optional[List[float]] = None  # [x, y, z]
    lattice_index: Optional[int] = None
    dimensional_depth: Optional[float] = None
    basis_seed: str = "crystal"


@dataclass
class CryptoShadowsDimension:
    """D14: CRYPTOGRAPHIC SHADOWS - Privacy/irreversible refs"""
    blinded_id: Optional[str] = None
    access_control: str = "public"
    encryption_key_ref: Optional[str] = None
    redacted_fields: List[str] = field(default_factory=list)


@dataclass
class HelixObject14D:
    """
    Universal 14D metadata envelope for any object in Helix/CDC.

    Every object - tensor, config, binary, log - gets wrapped in this
    structure, providing a unified way to address, verify, and decode
    anything in the ecosystem.
    """
    version: str = "14d-v1.0"
    object_id: str = ""
    object_type: str = "unknown"

    # Base 5D
    data: DataDimension = field(default_factory=lambda: DataDimension(payload_type="none"))
    folding: FoldingDimension = field(default_factory=lambda: FoldingDimension(codec="none"))
    growth: GrowthDimension = field(default_factory=GrowthDimension)
    entropy: EntropyDimension = field(default_factory=lambda: EntropyDimension(sha256=""))
    audit: AuditDimension = field(default_factory=AuditDimension)

    # Extended 14D
    cross_helix: CrossHelixDimension = field(default_factory=CrossHelixDimension)
    time_drift: TimeDriftDimension = field(default_factory=TimeDriftDimension)
    entropy_layers: EntropyLayersDimension = field(default_factory=EntropyLayersDimension)
    semantic_links: SemanticLinksDimension = field(default_factory=SemanticLinksDimension)
    network_coords: NetworkCoordsDimension = field(default_factory=NetworkCoordsDimension)
    energy_tokens: EnergyTokensDimension = field(default_factory=EnergyTokensDimension)
    quantum_hints: QuantumHintsDimension = field(default_factory=QuantumHintsDimension)
    spatial_geometry: SpatialGeometryDimension = field(default_factory=SpatialGeometryDimension)
    crypto_shadows: CryptoShadowsDimension = field(default_factory=CryptoShadowsDimension)

    def __post_init__(self):
        if not self.object_id:
            self.object_id = self._generate_id()

    def _generate_id(self) -> str:
        """Generate unique object ID from content hash."""
        content = f"{self.object_type}:{self.data.payload_type}:{self.entropy.sha256}:{self.audit.created_at}"
        return hashlib.sha256(content.encode()).hexdigest()[:16]

    def to_dict(self) -> Dict[str, Any]:
        """Convert to dictionary."""
        def convert(obj):
            if hasattr(obj, '__dataclass_fields__'):
                return {k: convert(v) for k, v in asdict(obj).items() if v is not None}
            elif isinstance(obj, list):
                return [convert(i) for i in obj]
            else:
                return obj
        return convert(self)

    def to_json(self, indent: int = 2) -> str:
        """Serialize to JSON string."""
        return json.dumps(self.to_dict(), indent=indent)

    def validate(self) -> bool:
        """Validate the object envelope."""
        errors = []

        if not self.object_id:
            errors.append("Missing object_id")

        if not self.entropy.sha256:
            errors.append("Missing entropy.sha256")

        if self.data.payload_type not in ["inline_base64", "file_ref", "crystal_encoded", "vault_ref", "external_url", "none"]:
            errors.append(f"Invalid payload_type: {self.data.payload_type}")

        if errors:
            raise ValueError(f"Validation failed: {errors}")

        return True

    @classmethod
    def from_file(cls, path: str, object_type: str = "binary") -> "HelixObject14D":
        """Create envelope from a file."""
        path = Path(path)

        # Read file and compute hash
        with open(path, "rb") as f:
            content = f.read()

        sha256 = hashlib.sha256(content).hexdigest()

        return cls(
            object_type=object_type,
            data=DataDimension(
                payload_type="file_ref",
                payload_ref=str(path.absolute()),
                size_bytes=len(content),
            ),
            folding=FoldingDimension(codec="none"),
            entropy=EntropyDimension(sha256=sha256),
            audit=AuditDimension(source=str(path)),
        )

    @classmethod
    def from_crystal_encoding(
        cls,
        encoding: Any,
        tensor_name: str,
        parent_model: Optional[str] = None,
    ) -> "HelixObject14D":
        """Create envelope from a crystal encoding."""
        from helix_cdc.crystal_codec import encoding_to_dict

        enc_dict = encoding_to_dict(encoding)

        return cls(
            object_type="tensor",
            data=DataDimension(
                payload_type="crystal_encoded",
                payload_ref=tensor_name,
                size_bytes=enc_dict.get("original_size"),
            ),
            folding=FoldingDimension(
                codec="crystal_fibpi3d",
                num_bases=enc_dict.get("num_bases", 64),
                adapter_mode=enc_dict.get("adapter_mode", "int8"),
                chunk_layout=list(enc_dict.get("original_shape", [])),
            ),
            entropy=EntropyDimension(
                sha256=enc_dict.get("checksum", ""),
                crc32=enc_dict.get("crc32"),
            ),
            cross_helix=CrossHelixDimension(
                parent_id=parent_model,
            ),
            spatial_geometry=SpatialGeometryDimension(
                basis_seed=enc_dict.get("seed", "crystal"),
            ),
            semantic_links=SemanticLinksDimension(
                tags=[tensor_name.split(".")[0]] if "." in tensor_name else [],
            ),
        )

    @classmethod
    def from_json(cls, json_str: str) -> "HelixObject14D":
        """Load from JSON string."""
        data = json.loads(json_str)
        return cls.from_dict(data)

    @classmethod
    def from_dict(cls, data: Dict) -> "HelixObject14D":
        """Load from dictionary."""
        return cls(
            version=data.get("version", "14d-v1.0"),
            object_id=data.get("object_id", ""),
            object_type=data.get("object_type", "unknown"),
            data=DataDimension(**data.get("data", {})),
            folding=FoldingDimension(**data.get("folding", {"codec": "none"})),
            growth=GrowthDimension(**data.get("growth", {})),
            entropy=EntropyDimension(**data.get("entropy", {"sha256": ""})),
            audit=AuditDimension(**data.get("audit", {})),
            cross_helix=CrossHelixDimension(**data.get("cross_helix", {})),
            time_drift=TimeDriftDimension(**data.get("time_drift", {})),
            entropy_layers=EntropyLayersDimension(**data.get("entropy_layers", {})),
            semantic_links=SemanticLinksDimension(**data.get("semantic_links", {})),
            network_coords=NetworkCoordsDimension(**data.get("network_coords", {})),
            energy_tokens=EnergyTokensDimension(**data.get("energy_tokens", {})),
            quantum_hints=QuantumHintsDimension(**data.get("quantum_hints", {})),
            spatial_geometry=SpatialGeometryDimension(**data.get("spatial_geometry", {})),
            crypto_shadows=CryptoShadowsDimension(**data.get("crypto_shadows", {})),
        )


def example_tensor_envelope() -> HelixObject14D:
    """Example: crystal-encoded tensor from Mistral."""
    import numpy as np

    PHI = (1 + np.sqrt(5)) / 2
    PI = np.pi
    GOLDEN_ANGLE = 2 * PI / (PHI ** 2)

    # FibPi3D coordinate for lattice index 0
    z = 1 - 1 / 64
    theta = 0
    r = np.sqrt(1 - z**2)
    coords = [float(r * np.cos(theta)), float(r * np.sin(theta)), float(z)]

    return HelixObject14D(
        object_type="tensor",
        data=DataDimension(
            payload_type="crystal_encoded",
            payload_ref="blk.0.attn_q.weight",
            size_bytes=67108864,  # 4096*4096*4
            mime_type="application/octet-stream",
        ),
        folding=FoldingDimension(
            codec="crystal_fibpi3d",
            num_bases=64,
            adapter_mode="int8",
            chunk_layout=[4096, 4096],
        ),
        growth=GrowthDimension(
            depends_on=[],
            decode_steps=["load_coefficients", "apply_basis", "add_adapter"],
            regrow_function="decode_crystal_to_tensor",
        ),
        entropy=EntropyDimension(
            sha256="a1b2c3d4e5f6789012345678901234567890123456789012345678901234abcd",
            crc32="0xABCD1234",
            entropy_bits=7.2,
        ),
        audit=AuditDimension(
            created_by="helix_cdc",
            source="mistral-7b-instruct-v0.2.Q4_K_M.gguf",
        ),
        cross_helix=CrossHelixDimension(
            parent_id="mistral-7b-crystal-pack",
            siblings=["blk.0.attn_k.weight", "blk.0.attn_v.weight"],
        ),
        semantic_links=SemanticLinksDimension(
            tags=["attention", "layer_0", "mistral", "query"],
        ),
        spatial_geometry=SpatialGeometryDimension(
            fibpi3d_coords=coords,
            lattice_index=0,
            dimensional_depth=26.6,
            basis_seed="crystal",
        ),
        energy_tokens=EnergyTokensDimension(
            encode_cost_ms=150.0,
            decode_cost_ms=50.0,
            storage_cost_bytes=16777216,  # Compressed size
            compute_tokens=1000,
        ),
    )


if __name__ == "__main__":
    # Demo
    print("=" * 60)
    print("Helix 14D Object Envelope Demo")
    print("=" * 60)

    obj = example_tensor_envelope()
    obj.validate()

    print("\nExample tensor envelope:")
    print(obj.to_json())

    print("\n" + "=" * 60)
    print("Dimensions present:")
    print("=" * 60)
    print("D1  DATA:           payload_type =", obj.data.payload_type)
    print("D2  FOLDING:        codec =", obj.folding.codec)
    print("D3  GROWTH:         steps =", obj.growth.decode_steps)
    print("D4  ENTROPY:        sha256 =", obj.entropy.sha256[:16] + "...")
    print("D5  AUDIT:          source =", obj.audit.source)
    print("D6  CROSS_HELIX:    parent =", obj.cross_helix.parent_id)
    print("D7  TIME_DRIFT:     mutations =", obj.time_drift.mutation_count)
    print("D8  ENTROPY_LAYERS: merkle =", obj.entropy_layers.merkle_root)
    print("D9  SEMANTIC:       tags =", obj.semantic_links.tags)
    print("D10 NETWORK:        node =", obj.network_coords.node_id)
    print("D11 ENERGY:         encode_ms =", obj.energy_tokens.encode_cost_ms)
    print("D12 QUANTUM:        seeds =", obj.quantum_hints.regrow_seeds)
    print("D13 SPATIAL:        coords =", obj.spatial_geometry.fibpi3d_coords)
    print("D14 CRYPTO:         access =", obj.crypto_shadows.access_control)
