"""Global task-worktree registry.

Tracks which `(project_id, task_group, task_id)` task-key owns which
on-disk worktree path and branch, across concurrent okstra runs on the
same machine. The registry lives under `OKSTRA_HOME` (default
`~/.okstra`) and is guarded by an `fcntl` exclusive lock so that two
processes cannot race to reserve the same path or branch.

Why a global registry:
  - A single task-key spans multiple phases (requirements-discovery →
    error-analysis → implementation-option-selection →
    implementation-planning → implementation). All phases must land in
    the **same** worktree on the **same** branch. Re-entry from any
    phase must look up the existing entry instead of creating a
    duplicate.
  - Two different task-keys must never collide on the same branch name.
    A global branch index makes that detectable cheaply.
  - Cleanup of stale entries (worktree dir removed manually) needs a
    single source of truth.

The registry is intentionally JSON-on-disk (no SQLite): the data set is
tiny (one row per active task on this machine) and the human-readable
file is useful for debugging.
"""
from __future__ import annotations

import time
from dataclasses import dataclass
from pathlib import Path
from typing import Optional

from okstra_project.dirs import okstra_home

from .json_registry import load_registry_json, registry_lock, save_registry_json


REGISTRY_FILENAME = "registry.json"
LOCK_FILENAME = "registry.lock"


def _okstra_worktrees_dir() -> Path:
    return okstra_home() / "worktrees"


def task_key(
    project_id: str, task_group: str, task_id: str,
    stage_number: Optional[int] = None,
    group_id: Optional[str] = None,
) -> str:
    """Canonical task-key. stage_number → `#stage-<N>` (per-stage worktree),
    group_id → `#group-<id>` (stage-group collector worktree). 둘은 상호배타."""
    if stage_number is not None and group_id is not None:
        raise ValueError("stage_number and group_id are mutually exclusive")
    base = f"{project_id}/{task_group}/{task_id}"
    if stage_number is not None:
        return f"{base}#stage-{stage_number}"
    if group_id is not None:
        return f"{base}#group-{group_id}"
    return base


@dataclass
class WorktreeEntry:
    task_key: str
    project_id: str
    task_group: str
    task_id: str
    worktree_path: str
    branch: str
    base_ref: str
    created_at: str
    last_phase: str = ""
    status: str = "active"  # "active" | "released"
    stage: Optional[int] = None
    implementation_base_commit: str = ""
    stages: Optional[list] = None


def _registry_lock():
    """Exclusive flock on `<worktrees>/registry.lock`."""
    return registry_lock(_okstra_worktrees_dir() / LOCK_FILENAME)


def _registry_path() -> Path:
    return _okstra_worktrees_dir() / REGISTRY_FILENAME


def _load() -> dict:
    data = load_registry_json(_registry_path(), lambda: {"tasks": {}, "branches": {}})
    data.setdefault("tasks", {})
    data.setdefault("branches", {})
    return data


def _save(data: dict) -> None:
    save_registry_json(_registry_path(), data)


def lookup(
    project_id: str, task_group: str, task_id: str,
    stage_number: Optional[int] = None,
    group_id: Optional[str] = None,
) -> Optional[WorktreeEntry]:
    key = task_key(project_id, task_group, task_id, stage_number, group_id)
    with _registry_lock():
        data = _load()
        row = data["tasks"].get(key)
        if not row:
            return None
        return WorktreeEntry(task_key=key, **row)


def reserve(
    *,
    project_id: str,
    task_group: str,
    task_id: str,
    worktree_path: str,
    branch: str,
    base_ref: str,
    phase: str = "",
    stage_number: Optional[int] = None,
    group_id: Optional[str] = None,
    stages: Optional[list] = None,
) -> WorktreeEntry:
    """Atomically insert a new entry. Raises RuntimeError if the
    task-key already exists or the branch is already owned by a
    different task-key. Callers should `lookup()` first when re-entry
    is expected.
    """
    key = task_key(project_id, task_group, task_id, stage_number, group_id)
    now = time.strftime("%Y-%m-%dT%H:%M:%S%z") or time.strftime("%Y-%m-%dT%H:%M:%S")
    with _registry_lock():
        data = _load()
        existing = data["tasks"].get(key)
        if existing and existing.get("status") != "released":
            raise RuntimeError(
                f"task-key already has a worktree registered: {key} → "
                f"{existing['worktree_path']} (branch {existing['branch']}). "
                "Use `lookup` to reuse it, or release it before reserving anew."
            )
        owner = data["branches"].get(branch)
        if owner and owner != key:
            raise RuntimeError(
                f"branch {branch!r} is already registered to a different "
                f"task-key: {owner}. Choose a different work-category or "
                "release the conflicting task first."
            )
        row = {
            "project_id": project_id,
            "task_group": task_group,
            "task_id": task_id,
            "worktree_path": worktree_path,
            "branch": branch,
            "base_ref": base_ref,
            "created_at": now,
            "last_phase": phase,
            "status": "active",
            "stage": stage_number,
            "stages": stages,
        }
        # Re-reserving a released task-key must not drop the anchor: the
        # implementation base commit is task-lifetime state, not per-reservation,
        # so a release()/reserve() cycle (e.g. local_checkout) keeps it.
        if existing and existing.get("implementation_base_commit"):
            row["implementation_base_commit"] = existing["implementation_base_commit"]
        data["tasks"][key] = row
        data["branches"][branch] = key
        _save(data)
        return WorktreeEntry(task_key=key, **row)


