"""Scoped task-branch grants and durable compare-and-swap ownership ledger."""
from __future__ import annotations

import base64
import hashlib
import hmac
import json
import os
import re
import sqlite3
import time
from pathlib import Path
from typing import Any

from gitea_api_policy import PolicyError, _validate_ref, _validate_repo, _validate_sha

MAX_GRANT_BYTES = 4096
MAX_GRANT_SECONDS = 300
MAX_ANCESTRY_STEPS = 1024
KEY_FILE = Path(os.environ.get("HERMES_SCM_TASK_GRANT_KEY_FILE", "/vault/secrets/scm-task-grant"))
LEDGER_PATH = Path(os.environ.get("HERMES_SCM_TASK_LEDGER", "/scm-state/task-branches.db"))
ADOPTIONS_PATH = Path(os.environ.get("HERMES_SCM_TASK_ADOPTIONS", "/scm-adoptions/task-branch-adoptions.json"))
REQUIRED = frozenset({"repo", "ref", "base", "board", "root_task_id", "assignment_task_id", "run", "ordinal", "expires", "expected_old", "new_head", "continuation_kind"})
IDENTIFIER = re.compile(r"[A-Za-z0-9][A-Za-z0-9_.-]{0,127}\Z")
ZERO_SHA = "0" * 40


def _key() -> bytes:
    try:
        value = KEY_FILE.read_text(encoding="ascii").strip()
    except OSError as exc:
        raise PolicyError("SCM task grant key is unavailable") from exc
    if not re.fullmatch(r"[0-9a-f]{64}", value):
        raise PolicyError("SCM task grant key is invalid")
    return value.encode("ascii")


def _b64(value: bytes) -> str:
    return base64.urlsafe_b64encode(value).rstrip(b"=").decode("ascii")


def _unb64(value: str) -> bytes:
    if not isinstance(value, str) or len(value) > MAX_GRANT_BYTES or not value.isascii():
        raise PolicyError("SCM task grant is invalid")
    try:
        return base64.urlsafe_b64decode(value + "=" * (-len(value) % 4))
    except ValueError as exc:
        raise PolicyError("SCM task grant is invalid") from exc


def sign_grant(claims: dict[str, Any], key: bytes | None = None) -> str:
    """Create a compact signed grant; mediator callers hold the key privately."""
    if set(claims) != REQUIRED:
        raise PolicyError("SCM task grant claims are invalid")
    payload = json.dumps(claims, sort_keys=True, separators=(",", ":")).encode("ascii")
    signature = hmac.new(key or _key(), payload, hashlib.sha256).digest()
    return _b64(payload) + "." + _b64(signature)


def verify_grant(token: str, *, now: int | None = None, key: bytes | None = None) -> dict[str, Any]:
    """Verify one short-lived grant and normalize only safe repository fields."""
    parts = token.split(".") if isinstance(token, str) else []
    if len(parts) != 2:
        raise PolicyError("SCM task grant is invalid")
    payload, signature = _unb64(parts[0]), _unb64(parts[1])
    if not hmac.compare_digest(hmac.new(key or _key(), payload, hashlib.sha256).digest(), signature):
        raise PolicyError("SCM task grant signature is invalid")
    try:
        claims = json.loads(payload)
    except (TypeError, ValueError) as exc:
        raise PolicyError("SCM task grant is invalid") from exc
    if not isinstance(claims, dict) or set(claims) != REQUIRED:
        raise PolicyError("SCM task grant claims are invalid")
    claims["repo"] = _validate_repo(claims["repo"])
    claims["ref"] = _validate_ref(claims["ref"], "task branch")
    claims["base"] = _validate_ref(claims["base"], "base branch")
    claims["expected_old"] = _validate_sha(claims["expected_old"], "expected old SHA")
    claims["new_head"] = _validate_sha(claims["new_head"], "new head SHA")
    if not isinstance(claims["ordinal"], int) or isinstance(claims["ordinal"], bool):
        raise PolicyError("SCM task grant ordinal is invalid")
    if claims["continuation_kind"] not in {"", "repair", "review"}:
        raise PolicyError("SCM task grant continuation is invalid")
    expiry = claims["expires"]
    current = int(time.time()) if now is None else now
    if not isinstance(expiry, int) or not current < expiry <= current + MAX_GRANT_SECONDS:
        raise PolicyError("SCM task grant is expired")
    if not all(isinstance(claims[name], str) and IDENTIFIER.fullmatch(claims[name]) for name in ("board", "root_task_id", "assignment_task_id", "run")):
        raise PolicyError("SCM task grant binding is invalid")
    return claims


