#!/usr/bin/env python3
import base64
import copy
import hashlib
import json
import os
import pathlib
import socket
import time
import urllib.error
import urllib.parse
import urllib.request

SOURCE_GOVERNANCE = "/opt/aios/aios-control-plane/current/active-runtime/app/governance"
if SOURCE_GOVERNANCE not in __import__("sys").path:
    __import__("sys").path.insert(0, SOURCE_GOVERNANCE)
from stage2_gate_migration import (
    CAMPAIGN_GATE,
    CAMPAIGN_ID,
    OWNER_GATE,
    STAGE1_SEMANTIC_SHA,
    Stage2GateError,
    build_target,
    semantic_digest as stage2_semantic_digest,
    validate_stage2_facts,
)

API = "http://127.0.0.1:3000"
REPO = "aios/aios"
SOCKET = "/run/aios-governance-executor/executor.sock"
JOURNAL = pathlib.Path("/var/lib/aios-gitops/governance-executor/executor-journal.jsonl")
RECONCILIATION_BINDING_PATH = pathlib.Path("/var/lib/aios-gitops/governance-executor/reconciliation-binding.json")
PR54_HEAD = "015185d7e5191d18c5c1596160807ac228782a0a"
PR54_REVIEW = 104
PR54_MANIFEST_SHA = "1b11ee42d53a9f6c1646e5f5ddd8ce5b7ecf301483b9f3d47a068290d541ba4b"
PR54_PREVIOUS_PROTECTION_SHA = "8b3370add21cd23c2b4f41d8d84d3eaaa09ddf0451549b72c7941bbd92d5d60a"
SEMANTIC_SCHEMA = "aios.semantic-protection.v1"
APPROVED_SEMANTIC_SHA = "ba15d512ed8c1f43e5f72c84125a3a71a7347ef59e18000bbfca24d6e69f1b19"
STAGE1_SEMANTIC_SHA = "d911e856f71de6105b8efa08e13338a448e0b6d37536332decd80c9f409ecf27"
USER_IDS = {"owner": 2, "aios-autonomy-agent": 153}


def canon(value):
    return json.dumps(value, sort_keys=True, separators=(",", ":"), ensure_ascii=False).encode()


def sha(value):
    return hashlib.sha256(canon(value)).hexdigest()


def token():
    value = (pathlib.Path(os.environ["CREDENTIALS_DIRECTORY"]) / "gitea-governance-executor-token").read_text().strip()
    if len(value) < 16:
        raise RuntimeError("EXECUTOR_CREDENTIAL_INVALID")
    return value


def api(method, path, value, body=None, expected=(200,)):
    raw = None if body is None else json.dumps(body, separators=(",", ":")).encode()
    req = urllib.request.Request(API + path, data=raw, method=method)
    req.add_header("Authorization", "token " + value)
    req.add_header("Accept", "application/json")
    if raw is not None:
        req.add_header("Content-Type", "application/json")
    try:
        with urllib.request.urlopen(req, timeout=20) as response:
            status = response.status
            payload = response.read()
    except urllib.error.HTTPError as exc:
        if exc.code in expected:
            payload = exc.read()
            return exc.code, json.loads(payload) if payload else None
        raise RuntimeError("GITEA_HTTP_" + str(exc.code))
    if status not in expected:
        raise RuntimeError("GITEA_HTTP_" + str(status))
    if not payload or not payload.strip():
        raise RuntimeError("GITEA_EMPTY_JSON:" + path)
    try:
        decoded = json.loads(payload)
    except (UnicodeDecodeError, json.JSONDecodeError) as exc:
        raise RuntimeError("GITEA_INVALID_JSON:" + path) from exc
    return status, decoded


