# Copyright 2026 The Fuchsia Authors. All rights reserved.
# Use of this source code is governed by a BSD-style license that can be
# found in the LICENSE file.

import os
import sys
import tempfile
import unittest
from io import StringIO
from pathlib import Path
from unittest.mock import MagicMock, patch

worktree_dir = os.path.dirname(os.path.dirname(os.path.abspath(__file__)))
sys.path.insert(0, worktree_dir)

import argparse

from subcommands import add as add_cmd
from subcommands import pool_add as pool_add_cmd
from subcommands import pool_remove as pool_remove_cmd
from subcommands import remove as remove_cmd
from worktree import NoFreeWorktreesError, SyncStatus, Worktree, WorktreeState
from worktree_pool import ADJECTIVES, NOUNS, WorktreePool
from worktree_printer import WorktreePrinter


class TestWorktreePool(unittest.TestCase):
    def setUp(self) -> None:
        self.temp_dir = tempfile.TemporaryDirectory()
        self.fuchsia_dir = Path(self.temp_dir.name)
        self.jiri_root = self.fuchsia_dir / ".jiri_root"
        self.jiri_root.mkdir(parents=True, exist_ok=True)
        self.pool = WorktreePool(fuchsia_dir=str(self.fuchsia_dir))

    def tearDown(self) -> None:
        self.temp_dir.cleanup()

    def test_empty(self) -> None:
        self.assertEqual(self.pool.get_worktrees(), [])

    def test_invalid_state_transitions(self) -> None:
        wt_path = self.jiri_root / "worktrees" / "wt1"
        wt_path.mkdir(parents=True, exist_ok=True)
        with open(self.pool.registry_file, "w") as f:
            f.write(f"{wt_path}\n")

        wt = self.pool.get_worktree_by_name("wt1")
        self.assertEqual(wt.get_state(), WorktreeState.FREE)

        # Cannot release if FREE
        with self.assertRaises(RuntimeError):
            wt.release_lease()

        # Lease it
        wt.acquire_lease(task_id="test")
        self.assertEqual(wt.get_state(), WorktreeState.LEASED)

        # Cannot lease again
        with self.assertRaises(RuntimeError):
            wt.acquire_lease(task_id="test2")

    def test_get_any_free_worktree(self) -> None:
        with self.assertRaises(NoFreeWorktreesError):
            self.pool.get_any_free_worktree()

        wt_path = self.jiri_root / "worktrees" / "wt1"
        wt_path.mkdir(parents=True, exist_ok=True)
        with open(self.pool.registry_file, "w") as f:
            f.write(f"{wt_path}\n")

        wt = self.pool.get_any_free_worktree()
        self.assertEqual(wt.name, "wt1")

    @patch("worktree.run_git")
    def test_release_detaches_head(self, mock_run_git: MagicMock) -> None:
        wt_path = self.jiri_root / "worktrees" / "wt1"
        wt_path.mkdir(parents=True, exist_ok=True)
        with open(self.pool.registry_file, "w") as f:
            f.write(f"{wt_path}\n")
        wt = self.pool.get_worktrees()[0]
        wt.acquire_lease("my-task")
        wt.release_lease()
        mock_run_git.assert_called_once_with(
            wt.path, ["checkout", "--detach"], quiet=True, check=True
        )

    def test_generate_random_pool_name(self) -> None:
        name = self.pool._generate_random_pool_name()
        self.assertIn("-", name)
        adj, noun = name.split("-", 1)
        self.assertIn(adj, ADJECTIVES)
        self.assertIn(noun, NOUNS)


