#!/usr/bin/env python3
"""Resolve a /next selection to its owning Beads store and act only there.

`stores` lists the usable stores read-only so consumers can route new work."""

from __future__ import annotations

import argparse
import fcntl
import hashlib
import importlib.util
import json
import os
import subprocess
import sys
import tempfile
from contextlib import contextmanager
from datetime import datetime, timezone
from pathlib import Path
from time import monotonic, sleep
from typing import Any, Iterator

sys.dont_write_bytecode = True

SCRIPT_DIR = Path(__file__).resolve().parent
HANDOFF_LIST = SCRIPT_DIR.parents[1] / "handoffs" / "scripts" / "list.sh"
COMMAND_TIMEOUT_SECONDS = 5
AMBIGUOUS_EXIT = 3
NOT_FOUND_EXIT = 4
UNAVAILABLE_EXIT = 5
STALE_EXIT = 6
CLAIM_VERSION = "agent-claim:v1"
SESSION_ID_VARIABLES = (
    "PI_SESSION_ID",
    "CLAUDE_SESSION_ID",
    "CLAUDE_CODE_SESSION_ID",
    "CODEX_SESSION_ID",
    "CODEX_THREAD_ID",
)


def load_collector() -> Any:
    spec = importlib.util.spec_from_file_location("next_collect", SCRIPT_DIR / "collect.py")
    module = importlib.util.module_from_spec(spec)
    sys.modules[spec.name] = module
    spec.loader.exec_module(module)
    return module


def usable_sources(collector: Any) -> tuple[bool, list[Any]]:
    return collector.discover_sources(Path.cwd())


def owns_local_issue(source: Any, issue_id: str) -> bool:
    try:
        result = subprocess.run(
            ["bd", "-C", str(source.directory), "show", issue_id, "--json", "--readonly"],
            capture_output=True,
            check=False,
            text=True,
            timeout=COMMAND_TIMEOUT_SECONDS,
        )
    except (OSError, subprocess.TimeoutExpired):
        return False
    return result.returncode == 0


def probe_issue(collector: Any, source: Any, issue_id: str) -> tuple[bool, str | None]:
    error = collector.store_error(source)
    if error is not None:
        return False, error
    try:
        result = subprocess.run(
            [
                "bd",
                "-C",
                str(source.directory),
                "list",
                "--id",
                issue_id,
                "--all",
                "--json",
                "--readonly",
            ],
            capture_output=True,
            check=False,
            text=True,
            timeout=COMMAND_TIMEOUT_SECONDS,
        )
    except subprocess.TimeoutExpired:
        return False, f"timed out after {COMMAND_TIMEOUT_SECONDS} seconds"
    except OSError as command_error:
        return False, collector.concise_text(str(command_error))
    if result.returncode != 0:
        return False, collector.diagnostic_text(result)
    try:
        issues = json.loads(result.stdout)
    except json.JSONDecodeError as json_error:
        return False, f"invalid bd JSON: {json_error}"
    if not isinstance(issues, list) or not all(isinstance(issue, dict) for issue in issues):
        return False, "invalid bd JSON: expected an issue list"
    return any(issue.get("id") == issue_id for issue in issues), None


def resolution(workspace: bool, source: Any, issue_id: str) -> dict[str, Any]:
    return {
        "status": "resolved",
        "workspace": workspace,
        "id": issue_id,
        "repository": source.name,
        "repository_path": source.relative_path,
        "directory": str(source.directory.resolve()),
    }


def store_listing(collector: Any) -> dict[str, Any]:
    workspace, sources = usable_sources(collector)
    stores = []
    for source in sources:
        error = collector.store_error(source)
        row = {
            "repository": source.name,
            "repository_path": source.relative_path,
            "directory": str(source.directory.resolve()),
            "usable": error is None,
            "error": error,
        }
        if source.store == "workspace":
            row.update(
                owner="workspace",
                directory=str(sources[0].directory.resolve()),
                repository_directory=str(source.directory.resolve()),
            )
        stores.append(row)
    return {"workspace": workspace, "stores": stores}