def journal(event, **fields):
    JOURNAL.parent.mkdir(parents=True, exist_ok=True)
    record = {"schema": "aios.governance-executor-journal.v2", "event": event, "unix": int(time.time()), **fields}
    with JOURNAL.open("a", encoding="utf-8") as stream:
        stream.write(json.dumps(record, sort_keys=True, separators=(",", ":"), ensure_ascii=False) + "\n")
        stream.flush()
        os.fsync(stream.fileno())


def typed_json(value):
    if value in (None, "null"):
        return []
    if isinstance(value, str):
        return json.loads(value)
    return value


def user_ids(names):
    result = []
    for name in names or []:
        if name not in USER_IDS:
            raise RuntimeError("UNKNOWN_GOVERNANCE_USER:" + str(name))
        result.append(USER_IDS[name])
    return sorted(result)


def semantic_policy(value):
    policy = {
        "schema_version": SEMANTIC_SCHEMA,
        "repo_id": 1,
        "protection_id": 2,
        "branch_name": value.get("branch_name"),
        "can_push": value.get("enable_push"),
        "can_force_push": value.get("enable_force_push"),
        "enable_merge_whitelist": value.get("enable_merge_whitelist"),
        "merge_whitelist_user_ids": user_ids(value.get("merge_whitelist_usernames")),
        "merge_whitelist_team_ids": [],
        "required_approvals": value.get("required_approvals"),
        "enable_status_check": value.get("enable_status_check"),
        "status_check_contexts": sorted(value.get("status_check_contexts") or []),
        "block_admin_merge_override": value.get("block_admin_merge_override"),
        "block_on_official_review_requests": value.get("block_on_official_review_requests"),
        "block_on_outdated_branch": value.get("block_on_outdated_branch"),
        "block_on_rejected_reviews": value.get("block_on_rejected_reviews"),
        "dismiss_stale_approvals": value.get("dismiss_stale_approvals"),
        "ignore_stale_approvals": value.get("ignore_stale_approvals"),
        "enable_approvals_whitelist": value.get("enable_approvals_whitelist"),
        "approvals_whitelist_user_ids": user_ids(value.get("approvals_whitelist_username")),
        "approvals_whitelist_team_ids": [],
        "enable_bypass_allowlist": value.get("enable_bypass_allowlist"),
        "bypass_allowlist_user_ids": user_ids(value.get("bypass_allowlist_usernames")),
        "bypass_allowlist_team_ids": [],
        "enable_force_push_allowlist": value.get("enable_force_push_allowlist"),
        "force_push_allowlist_user_ids": user_ids(value.get("force_push_allowlist_usernames")),
        "force_push_allowlist_team_ids": [],
        "enable_whitelist": value.get("enable_push_whitelist"),
        "whitelist_user_ids": user_ids(value.get("push_whitelist_usernames")),
        "whitelist_team_ids": [],
        "force_push_allowlist_deploy_keys": value.get("force_push_allowlist_deploy_keys"),
        "whitelist_deploy_keys": value.get("push_whitelist_deploy_keys"),
        "require_signed_commits": value.get("require_signed_commits"),
        "priority": value.get("priority"),
        "protected_file_patterns": value.get("protected_file_patterns"),
        "unprotected_file_patterns": value.get("unprotected_file_patterns"),
    }
    return policy


def semantic_sha(value):
    return sha(semantic_policy(value))


