"""
BloomOS Kernel - CDNA Handlers
==============================

Syscall handlers for CDNA encode/decode operations.

Syscalls:
- cdna.decode.v1: Decode CDNA shards to tensors
- cdna.encode.v1: Encode tensors to CDNA format
"""

from __future__ import annotations

import hashlib
import json
import time
from pathlib import Path
from typing import Any, Dict, List, Optional, TYPE_CHECKING

import numpy as np

if TYPE_CHECKING:
    from ..syscall_registry import SyscallRegistry


def cdna_decode_v1(inputs: Dict[str, Any]) -> Dict[str, Any]:
    """
    Decode CDNA shards to tensors.

    Wraps SubstrateOps.cdna_decode() from helix_cdc/substrate/executor.py

    Args (inputs dict):
        manifest_path: Path to CDNA manifest JSON
        tensor_names: Optional list of specific tensors to decode
        max_tensors: Optional limit on number of tensors (for testing)

    Returns:
        {
            "status": "OK" | "FAIL",
            "tensors_decoded": int,
            "total_elements": int,
            "bytes_read": int,
            "tensor_hashes": {name: sha256},
            "warnings": [],
        }
    """
    manifest_path = Path(inputs["manifest_path"])
    tensor_names = inputs.get("tensor_names")
    max_tensors = inputs.get("max_tensors", 0)

    if not manifest_path.exists():
        return {
            "status": "FAIL",
            "error": f"Manifest not found: {manifest_path}",
            "tensors_decoded": 0,
            "total_elements": 0,
            "bytes_read": 0,
            "tensor_hashes": {},
            "warnings": [],
        }

    # Import SubstrateOps
    try:
        from helix_cdc.substrate.executor import SubstrateOps, ExecutionContext
    except ImportError as e:
        return {
            "status": "FAIL",
            "error": f"Failed to import SubstrateOps: {e}",
            "tensors_decoded": 0,
            "total_elements": 0,
            "bytes_read": 0,
            "tensor_hashes": {},
            "warnings": [],
        }

    # Load manifest
    manifest = json.loads(manifest_path.read_text())

    # Determine base paths
    project_root = manifest_path.resolve().parent
    while project_root != project_root.parent:
        if (project_root / "helix_cdc").exists():
            break
        project_root = project_root.parent

    shard_base = project_root

    # Create execution context
    ctx = ExecutionContext(
        manifest_path=manifest_path,
        manifest=manifest,
        shard_base=shard_base,
    )

    # Execute decode
    result = SubstrateOps.cdna_decode(
        ctx,
        tensor_names=tensor_names,
        max_tensors=max_tensors,
    )

    # Compute tensor hashes
    tensor_hashes = {}
    for name, tensor in ctx.tensors.items():
        tensor_hashes[name] = hashlib.sha256(
            np.ascontiguousarray(tensor).tobytes()
        ).hexdigest()

    # Estimate bytes read from shards
    bytes_read = 0
    for shard in manifest.get("shards", []):
        shard_rel = shard.get("path", "")
        if Path(shard_rel).is_absolute():
            shard_path = Path(shard_rel)
        else:
            shard_path = shard_base / shard_rel

        if shard_path.exists():
            bytes_read += shard_path.stat().st_size

    return {
        "status": result.status,
        "tensors_decoded": result.tensors_decoded,
        "total_elements": result.elements_processed,
        "bytes_read": bytes_read,
        "tensor_hashes": tensor_hashes,
        "warnings": result.warnings,
        "memory_bytes": result.memory_bytes,
    }


def cdna_decode_v1_cost_estimator(inputs: Dict[str, Any]) -> float:
    """
    Estimate IT cost for cdna.decode.v1.

    Cost model:
    - Base cost: 1.0 IT
    - Per tensor: 0.1 IT
    - I/O cost computed from actual bytes

    Args:
        inputs: Syscall inputs

    Returns:
        Estimated IT tokens
    """
    base_cost = 1.0
    per_tensor_cost = 0.1

    manifest_path = Path(inputs.get("manifest_path", ""))
    if not manifest_path.exists():
        return base_cost

    try:
        manifest = json.loads(manifest_path.read_text())
        shard_count = len(manifest.get("shards", []))
        max_tensors = inputs.get("max_tensors", 0)

        if max_tensors > 0:
            shard_count = min(shard_count, max_tensors)

        return base_cost + (shard_count * per_tensor_cost)
    except Exception:
        return base_cost


