"""Tests para safetensors_inspect — parseo de header sin dependencias externas."""

import json
import os
import struct
import sys
import tempfile

sys.path.insert(0, "python/functions/ml")

import pytest

from safetensors_inspect import safetensors_inspect


def _write_safetensors(path: str, header: dict, data: bytes = b"") -> None:
    """Escribe un archivo safetensors mínimo siguiendo la spec oficial.

    Spec: https://github.com/huggingface/safetensors#format
      - 8 bytes: uint64 little-endian con la longitud N del header JSON
      - N bytes: JSON del header
      - (opcional) bytes de datos de tensores
    """
    header_bytes = json.dumps(header, separators=(",", ":")).encode("utf-8")
    header_len = len(header_bytes)
    with open(path, "wb") as f:
        f.write(struct.pack("<Q", header_len))  # uint64 LE
        f.write(header_bytes)
        f.write(data)


def _make_minimal_header(n_tensors: int = 2) -> dict:
    """Genera un header con n_tensors tensores sintéticos."""
    header = {
        "__metadata__": {"format": "pt", "creator": "test"},
    }
    for i in range(n_tensors):
        header[f"tensor_{i}"] = {
            "dtype": "F32",
            "shape": [4, 4],
            "data_offsets": [i * 64, (i + 1) * 64],
        }
    return header


def test_n_tensors_correcto(tmp_path):
    """n_tensors refleja el numero de tensores en el header"""
    path = str(tmp_path / "model.safetensors")
    _write_safetensors(path, _make_minimal_header(n_tensors=3))

    result = safetensors_inspect(path)
    assert result["n_tensors"] == 3


def test_total_size_bytes_correcto(tmp_path):
    """total_size_bytes refleja el tamaño real del archivo"""
    path = str(tmp_path / "model.safetensors")
    data = b"\x00" * 128  # 128 bytes de datos de tensor
    _write_safetensors(path, _make_minimal_header(2), data=data)

    file_size = os.path.getsize(path)
    result = safetensors_inspect(path)
    assert result["total_size_bytes"] == file_size


def test_metadata_campo_dunder_presente(tmp_path):
    """metadata devuelve el contenido de __metadata__"""
    path = str(tmp_path / "model.safetensors")
    header = {
        "__metadata__": {"format": "pt", "model_name": "test_model"},
        "weight": {"dtype": "BF16", "shape": [8], "data_offsets": [0, 16]},
    }
    _write_safetensors(path, header)

    result = safetensors_inspect(path)
    assert result["metadata"] == {"format": "pt", "model_name": "test_model"}


def test_tensors_lista_correcta(tmp_path):
    """tensors es lista con una entrada por tensor del header"""
    path = str(tmp_path / "model.safetensors")
    header = {
        "__metadata__": {},
        "embed.weight": {"dtype": "F16", "shape": [128, 64], "data_offsets": [0, 16384]},
        "proj.bias": {"dtype": "F32", "shape": [64], "data_offsets": [16384, 16640]},
    }
    _write_safetensors(path, header)

    result = safetensors_inspect(path)
    assert result["n_tensors"] == 2
    names = {t["name"] for t in result["tensors"]}
    assert "embed.weight" in names
    assert "proj.bias" in names


def test_tensor_campos_dtype_shape_offsets(tmp_path):
    """cada tensor tiene dtype, shape y data_offsets"""
    path = str(tmp_path / "model.safetensors")
    header = {
        "__metadata__": {},
        "my_tensor": {"dtype": "I32", "shape": [2, 3], "data_offsets": [0, 24]},
    }
    _write_safetensors(path, header)

    result = safetensors_inspect(path)
    t = result["tensors"][0]
    assert t["dtype"] == "I32"
    assert t["shape"] == [2, 3]
    assert t["data_offsets"] == [0, 24]


def test_path_absoluto_en_resultado(tmp_path):
    """result['path'] es la ruta absoluta del archivo"""
    path = str(tmp_path / "model.safetensors")
    _write_safetensors(path, _make_minimal_header(1))

    result = safetensors_inspect(path)
    assert os.path.isabs(result["path"])
    assert result["path"].endswith("model.safetensors")


def test_archivo_no_encontrado_levanta_file_not_found(tmp_path):
    """FileNotFoundError si el archivo no existe"""
    with pytest.raises(FileNotFoundError):
        safetensors_inspect(str(tmp_path / "nonexistent.safetensors"))


def test_header_invalido_levanta_value_error(tmp_path):
    """ValueError si el header no es JSON válido"""
    path = str(tmp_path / "bad.safetensors")
    with open(path, "wb") as f:
        bad_header = b"NOT JSON!!"
        f.write(struct.pack("<Q", len(bad_header)))
        f.write(bad_header)

    with pytest.raises(ValueError):
        safetensors_inspect(path)


def test_archivo_vacio_levanta_value_error(tmp_path):
    """ValueError si el archivo está vacío (< 8 bytes)"""
    path = str(tmp_path / "empty.safetensors")
    with open(path, "wb") as f:
        f.write(b"")

    with pytest.raises(ValueError):
        safetensors_inspect(path)


def test_sin_metadata_dunder(tmp_path):
    """si no hay __metadata__ en el header, metadata retorna dict vacio"""
    path = str(tmp_path / "no_meta.safetensors")
    header = {
        "weight": {"dtype": "F32", "shape": [4], "data_offsets": [0, 16]},
    }
    _write_safetensors(path, header)

    result = safetensors_inspect(path)
    assert result["metadata"] == {}
    assert result["n_tensors"] == 1