def ranked_candidates(options: list[str]) -> list[dict[str, Any]]:
    result = subprocess.run(
        [str(SCRIPT_DIR / "next-bd"), "--json", *options],
        capture_output=True,
        check=False,
        text=True,
    )
    if result.returncode != 0:
        return []
    try:
        candidates = json.loads(result.stdout)
    except json.JSONDecodeError:
        return []
    return candidates if isinstance(candidates, list) else []


def source_by_name(sources: list[Any], name: str) -> Any | None:
    for source in sources:
        if source.name == name:
            return source
    return None


def owner_match(source: Any, issue_id: str) -> dict[str, str]:
    return {
        "repository": source.name,
        "repository_path": source.relative_path,
        "selector": f"{source.name}:{issue_id}",
    }


def source_failure(collector: Any, source: Any, error: str) -> dict[str, str]:
    failure = {
        "repository": source.name,
        "repository_path": source.relative_path,
        "error": error,
    }
    if (
        source.relative_path != "."
        and not source.store_declared
        and error == collector.MISSING_STORE_ERROR
        and not os.path.lexists(source.directory / ".beads")
    ):
        failure["hint"] = (
            "beadsStore is omitted (defaults to local) and .beads is absent. "
            "Confirm tracking policy: if workspace-owned, declare "
            '"beadsStore": "workspace" for this member in workspace.json; '
            "otherwise restore its local store. This hint does not prove ownership "
            "or authorize configuration changes."
        )
    return failure


def resolve_index(
    workspace: bool, sources: list[Any], index: int, options: list[str]
) -> tuple[dict[str, Any], int]:
    candidates = ranked_candidates(options)
    if index < 1 or index > len(candidates):
        return {"status": "not-found", "selector": str(index)}, NOT_FOUND_EXIT
    candidate = candidates[index - 1]
    name = candidate.get("repository", sources[0].name)
    source = source_by_name(sources, name)
    if source is None:
        return {"status": "not-found", "selector": str(index)}, NOT_FOUND_EXIT
    return resolution(workspace, source, candidate["id"]), 0


def resolve_id(
    collector: Any,
    workspace: bool,
    sources: list[Any],
    issue_id: str,
    qualifier: str | None,
) -> tuple[dict[str, Any], int]:
    selector = f"{qualifier}:{issue_id}" if qualifier is not None else issue_id
    if not workspace and qualifier is None:
        if owns_local_issue(sources[0], issue_id):
            return resolution(False, sources[0], issue_id), 0
        return {"status": "not-found", "selector": issue_id}, NOT_FOUND_EXIT
    if qualifier is not None:
        source = source_by_name(sources, qualifier)
        if source is None:
            return {"status": "not-found", "selector": selector}, NOT_FOUND_EXIT
        if source.store == "workspace" and collector.store_error(source) is None:
            return {
                "status": "not-found", "selector": selector,
                "owner": "workspace", "owner_selector": f"workspace:{issue_id}",
            }, NOT_FOUND_EXIT
        owned, error = probe_issue(collector, source, issue_id)
        if error is not None:
            return {
                "status": "unavailable",
                "selector": selector,
                "id": issue_id,
                "failures": [source_failure(collector, source, error)],
            }, UNAVAILABLE_EXIT
        if not owned:
            return {"status": "not-found", "selector": selector}, NOT_FOUND_EXIT
        return resolution(workspace, source, issue_id), 0

    owners = []
    failures = []
    for source in sources:
        if source.store == "workspace" and collector.store_error(source) is None:
            continue
        owned, error = probe_issue(collector, source, issue_id)
        if error is not None:
            failures.append(source_failure(collector, source, error))
        elif owned:
            owners.append(source)
    matches = [owner_match(source, issue_id) for source in owners]
    if len(owners) > 1:
        payload: dict[str, Any] = {
            "status": "ambiguous",
            "id": issue_id,
            "matches": matches,
        }
        if failures:
            payload["failures"] = failures
        return payload, AMBIGUOUS_EXIT
    if failures:
        return {
            "status": "unavailable",
            "selector": issue_id,
            "id": issue_id,
            "matches": matches,
            "failures": failures,
        }, UNAVAILABLE_EXIT
    if not owners:
        return {"status": "not-found", "selector": issue_id}, NOT_FOUND_EXIT
    return resolution(workspace, owners[0], issue_id), 0