class TestAddSubcommand(unittest.TestCase):
    def setUp(self) -> None:
        self.temp_dir = tempfile.TemporaryDirectory()
        self.fuchsia_dir = Path(self.temp_dir.name)
        self.jiri_root = self.fuchsia_dir / ".jiri_root"
        self.jiri_root.mkdir(parents=True, exist_ok=True)
        self.pool = WorktreePool(fuchsia_dir=str(self.fuchsia_dir))

    def tearDown(self) -> None:
        self.temp_dir.cleanup()

    def test_add_claims_slot(self) -> None:
        wt_path = self.jiri_root / "worktrees" / "wt1"
        wt_path.mkdir(parents=True, exist_ok=True)
        self.pool.registry_file.write_text(f"{wt_path}\n")

        args = argparse.Namespace(
            name="my-feat", pool_name=None, sync=False, json=False
        )
        with patch("sys.stdout", new_callable=StringIO) as mock_out:
            add_cmd.run(args, self.pool)
            self.assertIn(".jiri_root/worktrees/my-feat", mock_out.getvalue())
        wt = self.pool.get_worktrees()[0]
        self.assertEqual(wt.get_state(), WorktreeState.LEASED)
        lease = wt.get_lease_info()
        assert lease is not None
        self.assertEqual(lease.task_id, "my-feat")

    @patch("subcommands.add.run_jiri")
    def test_add_with_sync_order(self, mock_run_jiri: MagicMock) -> None:
        wt_path = self.jiri_root / "worktrees" / "wt1"
        wt_path.mkdir(parents=True, exist_ok=True)
        self.pool.registry_file.write_text(f"{wt_path}\n")

        args = argparse.Namespace(
            name="my-feat", pool_name=None, sync=True, json=False
        )
        manager = MagicMock()
        manager.attach_mock(mock_run_jiri, "mock_run_jiri")

        with patch("sys.stdout", new_callable=StringIO):
            add_cmd.run(args, self.pool)

        expected_calls = [
            unittest.mock.call.mock_run_jiri(
                self.jiri_root,
                ["worktree", "sync", str(wt_path)],
                check=True,
            ),
        ]
        self.assertEqual(manager.mock_calls, expected_calls)

    @patch("worktree.run_git")
    def test_remove_by_task_id(self, mock_run_git: MagicMock) -> None:
        wt_path = self.jiri_root / "worktrees" / "wt1"
        wt_path.mkdir(parents=True, exist_ok=True)
        self.pool.registry_file.write_text(f"{wt_path}\n")
        wt = self.pool.get_worktrees()[0]
        wt.acquire_lease("my-task-123")

        args = argparse.Namespace(name="my-task-123")
        remove_cmd.run(args, self.pool)
        self.assertEqual(wt.get_state(), WorktreeState.FREE)

        args_pool = argparse.Namespace(name="my-task-123", force=False)
        with self.assertRaises(KeyError):
            pool_remove_cmd.run(args_pool, self.pool)

    @patch("worktree_pool.run_jiri")
    @patch("sys.stderr", new_callable=StringIO)
    def test_add_auto_provisions_when_no_free_slots(
        self, mock_stderr: MagicMock, mock_run_jiri: MagicMock
    ) -> None:
        from typing import Any

        def mock_run_jiri_side_effect(
            jiri_root: Path, args: list[str], **kwargs: Any
        ) -> MagicMock:
            if args[0:2] == ["worktree", "add"]:
                path = Path(args[2])
                path.mkdir(parents=True, exist_ok=True)
                with open(self.pool.registry_file, "a") as f:
                    f.write(f"{path}\n")
            return MagicMock()

        mock_run_jiri.side_effect = mock_run_jiri_side_effect

        args = argparse.Namespace(
            name="my-feat", pool_name=None, sync=False, json=False
        )

        with patch("sys.stdout", new_callable=StringIO) as mock_out:
            add_cmd.run(args, self.pool)
            self.assertIn(
                "No free worktrees available in the pool. Provisioning a new one...",
                mock_stderr.getvalue(),
            )
            self.assertIn(".jiri_root/worktrees/my-feat", mock_out.getvalue())

        mock_run_jiri.assert_called_once()
        call_args = mock_run_jiri.call_args[0][1]
        self.assertEqual(call_args[0:2], ["worktree", "add"])

        symlink_path = self.pool.worktrees_dir / "my-feat"
        self.assertTrue(symlink_path.is_symlink())

        worktrees = self.pool.get_worktrees()
        self.assertEqual(len(worktrees), 1)
        wt = worktrees[0]
        self.assertEqual(wt.get_state(), WorktreeState.LEASED)
        lease = wt.get_lease_info()
        self.assertIsNotNone(lease)
        assert lease is not None
        self.assertEqual(lease.task_id, "my-feat")

    def test_add_already_leased_task_fails(self) -> None:
        wt_path1 = self.jiri_root / "worktrees" / "wt1"
        wt_path2 = self.jiri_root / "worktrees" / "wt2"
        wt_path1.mkdir(parents=True, exist_ok=True)
        wt_path2.mkdir(parents=True, exist_ok=True)
        self.pool.registry_file.write_text(f"{wt_path1}\n{wt_path2}\n")

        args1 = argparse.Namespace(
            name="my-feat", pool_name=None, sync=False, json=False
        )
        add_cmd.run(args1, self.pool)

        args2 = argparse.Namespace(
            name="my-feat", pool_name=None, sync=False, json=False
        )
        with self.assertRaises(ValueError) as context:
            add_cmd.run(args2, self.pool)
        self.assertIn("already active", str(context.exception))


