"""Multi-tenant isolation layer -- stdlib only.

Provides:
- TenantContext: per-request tenant identification and scoping
- TenantStore: CRUD for tenant metadata
- TenantMiddleware: request-scoped tenant resolution from auth identity
- TenantDAGProxy: tenant-isolated WorldDAG access

All stdlib -- no external dependencies.
"""

from __future__ import annotations

import threading
import time
from contextvars import ContextVar
from dataclasses import dataclass, field
from typing import Any

from orxaq.kernel.types import CausalEdge, Entity
from orxaq.kernel.world_dag import WorldDAG

# ── Tenant Context ───────────────────────────────────────────

_current_tenant: ContextVar[TenantContext | None] = ContextVar("current_tenant", default=None)


@dataclass(frozen=True)
class TenantContext:
    """Immutable context identifying the current tenant."""

    tenant_id: str
    name: str = ""
    tier: str = "free"  # free | pro | enterprise
    max_entities: int = 100
    max_edges: int = 500
    metadata: dict[str, Any] = field(default_factory=dict)

    @property
    def is_unlimited(self) -> bool:
        """Enterprise tenants bypass quotas."""
        return self.tier == "enterprise"


def get_current_tenant() -> TenantContext | None:
    """Return the tenant context for the current execution scope."""
    return _current_tenant.get()


def set_current_tenant(ctx: TenantContext | None) -> None:
    """Set the tenant context for the current execution scope."""
    _current_tenant.set(ctx)


# ── Tenant Store ─────────────────────────────────────────────


class TenantStore:
    """Thread-safe in-memory store for tenant metadata.

    Production deployments would back this with a database.
    """

    def __init__(self) -> None:
        self._tenants: dict[str, TenantContext] = {}
        self._lock = threading.Lock()

    def create(
        self,
        tenant_id: str,
        name: str = "",
        tier: str = "free",
        max_entities: int = 100,
        max_edges: int = 500,
        metadata: dict[str, Any] | None = None,
    ) -> TenantContext:
        """Create a new tenant. Raises ValueError if ID already exists."""
        with self._lock:
            if tenant_id in self._tenants:
                msg = f"Tenant '{tenant_id}' already exists"
                raise ValueError(msg)
            ctx = TenantContext(
                tenant_id=tenant_id,
                name=name or tenant_id,
                tier=tier,
                max_entities=max_entities,
                max_edges=max_edges,
                metadata=metadata or {},
            )
            self._tenants[tenant_id] = ctx
        return ctx

    def get(self, tenant_id: str) -> TenantContext | None:
        """Retrieve tenant context by ID."""
        with self._lock:
            return self._tenants.get(tenant_id)

    def delete(self, tenant_id: str) -> bool:
        """Remove a tenant; return whether it existed."""
        with self._lock:
            return self._tenants.pop(tenant_id, None) is not None

    def list_tenants(self) -> list[TenantContext]:
        """Return all tenants sorted by ID."""
        with self._lock:
            return sorted(self._tenants.values(), key=lambda t: t.tenant_id)

    def count(self) -> int:
        """Number of registered tenants."""
        with self._lock:
            return len(self._tenants)


# ── Tenant DAG Proxy ─────────────────────────────────────────


class QuotaExceededError(Exception):
    """Raised when a tenant exceeds their entity or edge quota."""


class TenantDAGProxy:
    """Tenant-isolated view of a WorldDAG.

    Each tenant gets its own prefix-namespaced slice. Entities are
    stored as ``{tenant_id}:{entity_id}`` in the underlying DAG.
    """

    def __init__(self, dag: WorldDAG, tenant: TenantContext) -> None:
        self._dag = dag
        self._tenant = tenant
        self._prefix = f"{tenant.tenant_id}:"

    @property
    def tenant(self) -> TenantContext:
        return self._tenant

    def _scoped_id(self, entity_id: str) -> str:
        """Prefix entity_id with tenant namespace."""
        if entity_id.startswith(self._prefix):
            return entity_id
        return f"{self._prefix}{entity_id}"

    def _unscoped_id(self, scoped_id: str) -> str:
        """Strip tenant prefix from entity_id."""
        if scoped_id.startswith(self._prefix):
            return scoped_id[len(self._prefix) :]
        return scoped_id

    def add_entity(self, entity: Entity) -> None:
        """Add entity under tenant namespace with quota check."""
        if not self._tenant.is_unlimited:
            if self.entity_count() >= self._tenant.max_entities:
                msg = (
                    f"Tenant '{self._tenant.tenant_id}' entity quota "
                    f"exceeded ({self._tenant.max_entities})"
                )
                raise QuotaExceededError(msg)
        scoped = Entity(entity_id=self._scoped_id(entity.entity_id), name=entity.name)
        self._dag.add_entity(scoped)

    def get_entity(self, entity_id: str) -> Entity | None:
        """Retrieve entity, returning unscoped IDs."""
        scoped_id = self._scoped_id(entity_id)
        ent = self._dag.get_entity(scoped_id)
        if ent is None:
            return None
        return Entity(entity_id=self._unscoped_id(ent.entity_id), name=ent.name)

    def remove_entity(self, entity_id: str) -> None:
        """Remove entity from tenant's namespace."""
        self._dag.remove_entity(self._scoped_id(entity_id))

    def add_edge(self, edge: CausalEdge) -> None:
        """Add edge between tenant-scoped entities with quota check."""
        if not self._tenant.is_unlimited:
            if self.edge_count() >= self._tenant.max_edges:
                msg = (
                    f"Tenant '{self._tenant.tenant_id}' edge quota "
                    f"exceeded ({self._tenant.max_edges})"
                )
                raise QuotaExceededError(msg)
        scoped_edge = CausalEdge(
            source=self._scoped_id(edge.source),
            target=self._scoped_id(edge.target),
            mechanism=edge.mechanism,
            confidence=edge.confidence,
        )
        self._dag.add_edge(scoped_edge)

    def entity_count(self) -> int:
        """Count entities belonging to this tenant."""
        count = 0
        for eid in self._dag.entity_ids():
            if eid.startswith(self._prefix):
                count += 1
        return count

    def edge_count(self) -> int:
        """Count edges belonging to this tenant."""
        count = 0
        for eid in self._dag.edge_ids():
            # edge_id format: "source→target"; tenant-scoped sources start with prefix
            if self._prefix in eid:
                count += 1
        return count

    def list_entities(self) -> list[Entity]:
        """List all entities for this tenant with unscoped IDs."""
        result = []
        for eid in self._dag.entity_ids():
            if eid.startswith(self._prefix):
                ent = self._dag.get_entity(eid)
                if ent is not None:
                    result.append(
                        Entity(
                            entity_id=self._unscoped_id(ent.entity_id),
                            name=ent.name,
                        )
                    )
        return result


