import os
import sys
import unittest
import json
import tempfile
import threading
from unittest.mock import MagicMock, patch

sys.path.insert(0, os.path.dirname(os.path.abspath(__file__)))
import test_helper

from nodes.queue_control import PersistentQueueManager, PauseQueueManager
import nodes.queue_control as queue_control

# 36 Colors matching web/batch_queue.js
BATCH_COLORS = [
    "#10b981", "#3b82f6", "#8b5cf6", "#ec4899", "#f59e0b", "#06b6d4",
    "#ef4444", "#14b8a6", "#6366f1", "#f97316", "#84cc16", "#d946ef",
    "#0ea5e9", "#a855f7", "#e11d48", "#22c55e", "#eab308", "#64748b",
    "#2dd4bf", "#fb7185", "#38bdf8", "#c084fc", "#f43f5e", "#4ade80",
    "#fbbf24", "#818cf8", "#fb923c", "#a3e635", "#f472b6", "#22d3ee",
    "#e879f9", "#a7f3d0", "#fde047", "#fda4af", "#93c5fd", "#c4b5fd"
]

def calculate_segment_types(ordered_pids, registry):
    """
    Python mirror of the segment calculation logic in web/batch_queue.js
    """
    segments = {}
    for pid in ordered_pids:
        info = registry.get(pid)
        if not info or not info.get("batchId"):
            segments[pid] = None
            continue

        batch_id = info["batchId"]
        occurrences = [p for p in ordered_pids if registry.get(p, {}).get("batchId") == batch_id]

        if len(occurrences) <= 1:
            batch_count = info.get("batchCount", 1)
            if batch_count == 1:
                segments[pid] = "single"
            else:
                item_index = info.get("itemIndex", 0)
                if item_index == 0:
                    segments[pid] = "first"
                elif item_index == (batch_count - 1):
                    segments[pid] = "last"
                else:
                    segments[pid] = "middle"
            continue

        first_pid = occurrences[0]
        last_pid = occurrences[-1]

        if pid == first_pid:
            segments[pid] = "first"
        elif pid == last_pid:
            segments[pid] = "last"
        else:
            segments[pid] = "middle"

    return segments


class TestBatchQueueLogic(unittest.TestCase):
    def test_36_color_cycling(self):
        self.assertEqual(len(BATCH_COLORS), 36)
        
        # Batch 1 -> color index 0
        self.assertEqual((1 - 1) % 36, 0)
        self.assertEqual(BATCH_COLORS[(1 - 1) % 36], "#10b981")

        # Batch 36 -> color index 35
        self.assertEqual((36 - 1) % 36, 35)
        self.assertEqual(BATCH_COLORS[(36 - 1) % 36], "#c4b5fd")

        # Batch 37 -> cycles back to color index 0
        self.assertEqual((37 - 1) % 36, 0)
        self.assertEqual(BATCH_COLORS[(37 - 1) % 36], "#10b981")

        # Batch 73 -> cycles back to color index 0
        self.assertEqual((73 - 1) % 36, 0)

    def test_single_item_batch(self):
        registry = {
            "p1": {"batchId": "b1", "batchCount": 1, "itemIndex": 0}
        }
        segments = calculate_segment_types(["p1"], registry)
        self.assertEqual(segments["p1"], "single")

    def test_intact_contiguous_batch(self):
        registry = {
            "p1": {"batchId": "b1", "batchCount": 3, "itemIndex": 0},
            "p2": {"batchId": "b1", "batchCount": 3, "itemIndex": 1},
            "p3": {"batchId": "b1", "batchCount": 3, "itemIndex": 2},
        }
        segments = calculate_segment_types(["p1", "p2", "p3"], registry)
        self.assertEqual(segments["p1"], "first")
        self.assertEqual(segments["p2"], "middle")
        self.assertEqual(segments["p3"], "last")

    def test_interrupted_batch_inserted_in_between(self):
        # Batch A was [pA1, pA2], but prompt B was inserted in-between: [pA1, pB, pA2]
        registry = {
            "pA1": {"batchId": "bA", "batchCount": 2, "itemIndex": 0},
            "pB":  {"batchId": "bB", "batchCount": 1, "itemIndex": 0},
            "pA2": {"batchId": "bA", "batchCount": 2, "itemIndex": 1},
        }
        segments = calculate_segment_types(["pA1", "pB", "pA2"], registry)
        # pA1 is the start of Batch A (curves top, open bottom)
        self.assertEqual(segments["pA1"], "first")
        # pB is a standalone single item (curves top and bottom)
        self.assertEqual(segments["pB"], "single")
        # pA2 is the end of Batch A (open top, curves bottom)
        self.assertEqual(segments["pA2"], "last")

    def test_multiple_interruptions(self):
        # [pA1, pB, pA2, pC, pA3]
        registry = {
            "pA1": {"batchId": "bA", "batchCount": 3, "itemIndex": 0},
            "pB":  {"batchId": "bB", "batchCount": 1, "itemIndex": 0},
            "pA2": {"batchId": "bA", "batchCount": 3, "itemIndex": 1},
            "pC":  {"batchId": "bC", "batchCount": 1, "itemIndex": 0},
            "pA3": {"batchId": "bA", "batchCount": 3, "itemIndex": 2},
        }
        segments = calculate_segment_types(["pA1", "pB", "pA2", "pC", "pA3"], registry)
        self.assertEqual(segments["pA1"], "first")
        self.assertEqual(segments["pB"], "single")
        self.assertEqual(segments["pA2"], "middle")
        self.assertEqual(segments["pC"], "single")
        self.assertEqual(segments["pA3"], "last")


