# 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 random
import secrets
import xml.etree.ElementTree as ET
from pathlib import Path

from utils import run_jiri
from worktree import NoFreeWorktreesError, Worktree, WorktreeState

ADJECTIVES = (
    "amber",
    "blue",
    "bold",
    "calm",
    "cool",
    "coral",
    "crisp",
    "cyan",
    "dawn",
    "dry",
    "dune",
    "echo",
    "fast",
    "firm",
    "gray",
    "green",
    "hazy",
    "jade",
    "keen",
    "lime",
    "mint",
    "mist",
    "navy",
    "neat",
    "pale",
    "pink",
    "pure",
    "quiet",
    "red",
    "rich",
    "sage",
    "slate",
    "solar",
    "stone",
    "swift",
    "teal",
    "vivid",
    "warm",
    "wild",
    "yellow",
)

NOUNS = (
    "badger",
    "basalt",
    "canyon",
    "cavern",
    "crag",
    "crest",
    "falcon",
    "field",
    "forest",
    "glacier",
    "groove",
    "harbor",
    "haven",
    "heron",
    "island",
    "lagoon",
    "meadow",
    "mesa",
    "moss",
    "orchard",
    "osprey",
    "otter",
    "pebble",
    "pine",
    "prairie",
    "quarry",
    "rapids",
    "ravine",
    "reef",
    "ridge",
    "river",
    "summit",
    "thicket",
    "tundra",
    "valley",
)

# Pool directory used when `.jiri_root/config` doesn't set `worktrees_dir`,
# relative to the main checkout's root. Matches jiri's default.
DEFAULT_WORKTREES_DIR = ".jiri_root/worktrees"


def _read_jiri_config(checkout: Path) -> ET.Element | None:
    """Parses `<checkout>/.jiri_root/config`, or returns None if unreadable."""
    try:
        return ET.parse(checkout / ".jiri_root" / "config").getroot()
    except (OSError, ET.ParseError):
        return None


def _find_main_checkout(checkout: Path) -> Path:
    """Returns the main checkout that `checkout` is a worktree of.

    Jiri records a worktree's parent checkout as `<worktree><parent_root>` in
    the worktree's `.jiri_root/config`. A checkout without one is a main
    checkout, so it's returned as is.
    """
    seen = set()
    # `seen` guards against a corrupt config whose parents form a cycle.
    while checkout not in seen:
        seen.add(checkout)
        config = _read_jiri_config(checkout)
        parent = None
        if config is not None:
            parent = config.findtext("worktree/parent_root")
        if not parent:
            break
        checkout = Path(parent).resolve()
    return checkout


class WorktreePool:
    def __init__(self, fuchsia_dir: str | None = None):
        if fuchsia_dir is None:
            fuchsia_dir = os.environ.get("FUCHSIA_DIR")
        if fuchsia_dir is None:
            fuchsia_dir = str(
                Path(__file__).resolve().parent.parent.parent.parent
            )

        # When run from inside a worktree, operate on the pool of the main
        # checkout it belongs to.
        self.fuchsia_dir = _find_main_checkout(Path(fuchsia_dir).resolve())

        self.jiri_root = self.fuchsia_dir / ".jiri_root"
        self.registry_file = self.jiri_root / "worktrees_registry"

        # Configured with `jiri init -worktrees-dir=<path>`. A relative path
        # is relative to the main checkout's root.
        config = _read_jiri_config(self.fuchsia_dir)
        worktrees_dir = DEFAULT_WORKTREES_DIR
        if config is not None:
            worktrees_dir = config.findtext("worktrees_dir") or worktrees_dir
        self.worktrees_dir = (self.fuchsia_dir / worktrees_dir).resolve()

    def get_worktrees(self) -> list[Worktree]:
        if not self.registry_file.exists():
            return []
        worktrees = []
        with open(self.registry_file, "r") as f:
            for line in f:
                line = line.strip()
                if line:
                    path = Path(line)
                    # worktrees_dir only decides where fx worktree provisions
                    # slots. The jiri registry lists every worktree of the
                    # checkout, including ones created elsewhere (e.g. with
                    # `jiri worktree add <path>`, which takes any path) and
                    # slots left over from before worktrees_dir changed;
                    # those aren't pool slots.
                    if path.resolve().parent != self.worktrees_dir:
                        continue
                    worktrees.append(
                        Worktree(path.name, path, self.fuchsia_dir)
                    )
        return worktrees

    def get_worktree_by_name(self, name: str) -> Worktree:
        for wt in self.get_worktrees():
            if wt.name == name:
                return wt
        raise KeyError(f"Worktree '{name}' not found")

    def find_worktree(self, identifier: str) -> Worktree:
        for wt in self.get_worktrees():
            if wt.name == identifier:
                return wt
            if wt.get_state() == WorktreeState.LEASED:
                lease = wt.get_lease_info()
                if lease and lease.task_id == identifier:
                    return wt
        raise KeyError(f"Worktree '{identifier}' not found")

    def get_any_free_worktree(self) -> Worktree:
        free_worktrees = [
            wt
            for wt in self.get_worktrees()
            if wt.get_state() == WorktreeState.FREE
        ]
        if not free_worktrees:
            raise NoFreeWorktreesError("No free worktrees available")
        return random.choice(free_worktrees)

    def _generate_random_pool_name(self) -> str:
        existing = {wt.name for wt in self.get_worktrees()}
        for _ in range(100):
            candidate = f"{random.choice(ADJECTIVES)}-{random.choice(NOUNS)}"
            if (
                candidate not in existing
                and not (self.worktrees_dir / candidate).exists()
            ):
                return candidate
        return f"slot-{secrets.token_hex(3)}"

    def add_worktree(self, name: str | None = None) -> Worktree:
        if name is None:
            name = self._generate_random_pool_name()
        wt_path = self.worktrees_dir / name
        if wt_path.exists():
            raise FileExistsError(f"Worktree path '{wt_path}' already exists")

        run_jiri(self.jiri_root, ["worktree", "add", str(wt_path)], check=True)
        wt = Worktree(name, wt_path, self.fuchsia_dir)
        wt.meta_dir.mkdir(parents=True, exist_ok=True)
        return wt

    def remove_worktree(self, name: str, force: bool = False) -> None:
        wt_path = self.worktrees_dir / name
        if not wt_path.exists():
            raise FileNotFoundError(f"Worktree '{name}' does not exist on disk")

        cmd = ["worktree", "remove"]
        if force:
            cmd.append("-force")
        cmd.append(str(wt_path))
        run_jiri(self.jiri_root, cmd, check=True)

    def get_completion_choices(self, filter_type: str) -> list[str]:
        matches = []
        if filter_type in ("all", "leased"):
            if self.worktrees_dir.exists() and self.worktrees_dir.is_dir():
                for entry in self.worktrees_dir.iterdir():
                    if entry.is_symlink():
                        try:
                            target = entry.readlink()
                            if not target.is_absolute():
                                target = (self.worktrees_dir / target).resolve()
                            if target.is_dir():
                                matches.append(entry.name)
                        except OSError:
                            pass

        if filter_type != "leased":
            for wt in self.get_worktrees():
                is_leased = wt.get_state() == WorktreeState.LEASED
                if filter_type == "physical_free" and not is_leased:
                    matches.append(wt.name)
                elif filter_type == "physical_leased" and is_leased:
                    matches.append(wt.name)
                elif filter_type in ("physical_all", "all"):
                    matches.append(wt.name)

        seen = set()
        unique_matches = []
        for m in matches:
            if m not in seen:
                seen.add(m)
                unique_matches.append(m)
        return unique_matches