# ── Tenant Middleware ────────────────────────────────────────


@dataclass(frozen=True)
class TenantResolution:
    """Result of tenant resolution from an auth identity."""

    resolved: bool
    tenant: TenantContext | None = None
    error: str = ""


class TenantMiddleware:
    """Resolves tenant context from authenticated identity.

    Maps auth identities (API key IDs or JWT subjects) to tenants
    via a configurable mapping. Sets the contextvars-based
    TenantContext for the request scope.
    """

    def __init__(self, store: TenantStore) -> None:
        self._store = store
        self._identity_map: dict[str, str] = {}  # identity -> tenant_id
        self._lock = threading.Lock()

    def map_identity(self, identity: str, tenant_id: str) -> None:
        """Associate an auth identity with a tenant."""
        with self._lock:
            self._identity_map[identity] = tenant_id

    def unmap_identity(self, identity: str) -> bool:
        """Remove identity-to-tenant mapping."""
        with self._lock:
            return self._identity_map.pop(identity, None) is not None

    def resolve(self, identity: str) -> TenantResolution:
        """Resolve authenticated identity to a TenantContext.

        Sets the thread-local context var on success.
        """
        with self._lock:
            tenant_id = self._identity_map.get(identity)

        if tenant_id is None:
            return TenantResolution(resolved=False, error="no_tenant_mapping")

        tenant = self._store.get(tenant_id)
        if tenant is None:
            return TenantResolution(resolved=False, error="tenant_not_found")

        set_current_tenant(tenant)
        return TenantResolution(resolved=True, tenant=tenant)

    def clear(self) -> None:
        """Clear the current tenant context."""
        set_current_tenant(None)


# ── Tenant Audit Wrapper ────────────────────────────────────


@dataclass(frozen=True)
class TenantAuditEntry:
    """Audit log entry scoped to a tenant."""

    tenant_id: str
    action: str
    entity_id: str = ""
    detail: str = ""
    timestamp: float = 0.0

    def __post_init__(self) -> None:
        if self.timestamp == 0.0:
            # Workaround for frozen dataclass: use object.__setattr__.
            object.__setattr__(self, "timestamp", time.time())


class TenantAuditLog:
    """Per-tenant audit trail with configurable retention."""

    def __init__(self, max_entries: int = 10000) -> None:
        self._entries: dict[str, list[TenantAuditEntry]] = {}
        self._max_entries = max_entries
        self._lock = threading.Lock()

    def record(self, entry: TenantAuditEntry) -> None:
        """Record an audit entry under its tenant."""
        with self._lock:
            tenant_log = self._entries.setdefault(entry.tenant_id, [])
            tenant_log.append(entry)
            if len(tenant_log) > self._max_entries:
                self._entries[entry.tenant_id] = tenant_log[-self._max_entries :]

    def entries_for(self, tenant_id: str, last_n: int = 50) -> list[TenantAuditEntry]:
        """Retrieve recent entries for a tenant."""
        with self._lock:
            tenant_log = self._entries.get(tenant_id, [])
            return tenant_log[-last_n:]

    def count_for(self, tenant_id: str) -> int:
        """Number of audit entries for a tenant."""
        with self._lock:
            return len(self._entries.get(tenant_id, []))

    def purge(self, tenant_id: str) -> int:
        """Remove all entries for a tenant; return count removed."""
        with self._lock:
            entries = self._entries.pop(tenant_id, [])
            return len(entries)