def cdna_encode_v1(inputs: Dict[str, Any]) -> Dict[str, Any]:
    """
    Encode tensors to CDNA format.

    Args (inputs dict):
        tensor_data: Tensor as list (will be converted to numpy)
        tensor_name: Name for the tensor
        output_path: Path to write CDNA shard
        n_clusters: Number of k-means clusters (default: 256)

    Returns:
        {
            "status": "OK" | "FAIL",
            "output_path": str,
            "bytes_written": int,
            "input_hash": str,
            "output_hash": str,
            "compression_ratio": float,
        }
    """
    tensor_data = inputs.get("tensor_data")
    tensor_name = inputs.get("tensor_name", "tensor")
    output_path = inputs.get("output_path")
    n_clusters = inputs.get("n_clusters", 256)

    if tensor_data is None or output_path is None:
        return {
            "status": "FAIL",
            "error": "Missing required inputs: tensor_data, output_path",
        }

    try:
        # Import encoder
        from helix_cdc.regrow.cdna_encoder import encode_tensor_to_cdna
    except ImportError as e:
        return {
            "status": "FAIL",
            "error": f"Failed to import cdna_encoder: {e}",
        }

    # Convert to numpy if needed
    if isinstance(tensor_data, list):
        tensor = np.array(tensor_data, dtype=np.float32)
    else:
        tensor = np.asarray(tensor_data, dtype=np.float32)

    # Compute input hash
    input_hash = hashlib.sha256(
        np.ascontiguousarray(tensor).tobytes()
    ).hexdigest()

    # Encode
    try:
        output_path = Path(output_path)
        output_path.parent.mkdir(parents=True, exist_ok=True)

        encode_tensor_to_cdna(
            tensor=tensor,
            output_path=output_path,
            tensor_name=tensor_name,
            n_clusters=n_clusters,
        )

        bytes_written = output_path.stat().st_size
        output_hash = hashlib.sha256(output_path.read_bytes()).hexdigest()

        original_bytes = tensor.nbytes
        compression_ratio = original_bytes / bytes_written if bytes_written > 0 else 0.0

        return {
            "status": "OK",
            "output_path": str(output_path),
            "bytes_written": bytes_written,
            "input_hash": input_hash,
            "output_hash": output_hash,
            "compression_ratio": compression_ratio,
        }

    except Exception as e:
        return {
            "status": "FAIL",
            "error": str(e),
        }


def cdna_encode_v1_cost_estimator(inputs: Dict[str, Any]) -> float:
    """
    Estimate IT cost for cdna.encode.v1.

    Cost model:
    - Base cost: 2.0 IT (encoding is more expensive than decoding)
    - Per MB: 0.1 IT

    Args:
        inputs: Syscall inputs

    Returns:
        Estimated IT tokens
    """
    base_cost = 2.0
    per_mb_cost = 0.1

    tensor_data = inputs.get("tensor_data")
    if tensor_data is None:
        return base_cost

    try:
        if isinstance(tensor_data, list):
            # Estimate size from nested list
            import sys
            size_estimate = sys.getsizeof(tensor_data)
        else:
            size_estimate = np.asarray(tensor_data).nbytes

        mb = size_estimate / (1024 * 1024)
        return base_cost + (mb * per_mb_cost)
    except Exception:
        return base_cost


# Schema definitions
CDNA_DECODE_V1_SCHEMA = {
    "required": ["manifest_path"],
    "types": {
        "manifest_path": "str",
        "tensor_names": "list",
        "max_tensors": "int",
    },
}

CDNA_ENCODE_V1_SCHEMA = {
    "required": ["tensor_data", "output_path"],
    "types": {
        "tensor_data": "list",
        "tensor_name": "str",
        "output_path": "str",
        "n_clusters": "int",
    },
}


def register_cdna_handlers(registry: "SyscallRegistry") -> None:
    """Register CDNA handlers with the registry."""
    registry.register(
        canonical_name="cdna.decode.v1",
        handler=cdna_decode_v1,
        schema=CDNA_DECODE_V1_SCHEMA,
        cost_estimator=cdna_decode_v1_cost_estimator,
        description="Decode CDNA shards to tensors (wraps SubstrateOps.cdna_decode)",
    )

    registry.register(
        canonical_name="cdna.encode.v1",
        handler=cdna_encode_v1,
        schema=CDNA_ENCODE_V1_SCHEMA,
        cost_estimator=cdna_encode_v1_cost_estimator,
        description="Encode tensors to CDNA format",
    )
