#!/usr/bin/env python3
"""
Runtime Weights Policy Enforcement
Guards against unsafe pickle weights at load time + health checks.
"""

import os
from pathlib import Path
from typing import Dict, Any, Optional

# Import safetensor loader
try:
    from edge_app.safetensor_loader import (
        reject_pickle_weights,
        validate_safetensors,
        load_safetensors,
        SAFETENSORS_AVAILABLE
    )
except ImportError:
    SAFETENSORS_AVAILABLE = False
    # Stub functions if import fails
    def reject_pickle_weights(path: str):
        raise ImportError("safetensor_loader not available")
    def validate_safetensors(path: str):
        return {"valid": False, "error": "safetensor_loader not available"}
    def load_safetensors(path: str, device: str = "cpu"):
        raise ImportError("safetensor_loader not available")


# Policy: Reject pickle weights by default
SAFETENSORS_ONLY = os.getenv("SAFETENSORS_ONLY", "1") == "1"

# Allowed dtypes (fail closed)
ALLOWED_DTYPES = {
    "torch.float16",
    "torch.bfloat16",
    "torch.float32",
    "torch.int8",
    "torch.uint8"
}


def load_weights_safe(path: str, device: str = "cpu") -> tuple[Dict[str, Any], Dict[str, Any]]:
    """
    Load weights with policy enforcement.

    Args:
        path: Path to weights file (.safetensors or .hxz)
        device: Device to load to

    Returns:
        (state_dict, metadata)

    Raises:
        ValueError: If weights violate policy
        RuntimeError: If safetensors not available
    """
    path = Path(path)

    # Check if HXZ seed (allowed)
    if path.suffix == '.hxz':
        return load_hxz_seed(path)

    # Enforce safetensors-only policy
    if SAFETENSORS_ONLY:
        reject_pickle_weights(str(path))

    # Validate safetensors
    if not SAFETENSORS_AVAILABLE:
        raise RuntimeError(
            "safetensors not available but SAFETENSORS_ONLY=1. "
            "Install: pip install safetensors>=0.4.0"
        )

    # Validate file
    meta = validate_safetensors(str(path))
    if not meta["valid"]:
        raise ValueError(f"Invalid safetensors file: {meta.get('error')}")

    # Check dtypes (fail closed)
    for dtype in meta["dtypes"]:
        if dtype not in ALLOWED_DTYPES:
            raise ValueError(
                f"Dtype not allowed: {dtype}. "
                f"Allowed: {ALLOWED_DTYPES}"
            )

    # Load
    state = load_safetensors(str(path), device=device)

    return state, meta


def load_hxz_seed(path: Path) -> tuple[Dict[str, Any], Dict[str, Any]]:
    """
    Load HXZ seed (allowed format).

    Args:
        path: Path to .hxz file

    Returns:
        (empty_state_dict, metadata)
    """
    import hashlib

    # Compute SHA256
    sha256_hash = hashlib.sha256()
    with open(path, 'rb') as f:
        while chunk := f.read(8192):
            sha256_hash.update(chunk)

    meta = {
        "format": "hxz-seed",
        "seed_sha256": sha256_hash.hexdigest(),
        "size_bytes": path.stat().st_size,
        "valid": True
    }

    # HXZ seeds don't have state dicts (regenerative)
    return {}, meta


def get_weights_health() -> Dict[str, Any]:
    """
    Get weights health status for /health endpoint.

    Returns:
        Health report dict
    """
    health = {
        "policy": {
            "safetensors_only": SAFETENSORS_ONLY,
            "allowed_dtypes": list(ALLOWED_DTYPES)
        },
        "safetensors_available": SAFETENSORS_AVAILABLE
    }

    # Check FLOWTORCH_SEED_PATH if set
    seed_path = os.getenv("FLOWTORCH_SEED_PATH")
    if seed_path:
        seed_path = Path(seed_path)
        if seed_path.exists():
            try:
                if seed_path.suffix == '.hxz':
                    _, meta = load_hxz_seed(seed_path)
                    health["active_seed"] = {
                        "path": str(seed_path),
                        "format": meta["format"],
                        "sha256": meta["seed_sha256"],
                        "size_bytes": meta["size_bytes"],
                        "status": "ok"
                    }
                elif seed_path.suffix == '.safetensors':
                    if not SAFETENSORS_AVAILABLE:
                        health["active_seed"] = {
                            "path": str(seed_path),
                            "status": "error",
                            "error": "safetensors not available"
                        }
                    else:
                        meta = validate_safetensors(str(seed_path))
                        if meta["valid"]:
                            health["active_seed"] = {
                                "path": str(seed_path),
                                "format": "safetensors",
                                "sha256": meta["sha256"],
                                "num_tensors": meta["num_tensors"],
                                "dtypes": meta["dtypes"],
                                "total_params": meta["total_params"],
                                "size_bytes": meta["size_bytes"],
                                "status": "ok"
                            }
                        else:
                            health["active_seed"] = {
                                "path": str(seed_path),
                                "status": "error",
                                "error": meta.get("error")
                            }
                else:
                    health["active_seed"] = {
                        "path": str(seed_path),
                        "status": "error",
                        "error": f"Unsupported format: {seed_path.suffix}"
                    }
            except Exception as e:
                health["active_seed"] = {
                    "path": str(seed_path),
                    "status": "error",
                    "error": str(e)
                }
        else:
            health["active_seed"] = {
                "path": str(seed_path),
                "status": "missing"
            }

    return health


# Tripwire test
def test_tripwire():
    """
    Test that pickle rejection works.

    Returns:
        True if policy enforcement works
    """
    import tempfile

    # Create a fake .pt file
    with tempfile.NamedTemporaryFile(suffix='.pt', delete=False) as f:
        fake_pt = f.name
        f.write(b'fake')

    try:
        # Should raise ValueError
        reject_pickle_weights(fake_pt)
        return False  # Should not reach here
    except ValueError as e:
        if "Pickle-based weights rejected" in str(e):
            return True
        return False
    finally:
        Path(fake_pt).unlink(missing_ok=True)


if __name__ == '__main__':
    import sys

    # Test tripwire
    if test_tripwire():
        print("✅ Tripwire test passed (pickles rejected)")
    else:
        print("❌ Tripwire test failed")
        sys.exit(1)

    # Test health
    health = get_weights_health()
    import json
    print("\nHealth check:")
    print(json.dumps(health, indent=2))
