"""Signed tenant context and canonical message contract for AIOS V2R1."""
from __future__ import annotations

import base64
from dataclasses import asdict, dataclass
from datetime import datetime, timedelta, timezone
import hashlib
import hmac
import json
import re
import secrets
from typing import Any, Mapping

VERSION = "V2R1C0"
MAX_SAFE_INTEGER = 9_007_199_254_740_991
TENANT_RE = re.compile(r"^[a-z0-9](?:[a-z0-9-]{0,38}[a-z0-9])?$")
IDENT_RE = re.compile(r"^[A-Za-z0-9](?:[A-Za-z0-9._:@-]{0,126}[A-Za-z0-9])?$")


class V2R1Error(RuntimeError):
    pass


class ValidationError(V2R1Error):
    pass


class ScopeError(V2R1Error):
    pass


class ConflictError(V2R1Error):
    pass


def utc_now() -> datetime:
    return datetime.now(timezone.utc)


def isoformat(value: datetime | None = None) -> str:
    return (value or utc_now()).astimezone(timezone.utc).isoformat()


def parse_time(value: str) -> datetime:
    if not isinstance(value, str) or not value.strip():
        raise ValidationError("timestamp must be a non-empty string")
    text = value.strip()
    if text.endswith("Z"):
        text = text[:-1] + "+00:00"
    try:
        parsed = datetime.fromisoformat(text)
    except ValueError as exc:
        raise ValidationError(f"invalid timestamp: {value!r}") from exc
    if parsed.tzinfo is None:
        raise ValidationError("timestamp must include timezone")
    return parsed.astimezone(timezone.utc)


def require_ident(name: str, value: Any, *, tenant: bool = False) -> str:
    if not isinstance(value, str):
        raise ValidationError(f"{name} must be a string")
    value = value.strip()
    pattern = TENANT_RE if tenant else IDENT_RE
    if not pattern.fullmatch(value):
        raise ValidationError(f"invalid {name}: {value!r}")
    return value


def normalize_message_id(value: Any, *, field_name: str = "message_id") -> str:
    if isinstance(value, bool):
        raise ValidationError(f"{field_name} cannot be boolean")
    if isinstance(value, str):
        value = value.strip()
        if not value or len(value) > 256 or any(ord(ch) < 32 for ch in value):
            raise ValidationError(f"invalid {field_name}")
        return value
    if isinstance(value, int):
        if abs(value) > MAX_SAFE_INTEGER:
            raise ValidationError(f"{field_name} is an unsafe numeric ID")
        return str(value)
    raise ValidationError(f"{field_name} must be a lossless string or safe integer")


def canonical_json(value: Any) -> bytes:
    return json.dumps(
        value, ensure_ascii=False, separators=(",", ":"), sort_keys=True
    ).encode("utf-8")


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


def _unb64(value: str) -> bytes:
    return base64.urlsafe_b64decode(value + "=" * (-len(value) % 4))


@dataclass(frozen=True)
class TenantContext:
    tenant_id: str
    cell_id: str
    principal_id: str
    principal_type: str
    session_id: str
    task_id: str | None
    channel: str
    channel_account_id: str
    issued_at: str
    expires_at: str
    nonce: str

    def __post_init__(self) -> None:
        require_ident("tenant_id", self.tenant_id, tenant=True)
        for name in (
            "cell_id", "principal_id", "principal_type", "session_id",
            "channel", "channel_account_id", "nonce",
        ):
            require_ident(name, getattr(self, name))
        if self.task_id is not None:
            require_ident("task_id", self.task_id)
        issued = parse_time(self.issued_at)
        expires = parse_time(self.expires_at)
        if expires <= issued or expires - issued > timedelta(hours=1):
            raise ValidationError("invalid tenant context lifetime")

    @classmethod
    def issue(cls, *, lifetime_seconds: int = 300, now: datetime | None = None, **scope: Any) -> "TenantContext":
        if not 1 <= lifetime_seconds <= 3600:
            raise ValidationError("lifetime_seconds must be between 1 and 3600")
        current = (now or utc_now()).astimezone(timezone.utc)
        return cls(
            **scope,
            issued_at=isoformat(current),
            expires_at=isoformat(current + timedelta(seconds=lifetime_seconds)),
            nonce=secrets.token_hex(16),
        )

    def sign(self, key: bytes) -> str:
        if len(key) < 32:
            raise ValidationError("signing key must be at least 32 bytes")
        payload = _b64(canonical_json(asdict(self)))
        signature = _b64(hmac.new(key, payload.encode("ascii"), hashlib.sha256).digest())
        return f"{payload}.{signature}"

    @classmethod
    def verify(
        cls,
        token: str,
        key: bytes,
        *,
        now: datetime | None = None,
        expected_tenant_id: str | None = None,
        expected_cell_id: str | None = None,
    ) -> "TenantContext":
        if len(key) < 32:
            raise ValidationError("signing key must be at least 32 bytes")
        try:
            payload, encoded_signature = token.split(".", 1)
            actual = _unb64(encoded_signature)
            if encoded_signature != _b64(actual):
                raise ValidationError("non-canonical tenant context signature")
            data = json.loads(_unb64(payload))
        except Exception as exc:
            raise ValidationError("malformed tenant context token") from exc
        expected = hmac.new(key, payload.encode("ascii"), hashlib.sha256).digest()
        if not hmac.compare_digest(actual, expected):
            raise ValidationError("tenant context signature mismatch")
        context = cls(**data)
        current = (now or utc_now()).astimezone(timezone.utc)
        if current < parse_time(context.issued_at) - timedelta(seconds=60):
            raise ValidationError("tenant context issued in future")
        if current >= parse_time(context.expires_at):
            raise ValidationError("tenant context expired")
        if expected_tenant_id and context.tenant_id != expected_tenant_id:
            raise ScopeError("tenant mismatch")
        if expected_cell_id and context.cell_id != expected_cell_id:
            raise ScopeError("cell mismatch")
        return context