class TestPersistentQueueBatchMetadata(unittest.TestCase):
    def setUp(self):
        self.temp_dir = tempfile.TemporaryDirectory()
        self.temp_file = os.path.join(self.temp_dir.name, "test_persistent_queue.json")
        self.patcher = patch("nodes.queue_control.PERSISTENT_FILE", self.temp_file)
        self.patcher.start()
        self.enabled_patcher = patch("nodes.queue_control.is_persistent_queue_enabled", return_value=True)
        self.enabled_patcher.start()

    def tearDown(self):
        self.enabled_patcher.stop()
        self.patcher.stop()
        self.temp_dir.cleanup()

    def test_save_and_restore_batch_info(self):
        mgr = PersistentQueueManager()
        
        # Set batch info for a prompt
        batch_data = {
            "batchId": "b_12345",
            "batchNumber": 5,
            "batchCount": 3,
            "itemIndex": 0,
            "color": "#10b981",
            "colorIndex": 4
        }
        mgr.set_batch_info("prompt_99", batch_data)
        self.assertIn("prompt_99", mgr.batch_meta)

        # Add item to persistent queue
        item_tuple = (0, "prompt_99", {"prompt": {}}, {}, {})
        mgr.add_item(item_tuple)
        mgr.flush()

        # Verify batch_info is attached inside persistent_items
        saved_entry = [x for x in mgr.persistent_items if x["prompt_id"] == "prompt_99"][0]
        self.assertEqual(saved_entry.get("batch_info"), batch_data)

        # Simulate reload from disk
        new_mgr = PersistentQueueManager()
        self.assertIn("prompt_99", new_mgr.batch_meta)
        self.assertEqual(new_mgr.batch_meta["prompt_99"]["batchId"], "b_12345")

    def test_remove_item_cleans_batch_meta(self):
        mgr = PersistentQueueManager()
        mgr.set_batch_info("p1", {"batchId": "b1"})
        mgr.add_item((0, "p1", {}, {}, {}))
        self.assertIn("p1", mgr.batch_meta)

        mgr.remove_item("p1")
        mgr.flush()
        self.assertNotIn("p1", mgr.batch_meta)
        self.assertEqual(len(mgr.persistent_items), 0)

    def test_wipe_all_cleans_batch_meta(self):
        mgr = PersistentQueueManager()
        mgr.set_batch_info("p1", {"batchId": "b1"})
        mgr.set_batch_info("p2", {"batchId": "b2"})
        mgr.wipe_all()
        self.assertEqual(len(mgr.batch_meta), 0)
        self.assertEqual(len(mgr.persistent_items), 0)