def resolve(selector: str, options: list[str], expected_id: str) -> tuple[dict[str, Any], int]:
    collector = load_collector()
    workspace, sources = usable_sources(collector)
    if selector.isdigit():
        resolved, code = resolve_index(workspace, sources, int(selector), options)
    else:
        qualifier, separator, issue_id = selector.partition(":")
        if separator:
            resolved, code = resolve_id(
                collector, workspace, sources, issue_id, qualifier
            )
        else:
            resolved, code = resolve_id(
                collector, workspace, sources, selector, None
            )
    if code == 0 and expected_id and resolved["id"] != expected_id:
        return {
            "status": "stale",
            "selector": selector,
            "expected": expected_id,
            "actual": resolved["id"],
        }, STALE_EXIT
    return resolved, code


def run_handoff(directory: Path, issue_id: str, extra: list[str]) -> int:
    if not HANDOFF_LIST.is_file():
        print(f"next-select: handoff lookup unavailable: {HANDOFF_LIST}", file=sys.stderr)
        return 0
    result = subprocess.run(
        [str(HANDOFF_LIST), "--bead", issue_id, *extra],
        cwd=directory,
        check=False,
    )
    return result.returncode


def harness_session_id() -> str:
    for variable in SESSION_ID_VARIABLES:
        value = os.environ.get(variable, "").strip()
        if value:
            return " ".join(value.split())
    return "unavailable"


def claim_marker(directory: Path, session_id: str) -> str:
    identity = json.dumps(
        {
            "owning_store": str(directory.resolve()),
            "session_id": session_id,
        },
        sort_keys=True,
        separators=(",", ":"),
    )
    return f"{CLAIM_VERSION} {identity}"


def claim_lock_path(directory: Path, issue_id: str) -> Path:
    identity = f"{directory.resolve()}\0{issue_id}".encode()
    digest = hashlib.sha256(identity).hexdigest()
    return Path(tempfile.gettempdir()) / f"next-select-claim-{digest}.lock"


@contextmanager
def claim_lock(directory: Path, issue_id: str) -> Iterator[None]:
    descriptor = os.open(
        claim_lock_path(directory, issue_id),
        os.O_CREAT | os.O_RDWR,
        0o600,
    )
    lock = os.fdopen(descriptor, "w")
    deadline = monotonic() + COMMAND_TIMEOUT_SECONDS
    try:
        while True:
            try:
                fcntl.flock(lock, fcntl.LOCK_EX | fcntl.LOCK_NB)
                break
            except BlockingIOError:
                if monotonic() >= deadline:
                    raise TimeoutError("claim attribution lock timed out")
                sleep(0.05)
        yield
    finally:
        fcntl.flock(lock, fcntl.LOCK_UN)
        lock.close()