def verify_binding(binding, value):
    required = {
        "governance_request_PR", "exact_request_head", "Owner_review_id", "Owner_review_exact_head",
        "manifest_sha256", "expected_previous_snapshot_sha256", "approved_delta",
        "semantic_reconciliation_request_id", "semantic_reconciliation_head",
        "semantic_reconciliation_review_id", "semantic_reconciliation_review_head",
        "approved_semantic_digest",
    }
    if set(binding) != required:
        raise RuntimeError("GOVERNANCE_BINDING_SCHEMA_INVALID")
    if binding["governance_request_PR"] != 54 or binding["exact_request_head"] != PR54_HEAD:
        raise RuntimeError("GOVERNANCE_BINDING_PR54_MISMATCH")
    if binding["Owner_review_id"] != PR54_REVIEW or binding["Owner_review_exact_head"] != PR54_HEAD:
        raise RuntimeError("GOVERNANCE_BINDING_REVIEW_MISMATCH")
    if binding["manifest_sha256"] != PR54_MANIFEST_SHA:
        raise RuntimeError("GOVERNANCE_BINDING_MANIFEST_MISMATCH")
    if binding["expected_previous_snapshot_sha256"] != PR54_PREVIOUS_PROTECTION_SHA:
        raise RuntimeError("GOVERNANCE_BINDING_SNAPSHOT_MISMATCH")
    if binding["approved_delta"] != {"add_merge_whitelist_user_id": 153}:
        raise RuntimeError("GOVERNANCE_BINDING_DELTA_DENIED")
    if binding["approved_semantic_digest"] != APPROVED_SEMANTIC_SHA:
        raise RuntimeError("GOVERNANCE_BINDING_SEMANTIC_DIGEST_MISMATCH")
    if not RECONCILIATION_BINDING_PATH.is_file():
        raise RuntimeError("SEMANTIC_RECONCILIATION_BINDING_REQUIRED")
    try:
        reconciliation = json.loads(RECONCILIATION_BINDING_PATH.read_text(encoding="utf-8"))
    except (OSError, ValueError):
        raise RuntimeError("SEMANTIC_RECONCILIATION_BINDING_INVALID") from None
    required_reconciliation = {
        "schema", "request_id", "pr_number", "request_head", "owner_review_id",
        "owner_review_head", "approved_semantic_digest", "approved_delta",
    }
    if set(reconciliation) != required_reconciliation or reconciliation.get("schema") != "aios.semantic-reconciliation-binding.v1":
        raise RuntimeError("SEMANTIC_RECONCILIATION_BINDING_INVALID")
    if reconciliation["approved_semantic_digest"] != APPROVED_SEMANTIC_SHA or reconciliation["approved_delta"] != binding["approved_delta"]:
        raise RuntimeError("SEMANTIC_RECONCILIATION_BINDING_MISMATCH")
    if (
        reconciliation["request_id"] != binding["semantic_reconciliation_request_id"]
        or reconciliation["request_head"] != binding["semantic_reconciliation_head"]
        or reconciliation["owner_review_id"] != binding["semantic_reconciliation_review_id"]
        or reconciliation["owner_review_head"] != binding["semantic_reconciliation_review_head"]
    ):
        raise RuntimeError("SEMANTIC_RECONCILIATION_BINDING_MISMATCH")
    _, reconciliation_pr = api("GET", f"/api/v1/repos/{REPO}/pulls/{reconciliation['pr_number']}", value)
    if (
        reconciliation_pr.get("state") != "open"
        or reconciliation_pr.get("merged")
        or reconciliation_pr.get("head", {}).get("sha") != reconciliation["request_head"]
        or reconciliation_pr.get("base", {}).get("ref") != "governance/approvals"
    ):
        raise RuntimeError("SEMANTIC_RECONCILIATION_REQUEST_STALE")
    _, reconciliation_reviews = api("GET", f"/api/v1/repos/{REPO}/pulls/{reconciliation['pr_number']}/reviews", value)
    reconciliation_review = next((x for x in reconciliation_reviews if x.get("id") == reconciliation["owner_review_id"]), None)
    if (
        not reconciliation_review
        or reconciliation_review.get("user", {}).get("id") != 2
        or reconciliation_review.get("state") != "APPROVED"
        or reconciliation_review.get("commit_id") != reconciliation["owner_review_head"]
        or reconciliation_review.get("dismissed") is not False
    ):
        raise RuntimeError("SEMANTIC_RECONCILIATION_OWNER_REVIEW_STALE")
    _, pr = api("GET", f"/api/v1/repos/{REPO}/pulls/54", value)
    if pr.get("state") != "open" or pr.get("merged") or pr.get("head", {}).get("sha") != PR54_HEAD:
        raise RuntimeError("PR54_APPROVAL_STALE")
    _, reviews = api("GET", f"/api/v1/repos/{REPO}/pulls/54/reviews", value)
    review = next((x for x in reviews if x.get("id") == PR54_REVIEW), None)
    if not review or review.get("user", {}).get("id") != 2 or review.get("state") != "APPROVED" or review.get("commit_id") != PR54_HEAD or review.get("dismissed") is not False:
        raise RuntimeError("PR54_APPROVAL_STALE")
    manifest_path = "approvals/governance-policy/AIOS-AUTONOMOUS-MERGE-GOVERNANCE-SCOPE-EXTENSION-V1.json"
    _, content = api("GET", f"/api/v1/repos/{REPO}/contents/{urllib.parse.quote(manifest_path, safe='/')}?ref={PR54_HEAD}", value)
    raw = base64.b64decode("".join(content.get("content", "").split()))
    if hashlib.sha256(raw).hexdigest() != PR54_MANIFEST_SHA:
        raise RuntimeError("PR54_MANIFEST_STALE")