class TestPersistentQueueHardening(unittest.TestCase):
    def setUp(self):
        self.temp_dir = tempfile.TemporaryDirectory()
        self.temp_file = os.path.join(self.temp_dir.name, "test_persistent_queue.json")
        self.patch_file = patch("nodes.queue_control.PERSISTENT_FILE", self.temp_file)
        self.patch_user = patch("nodes.queue_control.USER_DIR", self.temp_dir.name)
        self.patch_enabled = patch("nodes.queue_control.is_persistent_queue_enabled", return_value=True)
        self.patch_file.start()
        self.patch_user.start()
        self.patch_enabled.start()

    def tearDown(self):
        self.patch_enabled.stop()
        self.patch_user.stop()
        self.patch_file.stop()
        self.temp_dir.cleanup()

    def test_finished_prompt_cannot_be_resynced(self):
        mgr = PersistentQueueManager()
        mgr.set_batch_info("p1", {"batchId": "b1"})
        mgr.add_item((0, "p1", {}, {}, {}))
        mgr.remove_item("p1")

        # A late /saturnnodes/batch_queue/sync from another client must not
        # resurrect metadata for a prompt that already finished.
        mgr.set_batch_info("p1", {"batchId": "b1"})
        self.assertNotIn("p1", mgr.batch_meta)

        # Flush so the debounce timer cannot fire after tearDown restores the
        # real PERSISTENT_FILE (it would write into the live user directory).
        mgr.flush()

    def test_batch_meta_is_capped_keeping_recent_entries(self):
        mgr = PersistentQueueManager()
        mgr.MAX_BATCH_META = 3
        with patch("nodes.queue_control.time.time", side_effect=[1, 2, 3, 4, 5]):
            for i in range(5):
                mgr.set_batch_info(f"p{i}", {"batchId": f"b{i}"})

        self.assertEqual(len(mgr.batch_meta), 3)
        self.assertNotIn("p0", mgr.batch_meta)
        self.assertNotIn("p1", mgr.batch_meta)
        self.assertIn("p2", mgr.batch_meta)
        self.assertIn("p4", mgr.batch_meta)

    def test_saves_are_debounced_compact_and_flushable(self):
        mgr = PersistentQueueManager()
        mgr.add_item((0, "p1", {"prompt": {}}, {}, {}))
        # The first write is deliberately deferred; memory is authoritative.
        self.assertTrue(mgr._save_dirty)
        mgr.flush()

        self.assertFalse(mgr._save_dirty)
        self.assertTrue(os.path.exists(self.temp_file))
        with open(self.temp_file, "r", encoding="utf-8") as f:
            raw = f.read()
        # Compact JSON (no indentation/newlines) keeps the file small.
        self.assertNotIn("\n", raw)
        payload = json.loads(raw)
        self.assertEqual(payload["schema_version"], queue_control.PERSISTENT_QUEUE_SCHEMA_VERSION)
        self.assertEqual(payload["items"][0]["prompt_id"], "p1")
        self.assertFalse(os.path.exists(self.temp_file + ".tmp"))

    def test_restore_retry_setting_keeps_then_drops_failed_prompts(self):
        mgr = PersistentQueueManager()
        mgr.persistent_items = [{
            "prompt_id": "p1",
            "item": [0, "p1", {"prompt": {}}, {}],
            "retries": 0,
        }]

        server = MagicMock()
        server.prompt_queue.get_history.return_value = {}
        server.prompt_queue.put.side_effect = RuntimeError("queue unavailable")
        server.number = 0

        def fake_setting(key, default):
            return "1" if key == "PERSISTENT_QUEUE_MAX_RETRIES" else default

        with patch("nodes.queue_control.PromptServer.instance", server), \
             patch("nodes.queue_control.get_env_setting", side_effect=fake_setting):
            mgr.restore_queue("test-client")
            self.assertEqual(len(mgr.persistent_items), 1)
            self.assertEqual(mgr.persistent_items[0]["retries"], 1)

            mgr.has_claimed_once = False
            mgr.restore_queue("test-client")
            self.assertEqual(len(mgr.persistent_items), 0)
            self.assertIn("p1", mgr.recently_finished)

    def test_pause_state_persists_when_enabled(self):
        def fake_setting(key, default):
            if key == "PERSIST_PAUSE_STATE":
                return "true"
            return default

        with patch("nodes.queue_control.get_env_setting", side_effect=fake_setting):
            first = PauseQueueManager()
            first.set_pause(True, "instantly")
            second = PauseQueueManager()

        self.assertTrue(second.paused)
        self.assertEqual(second.mode, "instantly")
        self.assertEqual(second.is_waiting, True)


class TestPauseQueuePatchedGet(unittest.TestCase):
    def test_finite_timeout_returns_instead_of_busy_spinning(self):
        class FakeQueue:
            def __init__(self):
                self.currently_running = {}
                self.mutex = threading.RLock()
                self.calls = 0
                self.last_timeout = object()

            def get(self, timeout=None):
                self.calls += 1
                self.last_timeout = timeout
                return None

        queue = FakeQueue()
        server = MagicMock()
        server.prompt_queue = queue

        def fake_setting(key, default):
            return default

        with patch("nodes.queue_control.get_env_setting", side_effect=fake_setting), \
             patch.object(queue_control.PromptServer, "instance", server):
            manager = PauseQueueManager()
            manager.patch_all()
            # ComfyUI main.py calls get(timeout=0.0) when a GC pass is due.
            # The wrapper must return immediately (not spin forever).
            self.assertIsNone(queue.get(timeout=0.0))
            self.assertEqual(queue.calls, 1)
            self.assertEqual(queue.last_timeout, 0.0)


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