def record_claim(directory: Path, issue_id: str) -> int:
    session_id = harness_session_id()
    marker = claim_marker(directory, session_id)
    try:
        existing = subprocess.run(
            [
                "bd",
                "-C",
                str(directory),
                "comments",
                issue_id,
                "--json",
                "--readonly",
            ],
            capture_output=True,
            check=False,
            text=True,
            timeout=COMMAND_TIMEOUT_SECONDS,
        )
    except (OSError, subprocess.TimeoutExpired) as error:
        print(f"next-select: claim lookup failed: {error}", file=sys.stderr)
        return 1
    if existing.returncode != 0:
        sys.stderr.write(existing.stderr)
        return existing.returncode
    try:
        comments = json.loads(existing.stdout)
    except json.JSONDecodeError as error:
        print(f"next-select: invalid comment JSON: {error}", file=sys.stderr)
        return 1
    if not isinstance(comments, list):
        print("next-select: invalid comment JSON: expected a list", file=sys.stderr)
        return 1
    if any(
        isinstance(comment, dict)
        and isinstance(comment.get("text"), str)
        and marker in comment["text"]
        for comment in comments
    ):
        return 0

    timestamp = datetime.now(timezone.utc).isoformat(timespec="seconds").replace(
        "+00:00", "Z"
    )
    body = "\n".join(
        (
            marker,
            "Agent-driven in_progress claim attribution.",
            f"session_id: {session_id}",
            f"claimed_at_utc: {timestamp}",
            f"owning_store: {directory.resolve()}",
            "Session activity is unverified.",
        )
    )
    try:
        added = subprocess.run(
            ["bd", "-C", str(directory), "comments", "add", issue_id, body],
            check=False,
            timeout=COMMAND_TIMEOUT_SECONDS,
        )
    except (OSError, subprocess.TimeoutExpired) as error:
        print(f"next-select: claim attribution failed: {error}", file=sys.stderr)
        return 1
    return added.returncode


def run_start(directory: Path, issue_id: str) -> int:
    try:
        with claim_lock(directory, issue_id):
            claimed = record_claim(directory, issue_id)
            if claimed != 0:
                return claimed
            update = subprocess.run(
                [
                    "bd",
                    "-C",
                    str(directory),
                    "update",
                    issue_id,
                    "--status=in_progress",
                ],
                check=False,
                timeout=COMMAND_TIMEOUT_SECONDS,
            )
            if update.returncode != 0:
                return update.returncode
            return subprocess.run(
                ["bd", "-C", str(directory), "show", issue_id],
                check=False,
                timeout=COMMAND_TIMEOUT_SECONDS,
            ).returncode
    except (OSError, subprocess.TimeoutExpired, TimeoutError) as error:
        print(f"next-select: claim failed: {error}", file=sys.stderr)
        return 1


def parse_arguments(argv: list[str]) -> argparse.Namespace:
    parser = argparse.ArgumentParser(description=__doc__)
    parser.add_argument("command", choices=("stores", "resolve", "handoff", "start"))
    parser.add_argument("selector", nargs="?", default="")
    parser.add_argument("--avoid-busy", action="store_true")
    parser.add_argument("--type", default="")
    parser.add_argument("--expect-id", default="")
    parser.add_argument("--check-branches", action="store_true")
    arguments = parser.parse_args(argv)
    if arguments.command == "stores" and arguments.selector:
        parser.error("stores takes no selector")
    if arguments.command != "stores" and not arguments.selector:
        parser.error(f"{arguments.command} requires a selector")
    return arguments


def main(argv: list[str]) -> int:
    arguments = parse_arguments(argv)
    if arguments.command == "stores":
        json.dump(store_listing(load_collector()), sys.stdout, separators=(",", ":"))
        sys.stdout.write("\n")
        return 0
    options = []
    if arguments.avoid_busy:
        options.append("--avoid-busy")
    if arguments.type:
        options.append(f"--type={arguments.type}")

    resolved, code = resolve(arguments.selector, options, arguments.expect_id)
    if code != 0 or arguments.command == "resolve":
        json.dump(resolved, sys.stdout, separators=(",", ":"))
        sys.stdout.write("\n")
        return code

    directory = Path(resolved["directory"])
    if arguments.command == "handoff":
        return run_handoff(
            directory,
            resolved["id"],
            ["--check-branches"] if arguments.check_branches else [],
        )
    return run_start(directory, resolved["id"])


if __name__ == "__main__":
    raise SystemExit(main(sys.argv[1:]))
