"""
SEE — Symbolic Execution Envelope

Wraps decompressed GX code in a symbolic context envelope that bridges
the XIC virtual machine with the LAIN 8-lane cognition engine.

The envelope serves as an immutable container that carries:
- Decompressed code bytes + manifest
- Glyph context (resonance data, superpowers, specialized types)
- Execution metadata (mode, epoch, invocation chain)
- Integrity hash for verification

Integration points:
  - XIC VM (xic_vm.py): run_xic_program consumes SEE envelopes
  - LAIN runtime (gx_lain/runtime.py): execute_with_lain works within envelopes
  - Symbolic pipeline (glyphos/symbolic_pipeline.py): run_symbolic_pipeline feeds envelopes
  - GSZ3 decompressor (xic_extensions/gsz3_decompressor.py): decompresses payloads
"""

from __future__ import annotations

import hashlib
import json
import time
import uuid
import logging
from dataclasses import dataclass, field
from typing import Any, Dict, List, Optional, Tuple

logger = logging.getLogger(__name__)


@dataclass
class SymbolicExecutionEnvelope:
    """
    Immutable envelope wrapping decompressed code with symbolic cognition context.

    Once constructed via build(), the envelope is read-only — LAIN and the XIC VM
    consume it without mutation. This guarantees deterministic execution.
    """

    code: bytes
    manifest: Dict[str, Any]
    glyph_context: Dict[str, Any]
    glyph_ids: List[str]
    resonance_map: Dict[str, float]
    mode: str
    epoch: Optional[str]
    invocation_id: str
    chain_label: Optional[str]
    integrity_hash: str
    built_at: float
    metadata: Dict[str, Any] = field(default_factory=dict)

    @classmethod
    def build(
        cls,
        code: bytes,
        manifest: Optional[Dict[str, Any]] = None,
        glyph_context: Optional[Dict[str, Any]] = None,
        glyph_ids: Optional[List[str]] = None,
        mode: str = "symbolic",
        epoch: Optional[str] = None,
        chain_label: Optional[str] = None,
        metadata: Optional[Dict[str, Any]] = None,
    ) -> "SymbolicExecutionEnvelope":
        """
        Construct a new envelope from raw components.

        Args:
            code: Decompressed code bytes.
            manifest: Optional GX manifest dict.
            glyph_context: Optional glyph cognition context.
            glyph_ids: Optional list of glyph IDs for multi-glyph resonance.
            mode: Execution mode ("symbolic", "analyze", "execute").
            epoch: Optional epoch identifier for time-aligned execution.
            chain_label: Optional chain label for jump-table routing.
            metadata: Optional extra metadata to embed.

        Returns:
            Fully constructed SymbolicExecutionEnvelope.
        """
        if manifest is None:
            manifest = {}
        if glyph_context is None:
            glyph_context = {}
        if glyph_ids is None:
            glyph_ids = []
        if metadata is None:
            metadata = {}

        glyph_resonance = cls._compute_glyph_resonance_map(glyph_context, glyph_ids)

        payload = {
            "code_len": len(code),
            "manifest_version": manifest.get("version", "unknown"),
            "glyph_ids": glyph_ids,
            "mode": mode,
        }
        integrity_hash = cls._hash_envelope(code, payload)

        return cls(
            code=code,
            manifest=manifest,
            glyph_context=glyph_context,
            glyph_ids=glyph_ids,
            resonance_map=glyph_resonance,
            mode=mode,
            epoch=epoch,
            invocation_id=metadata.get("invocation_id", str(uuid.uuid4())),
            chain_label=chain_label,
            integrity_hash=integrity_hash,
            built_at=time.time(),
            metadata=metadata,
        )

    def verify_integrity(self) -> bool:
        """Verify the envelope's integrity hash matches its contents."""
        payload = {
            "code_len": len(self.code),
            "manifest_version": self.manifest.get("version", "unknown"),
            "glyph_ids": self.glyph_ids,
            "mode": self.mode,
        }
        expected = self._hash_envelope(self.code, payload)
        return expected == self.integrity_hash

    def to_dict(self) -> Dict[str, Any]:
        """Serialize envelope to a dict (for telemetry, logging, transport)."""
        return {
            "code_size": len(self.code),
            "code_preview": self.code[:120].decode("utf-8", errors="replace"),
            "manifest_version": self.manifest.get("version", ""),
            "glyph_ids": self.glyph_ids,
            "glyph_count": len(self.glyph_ids),
            "resonance": self.resonance_map,
            "mode": self.mode,
            "epoch": self.epoch,
            "invocation_id": self.invocation_id,
            "chain_label": self.chain_label,
            "integrity_hash": self.integrity_hash,
            "built_at": self.built_at,
        }

    def resolve_glyph_context(
        self, glyph_id: str
    ) -> Optional[Dict[str, Any]]:
        """Resolve a single glyph's context data from the envelope.

        Args:
            glyph_id: The glyph identifier to look up.

        Returns:
            Glyph context dict or None if not found.
        """
        glyph_data = self.glyph_context.get(glyph_id)
        if glyph_data:
            return {
                "glyph_id": glyph_id,
                "data": glyph_data,
                "resonance_weight": self.resonance_map.get(glyph_id, 0.0),
            }
        raw_glyphs = self.glyph_context.get("glyphs", {})
        glyph_data = raw_glyphs.get(glyph_id)
        if glyph_data:
            return {
                "glyph_id": glyph_id,
                "data": glyph_data,
                "resonance_weight": self.resonance_map.get(glyph_id, 0.0),
            }
        return None

    @staticmethod
    def _compute_glyph_resonance_map(
        glyph_context: Dict[str, Any],
        glyph_ids: List[str],
    ) -> Dict[str, float]:
        """Compute a flat glyph_id → resonance_weight map.

        Extracts weights from glyph_context and supplements with
        even distribution for glyph_ids missing explicit weights.
        """
        resonance: Dict[str, float] = {}

        raw_glyphs: Dict[str, Any] = glyph_context.get("glyphs", {})
        for gid, data in raw_glyphs.items():
            if isinstance(data, dict):
                weight = data.get("resonance_weight") or data.get("weight") or data.get("score", 0)
                resonance[gid] = float(weight)

        for gid in glyph_ids:
            if gid not in resonance:
                direct = glyph_context.get(gid)
                if isinstance(direct, dict):
                    weight = direct.get("resonance_weight") or direct.get("weight") or direct.get("score", 0)
                    resonance[gid] = float(weight)
                else:
                    resonance[gid] = 0.0

        if resonance and not any(v > 0 for v in resonance.values()):
            fallback = 1.0 / max(len(resonance), 1)
            for gid in resonance:
                resonance[gid] = fallback

        return resonance

    @staticmethod
    def _hash_envelope(code: bytes, payload: Dict[str, Any]) -> str:
        """SHA-256 integrity hash covering code + metadata."""
        hasher = hashlib.sha256()
        hasher.update(code)
        hasher.update(json.dumps(payload, sort_keys=True).encode())
        return hasher.hexdigest()[:32]