def get_protection(value, branch):
    if branch not in {"main", "governance/approvals"}:
        raise RuntimeError("BRANCH_DENIED")
    _, protection = api("GET", f"/api/v1/repos/{REPO}/branch_protections/{urllib.parse.quote(branch, safe='')}", value)
    return protection


def mutation_receipt(before, after, binding, target_semantic):
    return {
        "before": {"transport_digest": sha(before), "semantic_digest": semantic_sha(before)},
        "approved": {
            "mutation_predicate_digest": sha({"operation": "ADD_MERGE_EXECUTOR", "subject_user_id": 153, "preserve_all_unmodified_security_fields": True}),
            "Owner_review_id": binding["Owner_review_id"],
            "semantic_policy_digest": APPROVED_SEMANTIC_SHA,
        },
        "after": {"transport_digest": sha(after), "semantic_digest": semantic_sha(after)},
        "delta": {"add_merge_whitelist_user_id": 153},
        "target_semantic_digest": sha(target_semantic),
        "unmodified_security_fields": {"verified": True},
    }


def apply_stage1(binding, value):
    verify_binding(binding, value)
    before = get_protection(value, "main")
    before_semantic = semantic_policy(before)
    before_semantic_sha = sha(before_semantic)
    names = list(before.get("merge_whitelist_usernames") or [])
    name_ids = user_ids(names)
    if name_ids not in ([2], [2, 153]):
        raise RuntimeError("PROTECTION_WHITELIST_PRESTATE_INVALID")
    target_names = ["owner", "aios-autonomy-agent"]
    target_semantic = copy.deepcopy(before_semantic)
    target_semantic["merge_whitelist_user_ids"] = [2, 153]
    if name_ids == [2, 153]:
        if before_semantic_sha != STAGE1_SEMANTIC_SHA:
            journal("APPLY_BLOCKED_STAGE1_IDEMPOTENCY_SEMANTIC_CONFLICT", transport_digest=sha(before), semantic_digest=before_semantic_sha, expected_semantic_digest=STAGE1_SEMANTIC_SHA)
            raise RuntimeError("PROTECTION_STAGE1_SEMANTIC_CONFLICT")
        receipt = mutation_receipt(before, before, binding, target_semantic)
        journal("APPLY_APPROVED_PROTECTION_DELTA_IDEMPOTENT_SUCCESS", receipt=receipt)
        return {"status": "IDEMPOTENT_SUCCESS", "protection": before, "transport_digest": sha(before), "semantic_digest": before_semantic_sha, "target_semantic_digest": sha(target_semantic)}
    if before_semantic_sha != APPROVED_SEMANTIC_SHA:
        journal("APPLY_BLOCKED_SEMANTIC_CAS", transport_digest=sha(before), semantic_digest=before_semantic_sha, expected_semantic_digest=APPROVED_SEMANTIC_SHA)
        raise RuntimeError("PROTECTION_SEMANTIC_CAS_CONFLICT")
    after_api = api("PATCH", f"/api/v1/repos/{REPO}/branch_protections/main", value, {"merge_whitelist_usernames": target_names})[1]
    after = get_protection(value, "main")
    for key in before:
        if key in {"updated_at", "merge_whitelist_usernames"}:
            continue
        if after.get(key) != before.get(key):
            raise RuntimeError("PROTECTION_UNAPPROVED_FIELD_CHANGED:" + key)
    if user_ids(after.get("merge_whitelist_usernames")) != [2, 153]:
        raise RuntimeError("PROTECTION_TARGET_READBACK_FAILED")
    after_semantic = semantic_policy(after)
    if after_semantic != target_semantic:
        raise RuntimeError("PROTECTION_SEMANTIC_TARGET_READBACK_FAILED")
    receipt = mutation_receipt(before, after, binding, target_semantic)
    journal("APPLY_APPROVED_PROTECTION_DELTA_COMPLETED", receipt=receipt)
    return {"status": "APPLIED", "protection": after, "previous_transport_digest": sha(before), "previous_semantic_digest": before_semantic_sha, "transport_digest": sha(after), "semantic_digest": sha(after_semantic), "receipt": receipt}