class TestPoolAddSubcommand(unittest.TestCase):
    def setUp(self) -> None:
        self.temp_dir = tempfile.TemporaryDirectory()
        self.fuchsia_dir = Path(self.temp_dir.name)
        self.jiri_root = self.fuchsia_dir / ".jiri_root"
        self.jiri_root.mkdir(parents=True, exist_ok=True)
        self.pool = WorktreePool(fuchsia_dir=str(self.fuchsia_dir))
        self.patcher_jiri = patch("worktree_pool.run_jiri")
        self.patcher_fx = patch("subcommands.pool_add.run_fx")
        self.mock_jiri = self.patcher_jiri.start()
        self.mock_fx = self.patcher_fx.start()

    def tearDown(self) -> None:
        self.patcher_jiri.stop()
        self.patcher_fx.stop()
        self.temp_dir.cleanup()

    def test_add_multiple_set_args(self) -> None:
        wt_path = self.pool.worktrees_dir / "wt1"
        self.pool.registry_file.write_text(f"{wt_path}\n")
        args = argparse.Namespace(
            name="wt1",
            set=["core.x64 --out out/core", "workbench.arm64 --out out/wb"],
        )
        pool_add_cmd.run(args, self.pool)
        self.assertEqual(self.mock_fx.call_count, 2)
        self.pool.get_worktree_by_name("wt1")

    def test_add_symlink_local_dir(self) -> None:
        src_local = self.fuchsia_dir / "local"
        src_local.mkdir(parents=True, exist_ok=True)
        (src_local / "file.txt").write_text("hello")

        wt_path = self.pool.worktrees_dir / "wt1"
        args = argparse.Namespace(
            name="wt1",
            set=None,
            symlink_local=True,
            copy_local=False,
        )
        pool_add_cmd.run(args, self.pool)

        dest_local = wt_path / "local"
        self.assertTrue(dest_local.is_symlink())
        self.assertEqual(dest_local.resolve(), src_local.resolve())

    def test_add_copy_local_dir(self) -> None:
        src_local = self.fuchsia_dir / "local"
        src_local.mkdir(parents=True, exist_ok=True)
        (src_local / "file.txt").write_text("hello")

        wt_path = self.pool.worktrees_dir / "wt1"
        args = argparse.Namespace(
            name="wt1",
            set=None,
            symlink_local=False,
            copy_local=True,
        )
        pool_add_cmd.run(args, self.pool)

        dest_local = wt_path / "local"
        self.assertTrue(dest_local.exists())
        self.assertFalse(dest_local.is_symlink())
        self.assertTrue((dest_local / "file.txt").exists())
        self.assertEqual((dest_local / "file.txt").read_text(), "hello")

    def test_add_local_dir_missing(self) -> None:
        args = argparse.Namespace(
            name="wt1",
            set=None,
            symlink_local=True,
            copy_local=False,
        )
        with self.assertRaises(SystemExit):
            pool_add_cmd.run(args, self.pool)