@dataclass(frozen=True)
class CanonicalMessage:
    tenant_id: str
    cell_id: str
    principal_id: str
    principal_type: str
    origin_type: str
    channel: str
    channel_account_id: str
    peer_id: str
    conversation_id: str
    session_id: str
    task_id: str | None
    event_id: str
    message_id: str
    reply_to_message_id: str | None
    direction: str
    command_capable: bool
    occurred_at: str
    received_at: str
    body: str
    quoted_body: str | None
    adapter_name: str
    adapter_version: str
    schema_version: str = "aios.message.v1"

    def __post_init__(self) -> None:
        if self.schema_version != "aios.message.v1":
            raise ValidationError("unsupported message schema version")
        require_ident("tenant_id", self.tenant_id, tenant=True)
        for name in (
            "cell_id", "principal_id", "principal_type", "origin_type", "channel",
            "channel_account_id", "peer_id", "conversation_id", "session_id",
            "event_id", "adapter_name", "adapter_version",
        ):
            require_ident(name, getattr(self, name))
        if self.task_id is not None:
            require_ident("task_id", self.task_id)
        normalize_message_id(self.message_id)
        if self.reply_to_message_id is not None:
            normalize_message_id(self.reply_to_message_id, field_name="reply_to_message_id")
        if self.direction not in {"inbound", "outbound", "internal"}:
            raise ValidationError("invalid direction")
        if not isinstance(self.command_capable, bool):
            raise ValidationError("command_capable must be boolean")
        parse_time(self.occurred_at)
        parse_time(self.received_at)
        if not isinstance(self.body, str):
            raise ValidationError("body must be a string")
        if self.quoted_body is not None and not isinstance(self.quoted_body, str):
            raise ValidationError("quoted_body must be a string or null")

    @classmethod
    def from_mapping(cls, data: Mapping[str, Any]) -> "CanonicalMessage":
        content = data.get("content") or {}
        adapter = data.get("adapter") or {}
        if not isinstance(content, Mapping) or not isinstance(adapter, Mapping):
            raise ValidationError("content and adapter must be objects")
        values = dict(data)
        return cls(
            schema_version=values.get("schema_version", "aios.message.v1"),
            tenant_id=values["tenant_id"],
            cell_id=values["cell_id"],
            principal_id=values["principal_id"],
            principal_type=values["principal_type"],
            origin_type=values["origin_type"],
            channel=values["channel"],
            channel_account_id=values["channel_account_id"],
            peer_id=values["peer_id"],
            conversation_id=values["conversation_id"],
            session_id=values["session_id"],
            task_id=values.get("task_id"),
            event_id=values["event_id"],
            message_id=normalize_message_id(values["message_id"]),
            reply_to_message_id=(
                normalize_message_id(values["reply_to_message_id"], field_name="reply_to_message_id")
                if values.get("reply_to_message_id") is not None else None
            ),
            direction=values["direction"],
            command_capable=values["command_capable"],
            occurred_at=values["occurred_at"],
            received_at=values["received_at"],
            body=content.get("body", ""),
            quoted_body=content.get("quoted_body"),
            adapter_name=adapter["name"],
            adapter_version=adapter["version"],
        )

    def assert_context(self, context: TenantContext) -> None:
        fields = (
            "tenant_id", "cell_id", "principal_id", "session_id", "task_id",
            "channel", "channel_account_id",
        )
        mismatches = [name for name in fields if getattr(self, name) != getattr(context, name)]
        if mismatches:
            raise ScopeError("message/context mismatch: " + ", ".join(mismatches))

    @property
    def is_authenticated_human_command(self) -> bool:
        return (
            self.direction == "inbound"
            and self.command_capable
            and self.principal_type == "authenticated_user"
            and self.origin_type == "channel_user"
        )


@dataclass(frozen=True)
class PolicyDecision:
    allowed: bool
    approval_required: bool
    code: str
    reason: str


READ_ONLY_TOOLS = frozenset({
    "read", "grep", "glob", "find", "ls", "web_search", "web_fetch",
    "memory_search", "sessions_list", "sessions_history", "session_status",
})


def evaluate_tool_call(tool_name: str, message: CanonicalMessage) -> PolicyDecision:
    normalized = tool_name.strip().lower() if isinstance(tool_name, str) else ""
    if not normalized:
        return PolicyDecision(False, False, "invalid_tool", "tool name missing")
    if normalized in READ_ONLY_TOOLS:
        return PolicyDecision(True, False, "read_only", "read-only operation")
    if not message.is_authenticated_human_command:
        return PolicyDecision(
            False, False, "origin_denied",
            "side effects require an authenticated current human command",
        )
    return PolicyDecision(
        True, True, "approval_required",
        "authenticated human side effect requires explicit approval",
    )