def touch_phase(project_id: str, task_group: str, task_id: str, phase: str) -> None:
    """Record the most recent phase observed on this worktree.
    Best-effort: silently no-ops if the task-key is not registered.
    """
    key = task_key(project_id, task_group, task_id)
    with _registry_lock():
        data = _load()
        row = data["tasks"].get(key)
        if not row:
            return
        row["last_phase"] = phase
        _save(data)


def set_implementation_base(
    project_id: str, task_group: str, task_id: str, commit: str,
) -> str:
    """Fix the shared base commit for this task's implementation stages,
    once. Idempotent: if already set, the existing value is returned and
    `commit` is ignored (so two concurrent first-stage runs converge).
    Raises RuntimeError when the task-key entry does not exist."""
    key = task_key(project_id, task_group, task_id)
    with _registry_lock():
        data = _load()
        row = data["tasks"].get(key)
        if row is None:
            raise RuntimeError(
                f"no task-key entry to anchor implementation base: {key}"
            )
        already = row.get("implementation_base_commit")
        if already:
            return already
        row["implementation_base_commit"] = commit
        _save(data)
        return commit


def reset_implementation_base(
    project_id: str, task_group: str, task_id: str, commit: str,
) -> str:
    """anchor 를 의식적으로 재고정한다. 유일한 호출자는 git-reconcile 의
    `--reset-anchor` — prepare 경로는 절대 anchor 를 움직이지 않는다."""
    key = task_key(project_id, task_group, task_id)
    with _registry_lock():
        data = _load()
        row = data["tasks"].get(key)
        if row is None:
            raise RuntimeError(
                f"no task-key entry to reset implementation base: {key}"
            )
        row["implementation_base_commit"] = commit
        _save(data)
        return commit


def get_implementation_base(
    project_id: str, task_group: str, task_id: str,
) -> Optional[str]:
    """Return the fixed implementation base commit, or None when unset /
    task-key missing."""
    key = task_key(project_id, task_group, task_id)
    with _registry_lock():
        data = _load()
        row = data["tasks"].get(key)
        if row is None:
            return None
        return row.get("implementation_base_commit") or None


def get_stage_row(
    project_id: str, task_group: str, task_id: str, stage: int,
) -> Optional[dict]:
    """Return the stage-key registry row (worktree_path / base_ref / branch)
    for `<task-key>#stage-<stage>`, or None when no such reservation exists."""
    key = task_key(project_id, task_group, task_id, stage)
    with _registry_lock():
        data = _load()
        return data["tasks"].get(key)


def list_active_stage_numbers(
    project_id: str, task_group: str, task_id: str,
) -> set:
    """Return the set of stage numbers with an active stage-key reservation
    for this task. Used by the stage resolver to exclude stages a concurrent
    run already holds (the occupancy SSOT). Excludes the task-key entry
    (stage is None) and released entries."""
    prefix = task_key(project_id, task_group, task_id) + "#stage-"
    with _registry_lock():
        data = _load()
        out = set()
        for key, row in data["tasks"].items():
            if (key.startswith(prefix)
                    and row.get("status") == "active"
                    and row.get("stage") is not None):
                out.add(row["stage"])
        return out


def release_status(
    project_id: str, task_group: str, task_id: str,
    stage_number: Optional[int] = None,
    group_id: Optional[str] = None,
) -> Optional[WorktreeEntry]:
    """Mark the entry as `released` WITHOUT freeing its branch index slot.

    Use when the caller must not touch the branch index — either the slot is
    still needed, or another caller owns its lifetime. Stage teardown is the
    latter case: `consumers._release_stage_reservation` already ran `release()`
    at stage-done time, so teardown only flips the status. Returns the prior
    entry, or None.
    """
    key = task_key(project_id, task_group, task_id, stage_number, group_id)
    with _registry_lock():
        data = _load()
        row = data["tasks"].get(key)
        if not row:
            return None
        row["status"] = "released"
        _save(data)
        return WorktreeEntry(task_key=key, **row)


def free_branch_slot(
    project_id: str, task_group: str, task_id: str,
    stage_number: Optional[int] = None,
    group_id: Optional[str] = None,
) -> bool:
    """Free only the branch index slot for this key. Returns True when a
    slot was freed.

    Call this ONLY after the physical branch is deleted, so the invariant
    "a branch name the registry marks free is always physically deleted"
    holds even under concurrent runs on the same machine.
    """
    key = task_key(project_id, task_group, task_id, stage_number, group_id)
    with _registry_lock():
        data = _load()
        row = data["tasks"].get(key)
        if not row:
            return False
        branch = row.get("branch")
        if branch and data["branches"].get(branch) == key:
            del data["branches"][branch]
            _save(data)
            return True
        return False


def release(
    project_id: str, task_group: str, task_id: str,
    stage_number: Optional[int] = None,
    group_id: Optional[str] = None,
) -> Optional[WorktreeEntry]:
    """Mark the entry as `released` and free its branch slot (worktree dir
    intact — preservation is the project's policy). Used at the points that
    end a reservation's life: done-time stage release
    (`consumers._release_stage_reservation`) and group cleanup. Both free the
    slot even when the physical branch survives — stage stack branches are
    deliberately kept, and their names are task-scoped, so a freed slot cannot
    collide with a live branch from another task.
    Returns the prior entry, or None when not found.
    """
    entry = release_status(project_id, task_group, task_id, stage_number, group_id)
    if entry is not None:
        free_branch_slot(project_id, task_group, task_id, stage_number, group_id)
    return entry