from __future__ import annotations

import base64
from concurrent.futures import ThreadPoolExecutor
import hashlib
import json
from pathlib import Path
import threading

from aios_signer_broker_service import (
    REQUEST_SCHEMA,
    RESPONSE_SCHEMA,
    SignerBrokerPolicy,
    SignerBrokerService,
    SignerBrokerServiceError,
)


class FakeRunner:
    def __init__(self, failure: str | None = None):
        self.failure = failure
        self.calls: list[tuple[str, str]] = []
        self.lock = threading.Lock()

    def self_test(self) -> None:
        with self.lock:
            self.calls.append(("self_test", ""))
        if self.failure:
            raise SignerBrokerServiceError(self.failure)

    def sign(self, payload_kind: str, payload: bytes) -> tuple[bytes, str]:
        with self.lock:
            self.calls.append(("sign", payload_kind))
        if self.failure:
            raise SignerBrokerServiceError(self.failure)
        return b"signature:" + payload, "test-broker-key"


def make_policy(tmp_path: Path, *, allowed_uids: frozenset[int] = frozenset({1000})) -> SignerBrokerPolicy:
    return SignerBrokerPolicy(
        socket_path=tmp_path / "signer.sock",
        signer_path=Path("/usr/local/libexec/aios-grant-signer-adapter"),
        audit_path=tmp_path / "audit.jsonl",
        state_root=tmp_path / "state",
        allowed_uids=allowed_uids,
    )


def sign_request(payload: bytes = b"canonical payload") -> dict[str, str]:
    return {
        "schema": REQUEST_SCHEMA,
        "operation": "sign",
        "payload_kind": "approval_plan",
        "payload_sha256": hashlib.sha256(payload).hexdigest(),
        "payload_base64": base64.b64encode(payload).decode("ascii"),
    }


def read_audit(path: Path) -> list[dict[str, object]]:
    return [json.loads(line) for line in path.read_text(encoding="utf-8").splitlines()]


def test_authorized_request_and_self_test_emit_metadata_only_audit(tmp_path: Path):
    service = SignerBrokerService(make_policy(tmp_path), FakeRunner())
    self_test = service.handle_request({"schema": REQUEST_SCHEMA, "operation": "self_test"}, 1000)
    signed = service.handle_request(sign_request(b"secret-shaped canonical input"), 1000)
    assert self_test == {"schema": RESPONSE_SCHEMA, "status": "PASS", "operation": "self_test"}
    assert signed["status"] == "PASS"
    assert base64.b64decode(signed["signature_base64"]) == b"signature:secret-shaped canonical input"
    assert signed["key_id"] == "test-broker-key"
    audit_text = (tmp_path / "audit.jsonl").read_text(encoding="utf-8")
    assert "secret-shaped canonical input" not in audit_text
    assert "payload_base64" not in audit_text
    allowed = {"schema", "timestamp", "peer_uid", "operation", "status", "error_code", "payload_sha256"}
    assert all(set(event) <= allowed for event in read_audit(tmp_path / "audit.jsonl"))


def test_unauthorized_caller_is_rejected_before_signer_use(tmp_path: Path):
    runner = FakeRunner()
    service = SignerBrokerService(make_policy(tmp_path), runner)
    result = service.handle_request(sign_request(), 1001)
    assert result["status"] == "FAIL"
    assert result["error_code"] == "CALLER_UNAUTHORIZED"
    assert runner.calls == []


def test_arbitrary_command_and_invalid_operation_are_rejected(tmp_path: Path):
    runner = FakeRunner()
    service = SignerBrokerService(make_policy(tmp_path), runner)
    injected = {"schema": REQUEST_SCHEMA, "operation": "sign", "command": ["/bin/sh", "-c", "id"]}
    invalid = {"schema": REQUEST_SCHEMA, "operation": "run", "command": ["id"]}
    assert service.handle_request(injected, 1000)["error_code"] == "REQUEST_FIELDS_REJECTED"
    assert service.handle_request(invalid, 1000)["error_code"] == "SIGNER_OPERATION_REJECTED"
    assert runner.calls == []


def test_invalid_payload_and_hash_fail_closed(tmp_path: Path):
    service = SignerBrokerService(make_policy(tmp_path), FakeRunner())
    invalid = sign_request()
    invalid["payload_base64"] = "not-base64"
    assert service.handle_request(invalid, 1000)["error_code"] == "SIGNER_PAYLOAD_INVALID"
    mismatch = sign_request()
    mismatch["payload_sha256"] = "0" * 64
    assert service.handle_request(mismatch, 1000)["error_code"] == "SIGNER_PAYLOAD_HASH_MISMATCH"


def test_signer_unavailable_and_timeout_are_stable_failures(tmp_path: Path):
    for code in ("SIGNER_UNAVAILABLE", "SIGNER_TIMEOUT"):
        result = SignerBrokerService(make_policy(tmp_path / code), FakeRunner(code)).handle_request(
            {"schema": REQUEST_SCHEMA, "operation": "self_test"}, 1000
        )
        assert result == {"schema": RESPONSE_SCHEMA, "status": "FAIL", "error_code": code}


def test_concurrent_calls_are_safe_and_audited(tmp_path: Path):
    service = SignerBrokerService(make_policy(tmp_path), FakeRunner())
    with ThreadPoolExecutor(max_workers=8) as pool:
        results = list(pool.map(lambda index: service.handle_request(sign_request(f"payload-{index}".encode()), 1000), range(24)))
    assert all(result["status"] == "PASS" for result in results)
    assert len(read_audit(tmp_path / "audit.jsonl")) == 24