#!/usr/bin/env python3
"""
FloWTorch Bridge - DLPack bridge with Helix64 preconditions

Enforces contiguity before DLPack export and generates receipts for all fixes.

Usage:
    from flowtorch.bridge import to_dlpack_safe, from_tf_dlpack_safe

    # Torch → TF
    dl, metadata = to_dlpack_safe(x)
    # metadata["contiguity_fix_applied"] = True if .contiguous() was called

    # TF → Torch
    torch_tensor, metadata = from_tf_dlpack_safe(tf_tensor)
    # metadata["zero_copy_expected"] = True if dtype/device compatible
"""

import os
import torch as th


def to_dlpack_safe(x: th.Tensor):
    """
    Helix64 precondition: ensure contiguous before DLPack export.

    Environment:
        HELIX64_PRECONDITIONS=1 (default) - Auto-fix non-contiguous tensors
        HELIX64_PRECONDITIONS=0 - Raise error on non-contiguous tensors

    Args:
        x: PyTorch tensor (possibly non-contiguous)

    Returns:
        (dlpack_capsule, metadata_dict)

    Metadata:
        contiguity_fix_applied: bool - True if .contiguous() was called
        original_contiguous: bool - Original contiguity state
        source_device: str - Device (cpu/cuda)
        source_dtype: str - Data type
        source_shape: tuple - Tensor shape
        zero_copy_expected: bool - Whether zero-copy is possible
        reason: str - Reason if not zero-copy

    Raises:
        RuntimeError: If HELIX64_PRECONDITIONS=0 and tensor is non-contiguous
    """
    metadata = {
        "original_contiguous": x.is_contiguous(),
        "source_device": str(x.device),
        "source_dtype": str(x.dtype),
        "source_shape": tuple(x.shape),
        "zero_copy_expected": True,
        "reason": "ok",
    }

    # Check preconditions toggle
    want_guard = os.environ.get("HELIX64_PRECONDITIONS", "1") == "1"

    # Apply contiguity fix if needed
    if not x.is_contiguous():
        if not want_guard:
            raise RuntimeError(
                "Non-contiguous tensor; set HELIX64_PRECONDITIONS=1 or call .contiguous(). "
                f"Tensor shape={x.shape}, stride={x.stride()}"
            )
        x = x.contiguous()
        metadata["contiguity_fix_applied"] = True
        metadata["zero_copy_expected"] = False
        metadata["reason"] = "non_contiguous_source"
    else:
        metadata["contiguity_fix_applied"] = False

    # Export to DLPack
    dl = th.utils.dlpack.to_dlpack(x)

    return dl, metadata


def from_dlpack_safe(dl):
    """
    Import from DLPack with metadata tracking.

    Args:
        dl: DLPack capsule

    Returns:
        (tensor, metadata_dict)
    """
    x = th.utils.dlpack.from_dlpack(dl)

    metadata = {
        "device": str(x.device),
        "dtype": str(x.dtype),
        "shape": tuple(x.shape),
        "contiguous": x.is_contiguous(),
    }

    return x, metadata


def from_tf_dlpack_safe(x):
    """
    TF tensor → Torch tensor via DLPack with Helix64 preconditions.

    Environment:
        HELIX64_PRECONDITIONS=1 (default) - Enforce dtype whitelist + finite check
        HELIX64_PRECONDITIONS=0 - Raise on violations

    Args:
        x: TensorFlow tensor

    Returns:
        (torch_tensor, metadata_dict)

    Metadata:
        source_dtype: str - Original TF dtype
        source_shape: tuple - Tensor shape
        source_device: str - Device location
        finite_check: bool - Whether finite check was performed
        dtype_allowed: bool - Whether dtype is in whitelist
        zero_copy_expected: bool - Whether zero-copy should be possible
        reason: str - Reason if not zero-copy

    Raises:
        RuntimeError: If HELIX64_PRECONDITIONS=0 and preconditions fail
    """
    try:
        import tensorflow as tf
    except ImportError:
        raise RuntimeError("TensorFlow not installed; cannot use from_tf_dlpack_safe")

    # Dtype whitelist (common cross-framework types)
    ALLOWED_DTYPES = {
        tf.bool,
        tf.int32, tf.int64,
        tf.float16, tf.bfloat16, tf.float32, tf.float64,
    }

    # Detect device location
    device_name = x.device if hasattr(x, 'device') else "unknown"
    is_gpu_tensor = "GPU" in device_name or "gpu" in device_name

    metadata = {
        "source_dtype": str(x.dtype),
        "source_shape": tuple(x.shape.as_list()),
        "source_device": device_name,
        "finite_check": False,
        "dtype_allowed": x.dtype in ALLOWED_DTYPES,
        "zero_copy_expected": True,
        "reason": "ok",
    }

    # Check preconditions toggle
    want_guard = os.environ.get("HELIX64_PRECONDITIONS", "1") == "1"

    # Check device compatibility (GPU tensor requires CUDA)
    if is_gpu_tensor and not th.cuda.is_available():
        metadata["zero_copy_expected"] = False
        metadata["reason"] = "device_mismatch_gpu_tensor_no_cuda"
        if not want_guard:
            raise RuntimeError(
                f"TF tensor on GPU but PyTorch CUDA unavailable. "
                f"TF device: {device_name}, torch.cuda.is_available()={th.cuda.is_available()}"
            )

    # Check dtype whitelist
    if x.dtype not in ALLOWED_DTYPES:
        metadata["zero_copy_expected"] = False
        metadata["reason"] = f"dtype_not_whitelisted_{x.dtype.name}"
        if not want_guard:
            raise RuntimeError(
                f"TF dtype {x.dtype} not in whitelist. "
                f"Allowed: {[d.name for d in ALLOWED_DTYPES]}"
            )

    # Finite check for floating point types
    if x.dtype.is_floating:
        metadata["finite_check"] = True
        if not tf.reduce_all(tf.math.is_finite(x)).numpy():
            metadata["zero_copy_expected"] = False
            metadata["reason"] = "non_finite_values"
            if not want_guard:
                raise RuntimeError("TF tensor contains NaN or Inf values")

    # Convert via DLPack
    try:
        dl = tf.experimental.dlpack.to_dlpack(x)
        torch_tensor = th.utils.dlpack.from_dlpack(dl)
    except Exception as e:
        metadata["zero_copy_expected"] = False
        metadata["reason"] = f"dlpack_conversion_failed: {e}"
        raise RuntimeError(f"DLPack conversion failed: {e}") from e

    # Verify result
    metadata["result_dtype"] = str(torch_tensor.dtype)
    metadata["result_shape"] = tuple(torch_tensor.shape)
    metadata["result_device"] = str(torch_tensor.device)
    metadata["result_contiguous"] = torch_tensor.is_contiguous()

    return torch_tensor, metadata
