import io
import json
from pathlib import Path
from types import SimpleNamespace
import struct
import tempfile
import time
import unittest
from unittest.mock import patch

from fastapi.testclient import TestClient
import ml_dtypes
import numpy as np

from tensor_pruning import prune_queries
from tensor_pruning.checkpoint import _queries, prune_checkpoint
from tensor_pruning.service import create_app


def _name(layer):
    return f"model.layers.{layer}.self_attn.q_proj.weight"


def _shard(path, entries, *, sparse=False):
    header, offset = {}, 0
    for name, dtype, shape, length, _ in entries:
        header[name] = dict(dtype=dtype, shape=shape, data_offsets=[offset, offset + length])
        offset += length
    raw = json.dumps(header).encode()
    raw += b" " * (-len(raw) % 8)
    with path.open("wb") as f:
        f.write(struct.pack("<Q", len(raw)))
        f.write(raw)
        if sparse:
            f.truncate(8 + len(raw) + offset)
        else:
            for _, _, _, _, payload in entries:
                f.write(payload)


class CheckpointTests(unittest.TestCase):
    def setUp(self):
        self.temp = tempfile.TemporaryDirectory()
        self.addCleanup(self.temp.cleanup)
        self.root = Path(self.temp.name)
        self.inputs = self.root / "inputs"
        self.models = self.root / "models"
        self.inputs.mkdir()
        self.models.mkdir()
        self.jobs = self.root / "jobs"

    def fixture(self):
        dims = dict(num_hidden_layers=2, num_attention_heads=2, hidden_size=4)
        (self.models / "config.json").write_text(json.dumps(dims))
        arrays = [np.arange(16, dtype=np.float16).reshape(2, 2, 4) - 4,
                  np.asarray(np.arange(16) / 3, dtype=ml_dtypes.bfloat16).reshape(2, 2, 4)]
        mapping = {}
        moments = [np.array([1, 2, 3, 4]), np.array([4, 3, 2, 1])]
        moment_names = {}
        for i, (array, dtype) in enumerate(zip(arrays, ("F16", "BF16"))):
            name = _name(i)
            filename = f"shard-{1-i}.safetensors"
            mapping[name] = filename
            noise = ("unused", "F16", [2], 4, b"\xff" * 4)
            query = (name, dtype, [4, 4], array.nbytes, array.tobytes())
            _shard(self.models / filename, [noise, query])
            moment_names[name] = f"moments-{i}.npy"
            np.save(self.inputs / moment_names[name], moments[i])
        (self.models / "model.safetensors.index.json").write_text(json.dumps({"weight_map": mapping}))
        (self.inputs / "moments.json").write_text(json.dumps(moment_names))
        return arrays, moments

    def test_service_selects_one_global_mask_from_fp16_and_bf16(self):
        arrays, moments = self.fixture()
        expected = prune_queries(arrays, 11, input_second_moments=moments, seed=47)
        app = create_app(self.inputs, self.jobs, checkpoint_root=self.models,
                         token="example", require_token=True)
        headers = {"X-Tensor-Pruning-Token": "example"}
        body = {"keep": 11, "seed": 47, "moments_manifest": "moments.json"}
        with TestClient(app) as client:
            self.assertEqual(client.post("/v1/checkpoint-jobs", json=body).status_code, 401)
            self.assertEqual(client.post("/v1/checkpoint-jobs", headers=headers,
                                         json={**body, "checkpoint": "../"}).status_code, 400)
            created = client.post("/v1/checkpoint-jobs", headers=headers, json=body)
            self.assertEqual(created.status_code, 202, created.text)
            url = "/v1/jobs/" + created.json()["id"]
            for _ in range(500):
                status = client.get(url, headers=headers).json()
                if status["state"] in {"completed", "failed"}:
                    break
                time.sleep(.01)
            self.assertEqual(status["state"], "completed", status)
            self.assertEqual(status["result"]["kept_slots"], 11)
            self.assertEqual([p["name"] for p in status["result"]["projections"]],
                             [_name(0), _name(1)])
            for i, want in enumerate(expected):
                got = client.get(f"{url}/masks/{i}", headers=headers)
                self.assertEqual(got.status_code, 200, got.text if got.status_code != 200 else "")
                np.testing.assert_array_equal(np.load(io.BytesIO(got.content)), want)
        self.assertEqual(list(self.jobs.glob("*/query-staging-*")), [])

    def test_job_submission_fails_closed_without_token(self):
        app = create_app(self.inputs, self.jobs, checkpoint_root=self.models, require_token=True)
        with TestClient(app) as client:
            self.assertEqual(client.get("/health").status_code, 200)
            self.assertEqual(client.post("/v1/checkpoint-jobs", json={"keep": 1}).status_code, 503)

    def test_disk_preflight_rejects_before_staging(self):
        self.fixture()
        usage = SimpleNamespace(total=100, used=100, free=0)
        with patch("tensor_pruning.checkpoint.shutil.disk_usage", return_value=usage):
            with self.assertRaisesRegex(OSError, "Insufficient local disk"):
                prune_checkpoint(self.models, self.root / "out", keep=11,
                                 workspace=self.root, max_weights=100)
        self.assertEqual(list(self.root.glob("query-staging-*")), [])

    def test_llama_70b_layout_exceeds_32_bit_global_indices(self):
        shape = [8192, 8192]
        byte_length = 2 * shape[0] * shape[1]
        config = dict(num_hidden_layers=80, num_attention_heads=64, hidden_size=8192)
        (self.models / "config.json").write_text(json.dumps(config))
        entries = [(_name(i), "BF16", shape, byte_length, None) for i in range(80)]
        _shard(self.models / "model.safetensors", entries, sparse=True)
        query, total = _queries(self.models, 6_000_000_000)
        self.assertEqual(total, 5_368_709_120)
        self.assertGreater(total, 1 << 32)
        self.assertEqual(query[-1].shape, (64, 128, 8192))
        self.assertGreater(query[-1].offset, 1 << 33)
        with self.assertRaisesRegex(ValueError, "eligible weight limit"):
            _queries(self.models, 1 << 32)