class TaskLedger:
    """Broker-owned branch ownership state; no silent claim of existing refs."""
    def __init__(self, path: Path = LEDGER_PATH) -> None:
        path.parent.mkdir(mode=0o700, parents=True, exist_ok=True)
        self.path = path
        with self._connect() as con:
            con.execute("create table if not exists task_branches (repo text not null, ref text not null, board text not null, root_task_id text not null, latest_head text not null, pr_number integer, primary key(repo, ref))")
            con.execute("create table if not exists task_roots (repo text not null, board text not null, root_task_id text not null, ref text not null, primary key(repo,board,root_task_id))")

    def _connect(self):
        return sqlite3.connect(self.path, timeout=5, isolation_level="IMMEDIATE")

    @staticmethod
    def _owner(claims: dict[str, Any]) -> tuple[str, str]:
        return claims["board"], claims["root_task_id"]

    def get(self, repo: str, ref: str) -> tuple[str, str, str] | None:
        with self._connect() as con:
            row = con.execute("select board,root_task_id,latest_head from task_branches where repo=? and ref=?", (repo, ref)).fetchone()
        return tuple(row) if row is not None else None

    @staticmethod
    def _adoption(claims: dict[str, Any], remote_head: str | None) -> bool:
        """Accept only an operator-reviewed exact migration record.

        The ConfigMap is mounted read-only and starts empty.  A task grant can
        consume a matching record but cannot create, edit, or broaden one.
        """
        try:
            raw = ADOPTIONS_PATH.read_bytes()
            value = json.loads(raw)
        except (OSError, ValueError, TypeError):
            return False
        if not isinstance(value, dict) or len(raw) > 64 * 1024:
            return False
        record = value.get(f"{claims['repo']}/{claims['ref']}")
        if not isinstance(record, dict):
            return False
        return (
            record.get("board") == claims["board"]
            and record.get("root_task_id") == claims["root_task_id"]
            and record.get("latest_head") == remote_head == claims["expected_old"]
        )

    def seed_adoption(self, record: dict[str, Any], remote_head: str | None) -> None:
        """Import one reviewed legacy branch only when Forgejo still matches it."""
        required = {"repo", "ref", "board", "root_task_id", "latest_head", "pr_number"}
        if set(record) != required or remote_head is None:
            raise PolicyError("task branch adoption record is invalid")
        repo = _validate_repo(record["repo"])
        ref = _validate_ref(record["ref"], "task branch")
        head = _validate_sha(record["latest_head"], "adoption head")
        if not all(isinstance(record[name], str) and IDENTIFIER.fullmatch(record[name]) for name in ("board", "root_task_id")):
            raise PolicyError("task branch adoption record is invalid")
        if not isinstance(record["pr_number"], int) or isinstance(record["pr_number"], bool) or record["pr_number"] < 1:
            raise PolicyError("task branch adoption pull request is invalid")
        with self._connect() as con:
            existing = con.execute("select board,root_task_id,latest_head from task_branches where repo=? and ref=?", (repo, ref)).fetchone()
            root = con.execute("select ref from task_roots where repo=? and board=? and root_task_id=?", (repo, record["board"], record["root_task_id"])).fetchone()
            if existing is not None:
                if tuple(existing[:2]) != (record["board"], record["root_task_id"]):
                    raise PolicyError("task branch adoption conflicts with ledger")
                # A broker-confirmed later revision is legitimate. The static
                # import record only proves the first seed, never rewrites it.
                return
            if head != remote_head:
                raise PolicyError("task branch adoption does not match its live head")
            if root is not None and root[0] != ref:
                raise PolicyError("task branch adoption forks a logical task")
            con.execute("insert into task_branches(repo,ref,board,root_task_id,latest_head,pr_number) values(?,?,?,?,?,?)", (repo, ref, record["board"], record["root_task_id"], head, record["pr_number"]))
            con.execute("insert into task_roots(repo,board,root_task_id,ref) values(?,?,?,?)", (repo, record["board"], record["root_task_id"], ref))

    def register(self, claims: dict[str, Any], *, remote_head: str | None) -> None:
        """Create only absent refs; adoption is an explicit operator action."""
        with self._connect() as con:
            row = con.execute("select board,root_task_id,latest_head from task_branches where repo=? and ref=?", (claims["repo"], claims["ref"])).fetchone()
            owner = self._owner(claims)
            root = con.execute("select ref from task_roots where repo=? and board=? and root_task_id=?", (claims["repo"], *owner)).fetchone()
            if root is not None and root[0] != claims["ref"]:
                raise PolicyError("logical task already owns a different branch")
            if row is not None and tuple(row[:2]) != owner:
                raise PolicyError("task branch is owned by another task")
            if row is None:
                if remote_head is not None:
                    if not self._adoption(claims, remote_head):
                        raise PolicyError("existing unregistered task branch requires operator adoption")
                    con.execute("insert into task_branches(repo,ref,board,root_task_id,latest_head) values(?,?,?,?,?)", (claims["repo"], claims["ref"], owner[0], owner[1], remote_head))
                    con.execute("insert into task_roots(repo,board,root_task_id,ref) values(?,?,?,?)", (claims["repo"], *owner, claims["ref"]))
                    return
                if claims["expected_old"] != ZERO_SHA:
                    raise PolicyError("new task branch must use the zero expected head")
                con.execute(
                    "insert into task_branches(repo,ref,board,root_task_id,latest_head) values(?,?,?,?,?)",
                    (claims["repo"], claims["ref"], owner[0], owner[1], ZERO_SHA),
                )
                con.execute("insert into task_roots(repo,board,root_task_id,ref) values(?,?,?,?)", (claims["repo"], *owner, claims["ref"]))
            elif remote_head != row[2]:
                raise PolicyError("task branch head changed; fetch and merge before retrying")

    def authorize_update(self, claims: dict[str, Any]) -> None:
        row = self.get(claims["repo"], claims["ref"])
        if row is None:
            raise PolicyError("task branch is not registered")
        if row[:2] != self._owner(claims):
            raise PolicyError("task branch is owned by another task")
        if row[2] != claims["expected_old"]:
            raise PolicyError("task branch head changed; fetch and merge before retrying")

    def authorize_current(self, claims: dict[str, Any]) -> None:
        """Authorize post-push PR prose only at the broker-confirmed new head."""
        row = self.get(claims["repo"], claims["ref"])
        if row is None or row[:2] != self._owner(claims) or row[2] != claims["new_head"]:
            raise PolicyError("task branch head changed; fetch and merge before retrying")

    def commit(self, claims: dict[str, Any]) -> None:
        self.authorize_update(claims)
        with self._connect() as con:
            changed = con.execute("update task_branches set latest_head=? where repo=? and ref=? and latest_head=?", (claims["new_head"], claims["repo"], claims["ref"], claims["expected_old"])).rowcount
        if changed != 1:
            raise PolicyError("task branch changed during broker update")