class TestWorktreeSelectedBuildDir(unittest.TestCase):
    def setUp(self) -> None:
        self.temp_dir = tempfile.TemporaryDirectory()
        self.wt_path = Path(self.temp_dir.name)
        self.wt = Worktree(
            name="test-wt",
            path=self.wt_path,
            main_checkout_dir=self.wt_path / "main",
        )

    def tearDown(self) -> None:
        self.temp_dir.cleanup()

    def test_default_selected_build_dir(self) -> None:
        with self.assertRaises(FileNotFoundError):
            self.wt.selected_build_dir()

    def test_empty_selected_build_dir(self) -> None:
        fx_build_dir_file = self.wt_path / ".fx-build-dir"
        fx_build_dir_file.write_text("")
        with self.assertRaises(ValueError):
            self.wt.selected_build_dir()

    def test_custom_selected_build_dir(self) -> None:
        fx_build_dir_file = self.wt_path / ".fx-build-dir"
        fx_build_dir_file.write_text("out/custom-dir")
        expected = (self.wt_path / "out" / "custom-dir").resolve()
        self.assertEqual(self.wt.selected_build_dir(), expected)

    def test_printer_highlights_selected_build_dir(self) -> None:
        other_dir = self.wt_path / "out" / "other"
        other_dir.mkdir(parents=True, exist_ok=True)
        (other_dir / "args.gn").write_text(
            'build_info_product = "core"\nbuild_info_board = "x64"\n'
        )

        default_symlink = self.wt_path / "out" / "default"
        default_symlink.symlink_to("other")

        fx_build_dir_file = self.wt_path / ".fx-build-dir"
        fx_build_dir_file.write_text("out/default")

        another_dir = self.wt_path / "out" / "another"
        another_dir.mkdir(parents=True, exist_ok=True)
        (another_dir / "args.gn").write_text(
            'build_info_product = "workbench"\nbuild_info_board = "arm64"\n'
        )

        with patch.object(
            self.wt, "get_sync_status", return_value=(SyncStatus.SYNCED, 0, 0)
        ):
            with patch("sys.stdout", new_callable=StringIO) as mock_out:
                WorktreePrinter.print_worktrees([self.wt])
                output = mock_out.getvalue()

        self.assertIn("out/other *:", output)
        self.assertIn("out/another:", output)
        self.assertNotIn("out/default", output)
        self.assertNotIn("out/another *:", output)

    def test_printer_highlights_selected_build_dir_with_color(self) -> None:
        other_dir = self.wt_path / "out" / "other"
        other_dir.mkdir(parents=True, exist_ok=True)
        (other_dir / "args.gn").write_text(
            'build_info_product = "core"\nbuild_info_board = "x64"\n'
        )

        default_symlink = self.wt_path / "out" / "default"
        default_symlink.symlink_to("other")

        fx_build_dir_file = self.wt_path / ".fx-build-dir"
        fx_build_dir_file.write_text("out/default")

        with patch.object(
            self.wt, "get_sync_status", return_value=(SyncStatus.SYNCED, 0, 0)
        ):
            with patch("utils.USE_COLORS", True):
                with patch("sys.stdout", new_callable=StringIO) as mock_out:
                    WorktreePrinter.print_worktrees([self.wt])
                    output = mock_out.getvalue()

        # Colors.GREEN is \033[92m, Colors.RESET is \033[0m
        self.assertIn("\033[92m", output)
        self.assertIn("out/other *:", output)
        self.assertIn("\033[0m", output)

    def test_printer_no_builds_no_trailing_newline(self) -> None:
        with patch.object(
            self.wt, "get_sync_status", return_value=(SyncStatus.SYNCED, 0, 0)
        ):
            with patch("utils.USE_COLORS", False):
                with patch("sys.stdout", new_callable=StringIO) as mock_out:
                    WorktreePrinter.print_worktrees([self.wt])
                    output = mock_out.getvalue()

        self.assertEqual(output, "test-wt\n")


class TestCompletionChoices(unittest.TestCase):
    def setUp(self) -> None:
        self.temp_dir = tempfile.TemporaryDirectory()
        self.fuchsia_dir = Path(self.temp_dir.name)
        self.jiri_root = self.fuchsia_dir / ".jiri_root"
        self.jiri_root.mkdir(parents=True, exist_ok=True)
        self.pool = WorktreePool(fuchsia_dir=str(self.fuchsia_dir))

        # Create physical worktree dirs
        self.wt1_path = self.pool.worktrees_dir / "wt1"
        self.wt2_path = self.pool.worktrees_dir / "wt2"
        self.wt3_path = self.pool.worktrees_dir / "wt3"

        self.wt1_path.mkdir(parents=True, exist_ok=True)
        self.wt2_path.mkdir(parents=True, exist_ok=True)
        self.wt3_path.mkdir(parents=True, exist_ok=True)

        # Register them
        with open(self.pool.registry_file, "w") as f:
            f.write(f"{self.wt1_path}\n")
            f.write(f"{self.wt2_path}\n")
            f.write(f"{self.wt3_path}\n")

        # wt1 is free
        # wt2 is leased
        wt2 = self.pool.get_worktree_by_name("wt2")
        wt2.acquire_lease(task_id="task-wt2")

        # wt3 is leased but also has leased symlink with different name
        wt3 = self.pool.get_worktree_by_name("wt3")
        wt3.acquire_lease(task_id="task-wt3")

        # Create leased symlink for task-wt3
        self.leased_link = self.pool.worktrees_dir / "task-wt3"
        self.leased_link.symlink_to("wt3")

    def tearDown(self) -> None:
        self.temp_dir.cleanup()

    def test_filter_physical_all(self) -> None:
        choices = self.pool.get_completion_choices("physical_all")
        self.assertEqual(sorted(choices), ["wt1", "wt2", "wt3"])

    def test_filter_physical_free(self) -> None:
        choices = self.pool.get_completion_choices("physical_free")
        self.assertEqual(choices, ["wt1"])

    def test_filter_physical_leased(self) -> None:
        choices = self.pool.get_completion_choices("physical_leased")
        self.assertEqual(sorted(choices), ["wt2", "wt3"])

    def test_filter_leased(self) -> None:
        choices = self.pool.get_completion_choices("leased")
        self.assertEqual(choices, ["task-wt3"])

    def test_filter_all(self) -> None:
        choices = self.pool.get_completion_choices("all")
        self.assertEqual(sorted(choices), ["task-wt3", "wt1", "wt2", "wt3"])


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