def apply_stage2(binding, facts, value):
    verify_binding(binding, value)
    validate_stage2_facts(facts)
    before = get_protection(value, "main")
    if OWNER_GATE in (before.get("status_check_contexts") or []) and stage2_semantic_digest(before) != STAGE1_SEMANTIC_SHA:
        raise RuntimeError("STAGE2_SEMANTIC_PRECONDITION_INVALID")
    try:
        status, target = build_target(before)
    except Stage2GateError as exc:
        journal("MIGRATE_APPROVED_GATE_BLOCKED", reason=str(exc), semantic_digest=stage2_semantic_digest(before))
        raise RuntimeError(str(exc)) from exc
    target_digest = stage2_semantic_digest(target)
    if status == "IDEMPOTENT_SUCCESS":
        journal(
            "MIGRATE_APPROVED_GATE_IDEMPOTENT_SUCCESS",
            before_semantic_digest=stage2_semantic_digest(before),
            target_semantic_digest=target_digest,
            campaign_id=CAMPAIGN_ID,
        )
        return {
            "status": status,
            "protection": before,
            "transport_digest": sha(before),
            "semantic_digest": stage2_semantic_digest(before),
            "target_semantic_digest": target_digest,
            "gate_transition": {"removed": OWNER_GATE, "required": CAMPAIGN_GATE},
        }
    after_api = api(
        "PATCH",
        f"/api/v1/repos/{REPO}/branch_protections/main",
        value,
        {"status_check_contexts": target["status_check_contexts"]},
    )[1]
    after = get_protection(value, "main")
    for key in before:
        if key in {"updated_at", "status_check_contexts"}:
            continue
        if after.get(key) != before.get(key):
            raise RuntimeError("PROTECTION_UNAPPROVED_FIELD_CHANGED:" + key)
    if after.get("status_check_contexts") != target["status_check_contexts"]:
        raise RuntimeError("STAGE2_GATE_TARGET_READBACK_FAILED")
    if stage2_semantic_digest(after) != target_digest:
        raise RuntimeError("STAGE2_SEMANTIC_TARGET_READBACK_FAILED")
    receipt = {
        "before_transport_digest": sha(before),
        "before_semantic_digest": stage2_semantic_digest(before),
        "after_transport_digest": sha(after),
        "after_semantic_digest": stage2_semantic_digest(after),
        "target_semantic_digest": target_digest,
        "campaign_id": CAMPAIGN_ID,
        "gate_transition": {"removed": OWNER_GATE, "required": CAMPAIGN_GATE},
        "preserved_security_fields": True,
        "facts_digest": sha(facts),
    }
    journal("MIGRATE_APPROVED_GATE_COMPLETED", receipt=receipt)
    return {
        "status": "APPLIED",
        "protection": after,
        "previous_transport_digest": sha(before),
        "previous_semantic_digest": stage2_semantic_digest(before),
        "transport_digest": sha(after),
        "semantic_digest": stage2_semantic_digest(after),
        "target_semantic_digest": target_digest,
        "receipt": receipt,
    }