def wrap_code(
    code_bytes: bytes,
    glyph_ids: Optional[List[str]] = None,
    mode: str = "symbolic",
    manifest: Optional[Dict[str, Any]] = None,
    glyph_context: Optional[Dict[str, Any]] = None,
    chain_label: Optional[str] = None,
) -> SymbolicExecutionEnvelope:
    """Convenience function: wrap raw decompressed code in an envelope.

    Args:
        code_bytes: Decompressed code bytes.
        glyph_ids: Optional glyph IDs for resonance.
        mode: Execution mode.
        manifest: Optional manifest dict.
        glyph_context: Optional glyph cognition context.
        chain_label: Optional chain label.

    Returns:
        SymbolicExecutionEnvelope ready for execution.
    """
    return SymbolicExecutionEnvelope.build(
        code=code_bytes,
        manifest=manifest,
        glyph_context=glyph_context,
        glyph_ids=glyph_ids,
        mode=mode,
        chain_label=chain_label,
    )


def unwrap_envelope(
    envelope: SymbolicExecutionEnvelope,
) -> Tuple[bytes, Dict[str, Any], List[str]]:
    """Extract the core execution components from an envelope.

    Returns (code_bytes, context_dict, glyph_ids).

    The context dict includes mode, epoch, invocation_id, chain_label,
    and the full resonance map for symbolic processing.
    """
    context = {
        "mode": envelope.mode,
        "epoch": envelope.epoch,
        "invocation_id": envelope.invocation_id,
        "chain_label": envelope.chain_label,
        "resonance_map": envelope.resonance_map,
        "manifest": envelope.manifest,
        "glyph_context": envelope.glyph_context,
    }
    return envelope.code, context, envelope.glyph_ids


def execute_with_envelope(
    envelope: SymbolicExecutionEnvelope,
) -> Dict[str, Any]:
    """Execute decompressed code through the full symbolic pipeline within the envelope.

    Pipeline:
      1. Verify envelope integrity
      2. Unwrap code + context
      3. Route through run_symbolic_pipeline with glyph data
      4. Return structured result

    Args:
        envelope: The execution envelope.

    Returns:
        Dict with keys: output_text, fused_symbol, steps, diagnostics.
    """
    if not envelope.verify_integrity():
        return {
            "output_text": "[SEE] Integrity verification failed — envelope tampered",
            "fused_symbol": None,
            "steps": [],
            "diagnostics": {"error": "integrity_check_failed"},
        }

    code, context, glyph_ids = unwrap_envelope(envelope)

    prompt = code.decode("utf-8", errors="replace")

    try:
        from glyphos.symbolic_pipeline import run_symbolic_pipeline

        result = run_symbolic_pipeline(
            prompt=prompt,
            context=context,
            glyph_ids=glyph_ids or None,
        )

        return {
            "output_text": result.output_text,
            "fused_symbol": result.fused_symbol,
            "steps": result.steps,
            "diagnostics": {
                "step_count": len(result.steps),
                "mode": envelope.mode,
                "integrity": "verified",
            },
        }
    except Exception as e:
        logger.exception(f"[SEE] Pipeline execution failed: {e}")
        return {
            "output_text": f"[SEE] Execution error: {e}",
            "fused_symbol": None,
            "steps": [],
            "diagnostics": {"error": str(e)},
        }