#!/usr/bin/env python3
"""
Safetensors Loader - Secure Weight Loading
No pickles, no code execution, memory-mapped reads.

Usage:
    from edge_app.safetensor_loader import load_safetensors, validate_safetensors

    # Full load
    state = load_safetensors("model.safetensors", device="cpu")
    model.load_state_dict(state, strict=False)

    # Streaming / partial load
    with stream_safetensors("model.safetensors") as f:
        for key in f.keys():
            tensor = f.get_tensor(key)

    # Validation
    report = validate_safetensors("model.safetensors")
    assert report["valid"], report["error"]
"""

import os
import hashlib
from pathlib import Path
from typing import Dict, Any, Optional
from contextlib import contextmanager

try:
    from safetensors import safe_open
    from safetensors.torch import load_file, save_file
    SAFETENSORS_AVAILABLE = True
except ImportError:
    SAFETENSORS_AVAILABLE = False


def check_available():
    """Ensure safetensors is installed."""
    if not SAFETENSORS_AVAILABLE:
        raise ImportError(
            "safetensors not installed. Install: pip install safetensors>=0.4.0"
        )


def load_safetensors(
    path: str,
    device: str = "cpu",
    dtype: Optional[str] = None
) -> Dict[str, Any]:
    """
    Load safetensors file into memory.

    Args:
        path: Path to .safetensors file
        device: Device to load tensors to (cpu, cuda, cuda:0, etc)
        dtype: Optional dtype override (e.g., "float16")

    Returns:
        Dict mapping tensor names to tensors

    Raises:
        FileNotFoundError: If file doesn't exist
        ValueError: If file is not safetensors format
    """
    check_available()

    if not os.path.exists(path):
        raise FileNotFoundError(f"Safetensors file not found: {path}")

    if not path.endswith('.safetensors'):
        raise ValueError(f"File must be .safetensors format: {path}")

    return load_file(path, device=device)


@contextmanager
def stream_safetensors(path: str, framework: str = "pt", device: str = "cpu"):
    """
    Stream safetensors file (memory-mapped, partial reads).

    Usage:
        with stream_safetensors("model.safetensors") as f:
            keys = f.keys()
            tensor = f.get_tensor("layer.0.weight")

    Args:
        path: Path to .safetensors file
        framework: Framework ("pt", "np", "tf", "flax")
        device: Device for tensors

    Yields:
        SafetensorFile handle for streaming reads
    """
    check_available()

    if not os.path.exists(path):
        raise FileNotFoundError(f"Safetensors file not found: {path}")

    with safe_open(path, framework=framework, device=device) as f:
        yield f


def save_safetensors(state_dict: Dict[str, Any], path: str):
    """
    Save state dict to safetensors file.

    Args:
        state_dict: Dict of tensor name → tensor
        path: Output path (.safetensors)

    Raises:
        ValueError: If tensors are not contiguous
    """
    check_available()

    if not path.endswith('.safetensors'):
        raise ValueError(f"Output must be .safetensors: {path}")

    # Ensure all tensors are contiguous
    state_dict_contiguous = {
        k: v.contiguous() if hasattr(v, 'contiguous') else v
        for k, v in state_dict.items()
    }

    save_file(state_dict_contiguous, path)


def validate_safetensors(path: str) -> Dict[str, Any]:
    """
    Validate safetensors file and extract metadata.

    Args:
        path: Path to .safetensors file

    Returns:
        Validation report with:
        - valid: bool
        - error: Optional error message
        - sha256: File hash
        - num_tensors: Tensor count
        - dtypes: List of unique dtypes
        - total_params: Total parameter count
        - size_bytes: File size
    """
    report = {
        "valid": False,
        "error": None,
        "sha256": None,
        "num_tensors": 0,
        "dtypes": [],
        "total_params": 0,
        "size_bytes": 0
    }

    try:
        check_available()

        if not os.path.exists(path):
            report["error"] = f"File not found: {path}"
            return report

        if not path.endswith('.safetensors'):
            report["error"] = f"Not .safetensors format: {path}"
            return report

        # Compute SHA256
        sha256 = hashlib.sha256()
        with open(path, 'rb') as f:
            while chunk := f.read(8192):
                sha256.update(chunk)
        report["sha256"] = sha256.hexdigest()

        # Get file size
        report["size_bytes"] = os.path.getsize(path)

        # Stream tensors for metadata
        dtypes_set = set()
        num_tensors = 0
        total_params = 0

        with stream_safetensors(path, framework="pt", device="cpu") as f:
            keys = f.keys()
            num_tensors = len(keys)

            for key in keys:
                tensor = f.get_tensor(key)
                dtypes_set.add(str(tensor.dtype))
                total_params += tensor.numel()

        report["num_tensors"] = num_tensors
        report["dtypes"] = sorted(list(dtypes_set))
        report["total_params"] = total_params
        report["valid"] = True

    except Exception as e:
        report["error"] = str(e)

    return report


def reject_pickle_weights(path: str):
    """
    Reject pickle-based weight files (.pt, .bin, .pth).

    Args:
        path: Path to check

    Raises:
        ValueError: If file is pickle format
    """
    dangerous_extensions = ['.pt', '.bin', '.pth', '.pkl', '.pickle']

    if any(path.endswith(ext) for ext in dangerous_extensions):
        raise ValueError(
            f"Pickle-based weights rejected for security: {path}\n"
            f"Use .safetensors format instead.\n"
            f"Convert: python -m safetensors.torch.convert {path}"
        )


# Example CLI converter
if __name__ == '__main__':
    import sys
    import argparse

    parser = argparse.ArgumentParser(description='Safetensors utilities')
    parser.add_argument('command', choices=['validate', 'convert'],
                        help='Command to run')
    parser.add_argument('input', help='Input file')
    parser.add_argument('--output', help='Output file (for convert)')
    parser.add_argument('--device', default='cpu', help='Device for loading')

    args = parser.parse_args()

    if args.command == 'validate':
        report = validate_safetensors(args.input)
        import json
        print(json.dumps(report, indent=2))
        sys.exit(0 if report["valid"] else 1)

    elif args.command == 'convert':
        if not args.output:
            print("Error: --output required for convert")
            sys.exit(1)

        # Load pickle weights (last time!)
        import torch
        state = torch.load(args.input, map_location='cpu')

        # Save as safetensors
        save_safetensors(state, args.output)

        # Validate
        report = validate_safetensors(args.output)
        print(f"✓ Converted: {args.input} → {args.output}")
        print(f"  SHA256: {report['sha256']}")
        print(f"  Tensors: {report['num_tensors']}")
        print(f"  Params: {report['total_params']:,}")
        print(f"  Size: {report['size_bytes']:,} bytes")
