import json
import os
import tempfile
import unittest
from unittest.mock import patch

from core import worker_utils
from core.metadata_audit import (
    _empty_audit_counts,
    audit_metadata_sizes,
)
from core.metadata_utils import merge_counted_payload
from core.path_utils import get_model_resolver_sidecar_path


class MetadataSizeAuditTests(unittest.TestCase):
    def test_merge_counted_payload_sums_audit_numbers_and_preserves_list_order(self):
        target = _empty_audit_counts()
        merge_counted_payload(
            target,
            {
                "scanned_models": 2,
                "checked_metadata": 1,
                "errors": [{"message": "first"}],
                "mismatches": [{"filename": "first.safetensors"}],
            },
            numeric_keys=(
                "scanned_models",
                "metadata_files",
                "checked_metadata",
                "missing_metadata",
                "missing_size",
                "invalid_metadata",
                "skipped_directories",
                "skipped_non_model_files",
            ),
            list_keys=("errors", "mismatches"),
        )
        merge_counted_payload(
            target,
            {
                "scanned_models": 3,
                "checked_metadata": 2,
                "errors": [{"message": "second"}],
                "mismatches": [{"filename": "second.safetensors"}],
            },
            numeric_keys=(
                "scanned_models",
                "metadata_files",
                "checked_metadata",
                "missing_metadata",
                "missing_size",
                "invalid_metadata",
                "skipped_directories",
                "skipped_non_model_files",
            ),
            list_keys=("errors", "mismatches"),
        )

        self.assertEqual(5, target["scanned_models"])
        self.assertEqual(3, target["checked_metadata"])
        self.assertEqual(
            [{"message": "first"}, {"message": "second"}],
            target["errors"],
        )
        self.assertEqual(
            [
                {"filename": "first.safetensors"},
                {"filename": "second.safetensors"},
            ],
            target["mismatches"],
        )

    def _write_model(self, directory, filename, content):
        path = os.path.join(directory, filename)
        with open(path, "wb") as handle:
            handle.write(content)
        return path

    def _write_metadata(self, model_path, payload):
        base_path, _extension = os.path.splitext(model_path)
        metadata_path = f"{base_path}.metadata.json"
        with open(metadata_path, "w", encoding="utf-8") as handle:
            json.dump(payload, handle)
        return metadata_path

    def _model_info(self, model_path, category="checkpoints"):
        return {
            "filename": os.path.basename(model_path),
            "relative_path": os.path.basename(model_path),
            "path": model_path,
            "category": category,
            "base_directory": os.path.dirname(model_path),
        }

    def test_worker_count_preserves_audit_defaults_and_empty_input_behavior(self):
        with patch.object(worker_utils.os, "cpu_count", return_value=8):
            self.assertEqual(
                worker_utils.resolve_worker_count(
                    100,
                    default_worker_multiplier=4,
                    default_worker_limit=64,
                    cpu_count_fallback=4,
                    empty_total_workers=1,
                ),
                (32, 8),
            )
            self.assertEqual(
                worker_utils.resolve_worker_count(
                    20,
                    6,
                    default_worker_multiplier=4,
                    default_worker_limit=64,
                    cpu_count_fallback=4,
                    empty_total_workers=1,
                ),
                (6, 8),
            )
            self.assertEqual(
                worker_utils.resolve_worker_count(
                    0,
                    default_worker_multiplier=4,
                    default_worker_limit=64,
                    cpu_count_fallback=4,
                    empty_total_workers=1,
                ),
                (1, 8),
            )
            self.assertEqual(
                worker_utils.resolve_worker_count(
                    0,
                    6,
                    default_worker_multiplier=4,
                    default_worker_limit=64,
                    cpu_count_fallback=4,
                    empty_total_workers=1,
                ),
                (1, 8),
            )

    def test_reports_top_level_size_mismatch(self):
        with tempfile.TemporaryDirectory() as tmpdir:
            model_path = self._write_model(tmpdir, "mismatch.safetensors", b"12345")
            metadata_path = self._write_metadata(model_path, {"size": 3})

            result = audit_metadata_sizes(models=[self._model_info(model_path)])

        self.assertTrue(result["success"])
        self.assertEqual(1, result["scanned_models"])
        self.assertEqual(1, result["checked_metadata"])
        self.assertEqual(1, result["mismatch_count"])
        mismatch = result["mismatches"][0]
        self.assertEqual(metadata_path, mismatch["metadata_path"])
        self.assertEqual(3, mismatch["metadata_size"])
        self.assertEqual(5, mismatch["actual_size"])
        self.assertEqual(2, mismatch["difference"])
        self.assertEqual("size", mismatch["size_field"])

    def test_ignores_matching_size_and_counts_missing_size(self):
        with tempfile.TemporaryDirectory() as tmpdir:
            matching_path = self._write_model(tmpdir, "matching.safetensors", b"abc")
            missing_size_path = self._write_model(tmpdir, "missing.safetensors", b"abcd")
            self._write_metadata(matching_path, {"size": 3})
            self._write_metadata(missing_size_path, {"model_name": "Missing Size"})

            result = audit_metadata_sizes(
                models=[
                    self._model_info(matching_path),
                    self._model_info(missing_size_path),
                ]
            )

        self.assertEqual(2, result["scanned_models"])
        self.assertEqual(1, result["checked_metadata"])
        self.assertEqual(1, result["missing_size"])
        self.assertEqual(0, result["mismatch_count"])

    def test_uses_matching_nested_file_size_when_top_level_size_is_missing(self):
        with tempfile.TemporaryDirectory() as tmpdir:
            model_path = self._write_model(tmpdir, "nested.safetensors", b"abcdef")
            self._write_metadata(
                model_path,
                {
                    "files": [
                        {"name": "other.safetensors", "size": 6},
                        {"name": "nested.safetensors", "sizeKB": 1},
                    ]
                },
            )

            result = audit_metadata_sizes(models=[self._model_info(model_path)])

        self.assertEqual(1, result["checked_metadata"])
        self.assertEqual(1, result["mismatch_count"])
        mismatch = result["mismatches"][0]
        self.assertEqual(1024, mismatch["metadata_size"])
        self.assertEqual(6, mismatch["actual_size"])
        self.assertEqual("files[].sizeKB", mismatch["size_field"])

    def test_collision_audits_only_each_models_exact_sidecar(self):
        with tempfile.TemporaryDirectory() as tmpdir:
            bin_path = self._write_model(tmpdir, "adapter.bin", b"123")
            safetensors_path = self._write_model(
                tmpdir,
                "adapter.safetensors",
                b"12345",
            )
            with open(f"{bin_path}.metadata.json", "w", encoding="utf-8") as handle:
                json.dump({"filename": "adapter.bin", "size": 3}, handle)
            with open(
                f"{safetensors_path}.metadata.json",
                "w",
                encoding="utf-8",
            ) as handle:
                json.dump(
                    {"filename": "adapter.safetensors", "size": 5},
                    handle,
                )
            self._write_metadata(
                bin_path,
                {"filename": "adapter.bin", "size": 999},
            )

            result = audit_metadata_sizes(
                models=[
                    self._model_info(bin_path),
                    self._model_info(safetensors_path),
                ]
            )

        self.assertTrue(result["success"])
        self.assertEqual(2, result["metadata_files"])
        self.assertEqual(2, result["checked_metadata"])
        self.assertEqual(0, result["mismatch_count"])

    def test_model_resolver_technical_size_overrides_external_size(self):
        with tempfile.TemporaryDirectory() as tmpdir:
            model_path = self._write_model(
                tmpdir,
                "layered.safetensors",
                b"12345",
            )
            external_path = self._write_metadata(
                model_path,
                {
                    "size": 999,
                    "notes": "external user note",
                },
            )
            resolver_path = get_model_resolver_sidecar_path(model_path)
            with open(resolver_path, "w", encoding="utf-8") as handle:
                json.dump(
                    {
                        "managed_by": "comfyui-model-resolver",
                        "size": 4,
                    },
                    handle,
                )

            result = audit_metadata_sizes(models=[self._model_info(model_path)])
            with open(external_path, "r", encoding="utf-8") as handle:
                external_after = json.load(handle)

        self.assertEqual(
            {"size": 999, "notes": "external user note"},
            external_after,
        )
        self.assertEqual(1, result["checked_metadata"])
        self.assertEqual(1, result["mismatch_count"])
        self.assertEqual(resolver_path, result["mismatches"][0]["metadata_path"])
        self.assertEqual(4, result["mismatches"][0]["metadata_size"])

    def test_skips_metadata_files_returned_by_scanner(self):
        with tempfile.TemporaryDirectory() as tmpdir:
            metadata_path = os.path.join(tmpdir, "orphan.metadata.json")
            with open(metadata_path, "w", encoding="utf-8") as handle:
                json.dump({"size": 3}, handle)

            result = audit_metadata_sizes(models=[self._model_info(metadata_path)])

        self.assertEqual(0, result["scanned_models"])
        self.assertEqual(1, result["skipped_non_model_files"])
        self.assertEqual(0, result["mismatch_count"])

    def test_parallel_batches_process_all_models(self):
        with tempfile.TemporaryDirectory() as tmpdir:
            models = []
            for index in range(5):
                model_path = self._write_model(
                    tmpdir,
                    f"model_{index}.safetensors",
                    b"x" * (index + 1),
                )
                stored_size = index + 1
                if index == 3:
                    stored_size = 999
                self._write_metadata(model_path, {"size": stored_size})
                models.append(self._model_info(model_path))

            result = audit_metadata_sizes(
                models=models,
                worker_count=2,
                batch_size=2,
            )

        self.assertEqual(2, result["worker_count"])
        self.assertEqual(2, result["batch_size"])
        self.assertEqual(3, result["batch_count"])
        self.assertEqual(5, result["scanned_models"])
        self.assertEqual(5, result["checked_metadata"])
        self.assertEqual(1, result["mismatch_count"])
        self.assertEqual("model_3.safetensors", result["mismatches"][0]["filename"])


if __name__ == "__main__":
    unittest.main()