def handle(request, value):
    if not isinstance(request, dict) or set(request) != {"operation", "payload"}:
        raise RuntimeError("REQUEST_SCHEMA_INVALID")
    operation = request["operation"]
    payload = request["payload"]
    if not isinstance(payload, dict) or payload.get("repository") != REPO:
        raise RuntimeError("REPOSITORY_DENIED")
    if operation == "GET_BRANCH_PROTECTION":
        protection = get_protection(value, payload.get("branch"))
        journal("GET_BRANCH_PROTECTION", branch=payload.get("branch"), transport_digest=sha(protection), semantic_digest=semantic_sha(protection))
        return {"status": "PASS", "protection": protection, "snapshot_sha256": sha(protection), "transport_digest": sha(protection), "semantic_digest": semantic_sha(protection), "semantic_schema": SEMANTIC_SCHEMA}
    if operation == "READBACK_BRANCH_PROTECTION":
        protection = get_protection(value, payload.get("branch"))
        return {"status": "PASS", "protection": protection, "snapshot_sha256": sha(protection), "transport_digest": sha(protection), "semantic_digest": semantic_sha(protection), "semantic_schema": SEMANTIC_SCHEMA}
    if operation == "APPLY_APPROVED_PROTECTION_DELTA":
        binding = payload.get("GovernanceExecutionBinding")
        if not isinstance(binding, dict):
            raise RuntimeError("GOVERNANCE_BINDING_REQUIRED")
        return apply_stage1(binding, value)
    if operation == "MIGRATE_APPROVED_GATE":
        binding = payload.get("GovernanceExecutionBinding")
        facts = payload.get("Stage2Facts")
        if not isinstance(binding, dict) or not isinstance(facts, dict):
            raise RuntimeError("STAGE2_BINDING_OR_FACTS_REQUIRED")
        return apply_stage2(binding, facts, value)
    raise RuntimeError("OPERATION_DENIED")


def main():
    value = token()
    sock_path = pathlib.Path(SOCKET)
    sock_path.unlink(missing_ok=True)
    server = socket.socket(socket.AF_UNIX, socket.SOCK_STREAM)
    server.bind(SOCKET)
    os.chmod(SOCKET, 0o660)
    server.listen(16)
    journal("EXECUTOR_STARTED", api_surface=["GET_BRANCH_PROTECTION", "APPLY_APPROVED_PROTECTION_DELTA", "MIGRATE_APPROVED_GATE", "READBACK_BRANCH_PROTECTION"], semantic_schema=SEMANTIC_SCHEMA, approved_semantic_digest=APPROVED_SEMANTIC_SHA, stage1_semantic_digest=STAGE1_SEMANTIC_SHA, campaign_gate=CAMPAIGN_GATE)
    try:
        while True:
            conn, _ = server.accept()
            with conn:
                data = conn.recv(1024 * 1024 + 1)
                if len(data) > 1024 * 1024:
                    result = {"status": "FAIL", "code": "REQUEST_TOO_LARGE"}
                else:
                    try:
                        request = json.loads(data.decode())
                        result = handle(request, value)
                    except Exception as exc:
                        result = {"status": "FAIL", "code": str(exc)}
                conn.sendall((json.dumps(result, sort_keys=True, separators=(",", ":"), ensure_ascii=False) + "\n").encode())
    finally:
        server.close()
        sock_path.unlink(missing_ok=True)


if __name__ == "__main__":
    main()
