from __future__ import annotations

import contextvars
import logging
import time
import uuid
from dataclasses import dataclass
from typing import Any, Iterable, Mapping, cast

from keiro.gateway.infra.billing.cost import PricingDimensions, default_pricing
from keiro.gateway.infra.concurrency import (
    BaseConcurrencyLimiter,
    InMemoryConcurrencyLimiter,
    RedisConcurrencyLimiter,
)
from keiro.gateway.infra.model_alias import _ALIASES as LEGACY_ALIASES
from keiro.gateway.infra.model_alias import ModelAlias
from keiro.gateway.infra.models import new_request_id
from keiro.gateway.infra.pricing.plan_cache import OrgPlanCache
from keiro.gateway.infra.provider_pool import ProviderPool
from keiro.gateway.infra.provider_registry import (
    get_capability_map,
    get_providers_for_model,
    get_routing_preferences,
)
from keiro.gateway.infra.providers import ProviderAdapter, default_adapters
from keiro.gateway.infra.rate_limit import RateLimiter as GatewayRateLimiter
from keiro.gateway.infra.rbc import jws_sign_hs256
from keiro.gateway.infra.routing import RouterLiteAdapter
from keiro.gateway.infra.secrets.store import InMemorySecretStore
from keiro.gateway.infra.stores import (
    BillingStore as LegacyBillingStore,
)
from keiro.gateway.infra.stores import (
    IdempotencyStore as LegacyIdempotencyStore,
)
from keiro.gateway.infra.stores import (
    InMemoryBillingStore,
    InMemoryIdempotencyStore,
    PostgresBillingStore,
    RedisIdempotencyStore,
)
from keiro.gateway.infra.streaming.replay import (
    InMemoryReplayStore,
    RedisReplayStore,
)
from keiro.gateway.infra.streaming.replay import (
    ReplayStore as LegacyReplayStore,
)
from keiro.gateway.infra.tokenizer import sum_message_tokens
from keiro.gateway.infra.tpm_limiter import DualTPMLimiter, InMemoryTPMLimiter, RedisTPMLimiter

from ..adapters.sur import FileSurStore
from ..pipeline.container import Services, Settings
from ..pipeline.executor import Executor, ProviderInvoker
from ..pipeline.types import OrgLimits
from ..boot.settings import GatewaySettings
from ..domain.admission import AdmissionController, RbcSigner
from ..domain.alias import AliasCatalog
from ..domain.errors import NoCapacityError, RateLimitedError
from ..domain.model import Message
from ..ports.auth import AuthContext
from ..ports.billing import FBR, BillingStore
from ..ports.creds import CredentialPool
from ..ports.idem import IdempotencyStore
from ..ports.plans import OrgPlanStore
from ..ports.pricing import PricingEngine
from ..ports.providers import ProviderRegistry
from ..ports.replay import ReplayStore
from ..ports.sur import SurStore
from ..ports.time_uuid import Clock, IdFactory

LegacyAliasMap = Mapping[str, ModelAlias]
LEGACY_ALIAS_MAP: LegacyAliasMap = cast(LegacyAliasMap, LEGACY_ALIASES)

LOGGER = logging.getLogger(__name__)


class RuntimeConfigError(RuntimeError):
    """Raised when gateway runtime dependencies are misconfigured."""


@dataclass(frozen=True)
class RuntimeOptions:
    """Toggle runtime behaviours that differ between production and tests."""

    allow_in_memory_fallback: bool = False
    enable_judge_stream: bool = False


def _load_settings(options: RuntimeOptions) -> GatewaySettings | None:
    try:
        from ..boot.settings import get_gateway_settings

        return get_gateway_settings()
    except Exception as exc:
        if options.allow_in_memory_fallback:
            LOGGER.debug("Falling back to in-memory stores: settings unavailable", exc_info=exc)
            return None
        raise RuntimeConfigError(
            "Gateway configuration is required for gateway runtime bootstrap"
        ) from exc


def _select_billing_store(
    settings: GatewaySettings | None,
    *,
    options: RuntimeOptions,
) -> LegacyBillingStore:
    if settings is None:
        if options.allow_in_memory_fallback:
            return InMemoryBillingStore()
        raise RuntimeConfigError("Postgres DSN must be configured for billing.")
    store_settings = settings.stores
    if not store_settings.postgres_dsn:
        if options.allow_in_memory_fallback:
            return InMemoryBillingStore()
        raise RuntimeConfigError("Postgres DSN must be configured for billing.")
    return PostgresBillingStore(
        store_settings.postgres_dsn,
        auto_init=store_settings.auto_init,
    )


def _select_idempotency_store(
    settings: GatewaySettings | None,
    *,
    options: RuntimeOptions,
) -> LegacyIdempotencyStore:
    if settings is None:
        if options.allow_in_memory_fallback:
            return InMemoryIdempotencyStore()
        raise RuntimeConfigError("Redis URL must be configured for idempotency tracking.")
    store_settings = settings.stores
    if not store_settings.redis_url:
        if options.allow_in_memory_fallback:
            return InMemoryIdempotencyStore()
        raise RuntimeConfigError("Redis URL must be configured for idempotency tracking.")
    import redis  # type: ignore

    client = redis.Redis.from_url(store_settings.redis_url)
    return RedisIdempotencyStore(client)


def _default_tpm_limit(settings: GatewaySettings | None) -> int | None:
    if settings is None:
        return None
    limit_settings = settings.limits()
    return limit_settings.default_tpm


def _select_tpm_limiter(
    settings: GatewaySettings | None,
    *,
    options: RuntimeOptions,
) -> DualTPMLimiter:
    default_limit = _default_tpm_limit(settings)
    if settings is None:
        if not options.allow_in_memory_fallback:
            raise RuntimeConfigError("Redis URL must be configured for TPM enforcement.")
        return DualTPMLimiter(
            InMemoryTPMLimiter(default_limit),
            InMemoryTPMLimiter(default_limit),
        )
    store_settings = settings.stores
    if not store_settings.redis_url:
        if not options.allow_in_memory_fallback:
            raise RuntimeConfigError("Redis URL must be configured for TPM enforcement.")
        return DualTPMLimiter(
            InMemoryTPMLimiter(default_limit),
            InMemoryTPMLimiter(default_limit),
        )
    import redis  # type: ignore

    client = redis.Redis.from_url(store_settings.redis_url)
    tpm_in = RedisTPMLimiter(client, prefix="itpm", default_limit=default_limit)
    tpm_out = RedisTPMLimiter(client, prefix="otpm", default_limit=default_limit)
    return DualTPMLimiter(tpm_in, tpm_out)


def _select_concurrency_limiter(
    settings: GatewaySettings | None,
    *,
    options: RuntimeOptions,
) -> BaseConcurrencyLimiter:
    if settings is None:
        if options.allow_in_memory_fallback:
            return InMemoryConcurrencyLimiter()
        raise RuntimeConfigError("Redis URL must be configured for concurrency enforcement.")
    store_settings = settings.stores
    if not store_settings.redis_url:
        if options.allow_in_memory_fallback:
            return InMemoryConcurrencyLimiter()
        raise RuntimeConfigError("Redis URL must be configured for concurrency enforcement.")
    import redis  # type: ignore

    client = redis.Redis.from_url(store_settings.redis_url)
    return RedisConcurrencyLimiter(client)


def _select_replay_store(
    settings: GatewaySettings | None,
    *,
    options: RuntimeOptions,
) -> InMemoryReplayStore | RedisReplayStore:
    """Select replay store implementation based on available infrastructure.

    Uses Redis when available for cross-node resume support; falls back to
    in-memory when Redis is unavailable and allow_in_memory_fallback is True.
    """
    if settings is None:
        if options.allow_in_memory_fallback:
            LOGGER.warning(
                "Using in-memory replay store; stream resume will not work across nodes"
            )
            return InMemoryReplayStore()
        raise RuntimeConfigError("Redis URL must be configured for durable stream replay.")
    store_settings = settings.stores
    if not store_settings.redis_url:
        if options.allow_in_memory_fallback:
            LOGGER.warning(
                "Using in-memory replay store; stream resume will not work across nodes"
            )
            return InMemoryReplayStore()
        raise RuntimeConfigError("Redis URL must be configured for durable stream replay.")
    import redis  # type: ignore

    client = redis.Redis.from_url(store_settings.redis_url)
    return RedisReplayStore(client)


def _default_plan_fetch_fn(org_id: str) -> dict[str, object] | None:
    """Stub fetch function that logs a warning about missing plan configuration.

    This function serves as a placeholder until a proper control-plane backed
    plan store is implemented. See P0 Gap #3 in gateway2_production_readiness.md.

    Args:
        org_id: Organization identifier.

    Returns:
        None, causing OrgPlanCache to use default tier values.
    """
    # TODO(P0): Replace with control-plane API call to fetch org billing config.
    # Currently all orgs get the same default tier, violating multi-tenancy.
    LOGGER.debug(
        "Plan fetch for org %r returning defaults (control-plane integration pending)",
        org_id,
    )
    return None


class _Signer(RbcSigner):
    """RBC signer that delegates to the legacy HS256 helper."""

    def sign(self, claims: Mapping[str, Any]) -> str:
        """Generate a signed RBC token.

        Args:
            claims: JWT-compatible payload to sign.

        Returns:
            Serialized token string.
        """
        return jws_sign_hs256(dict(claims))


class _ProviderRegistry(ProviderRegistry):
    """Adapter that exposes legacy provider metadata through the v2 ports."""

    def providers_for(self, canonical_model: str) -> list[str]:
        """Return registered providers that advertise the canonical model."""
        return get_providers_for_model(canonical_model)

    def capabilities(self) -> dict[str, dict[str, bool]]:
        """Return provider capability flags keyed by provider name."""
        return get_capability_map()

    def routing_preferences(self) -> dict[str, float]:
        """Return provider routing preference scores."""
        return get_routing_preferences()


class _CredentialPool(CredentialPool):
    """Credential pool adapter over the legacy ProviderPool.

    WARNING: The current implementation does NOT use org_id for credential
    isolation. All organizations share the same reseller credential pool.
    This is a known P0 gap documented in docs/design/gateway2_production_readiness.md.

    To properly support BYOK (Bring Your Own Key) and per-tenant isolation:
    1. ProviderPool must be extended to resolve credentials by {org_id, provider}
    2. A MultiTenantCredentialStore should replace this adapter
    3. See P0 Gap #2 "Credential isolation & BYOK" in the design doc
    """

    def __init__(self, pool: ProviderPool) -> None:
        self._pool = pool

    def acquire(self, org_id: str, provider: str) -> str | None:
        """Return an API key for the requested provider.

        Args:
            org_id: Organization identifier requesting the credential.
                WARNING: Currently ignored - all orgs share the reseller pool.
            provider: Provider slug.

        Returns:
            API key string, or None if no credential is available.
        """
        # TODO(P0): Implement per-org credential lookup. Currently org_id is
        # ignored and all requests use shared reseller credentials.
        # See: docs/design/gateway2_production_readiness.md "Credential isolation"
        del org_id  # Explicitly mark as unused until BYOK support is added.
        return self._pool.acquire(provider)

    def release(self, org_id: str, provider: str, credential: str | None) -> None:
        """Release a credential back to the pool.

        Args:
            org_id: Organization identifier.
            provider: Provider slug.
            credential: The credential to release.
        """
        # ProviderPool does not currently support release; credentials are reused.
        pass


class _OrgPlanStore(OrgPlanStore):
    """Plan store backed by the legacy gateway plan cache."""

    def __init__(self, cache: OrgPlanCache) -> None:
        self._cache = cache

    def get(self, org_id: str, *, fallback_tier: str) -> dict[str, object]:
        """Return the resolved plan for the given organization."""
        return self._cache.get(org_id, fallback_tier=fallback_tier)

    def billing_mode(self, plan: dict[str, object], provider: str) -> str:
        """Resolve the billing mode for a provider/plan combination."""
        default_mode = str(plan.get("default_billing_mode", "resell"))
        overrides = plan.get("provider_overrides")
        if isinstance(overrides, Mapping):
            candidate = overrides.get(provider)
            if isinstance(candidate, str) and candidate:
                return candidate
        return default_mode


class _Pricing(PricingEngine):
    """Pricing facade that reuses the legacy cost calculator."""

    def __init__(self) -> None:
        self._impl = default_pricing()

    def calculate(
        self,
        *,
        model: str,
        prompt_tokens: int,
        completion_tokens: int,
        price_book_version: str,
        billing_mode: str,
    ) -> float:
        """Compute the cost for a provider invocation."""
        dims = PricingDimensions(
            model=model,
            prompt_tokens=prompt_tokens,
            completion_tokens=completion_tokens,
            price_book_version=price_book_version,
            billing_mode=billing_mode,
        )
        return float(self._impl.calculate(dims))


class _BillingStore(BillingStore):
    """Shim exposing the legacy billing store through the v2 interface."""

    def __init__(self, store: LegacyBillingStore) -> None:
        self._store = store

    @property
    def underlying(self) -> LegacyBillingStore:
        """Return the wrapped legacy billing store for admin route wiring."""
        return self._store

    def seed(
        self,
        fbr_id: str,
        org_id: str,
        idem_key: str,
        price_book_version: str,
        *,
        billing_mode: str,
        api_key_id: str | None,
    ) -> FBR:
        """Create or fetch a billing record."""
        return self._store.seed_fbr(
            fbr_id,
            org_id,
            idem_key,
            price_book_version,
            billing_mode=billing_mode,
            api_key_id=api_key_id,
        )

    def record_delta(self, fbr_id: str, prompt: int, completion: int, cost_usd: float) -> None:
        """Record incremental usage for the active billing record."""
        self._store.record_usage_delta(fbr_id, prompt, completion, cost_usd)

    def finalize(self, fbr_id: str) -> FBR | None:
        """Finalize the billing record once usage accounting completes."""
        return self._store.finalize(fbr_id)

    def record_usage_event(
        self,
        fbr_id: str,
        event_seq: int,
        prompt_tokens: int,
        completion_tokens: int,
        cost_usd: float,
    ) -> None:
        """Store a streaming usage event for later reconciliation."""
        self._store.record_usage_event(
            fbr_id,
            event_seq,
            prompt_tokens,
            completion_tokens,
            cost_usd,
        )


class _Idempotency(IdempotencyStore):
    """Idempotency wrapper backed by the legacy store implementations."""

    def __init__(self, store: LegacyIdempotencyStore) -> None:
        self._store = store

    def claim(self, org_id: str, idem_key: str, request_id: str) -> bool:
        """Attempt to claim an idempotency key for the request."""
        return self._store.claim(org_id, idem_key, request_id)

    def get_request_id(self, org_id: str, idem_key: str) -> str | None:
        """Return the request id previously associated with the key, if any."""
        return self._store.get_request_id(org_id, idem_key)



# NOTE: The legacy limiter types (GatewayRateLimiter, DualTPMLimiter,
# BaseConcurrencyLimiter) structurally satisfy the corresponding port
# protocols (RateLimiter, TPMLimiter, ConcurrencyLimiter) — their method
# signatures match exactly. No adapter wrappers are needed; the types are
# stored directly in the Services container.


class _Replay(ReplayStore):
    """Replay store wrapper supporting both in-memory and Redis backends."""

    def __init__(self, store: LegacyReplayStore) -> None:
        self._store = store

    def append(
        self,
        request_id: str,
        *,
        event_id: int,
        data: str,
        prompt_tokens: int,
        completion_tokens: int,
        ts: int,
    ) -> None:
        """Persist a streaming event."""
        from keiro.gateway.infra.streaming.replay import StreamEvent, Usage  # lazy import

        self._store.append(
            request_id,
            StreamEvent(
                event_id=event_id,
                data=data,
                usage_delta=Usage(
                    prompt_tokens=prompt_tokens,
                    completion_tokens=completion_tokens,
                    total_tokens=prompt_tokens + completion_tokens,
                ),
            ),
        )

    def replay_from(self, request_id: str, *, after: int) -> Iterable[dict[str, object]]:
        """Yield serialized streaming events after the provided cursor."""
        events = self._store.replay_from(request_id, after_event_id=after)
        return [
            {
                "event_id": event.event_id,
                "data": event.data,
                "prompt_tokens": event.usage_delta.prompt_tokens,
                "completion_tokens": event.usage_delta.completion_tokens,
            }
            for event in events
        ]


class _Tokenizer:
    def estimate_messages(self, messages: list[Message]) -> int:
        """Estimate token usage for the provided message history."""
        return sum_message_tokens(msg.content for msg in messages)


class _Clock(Clock):
    def now_s(self) -> int:
        """Return the current epoch time in whole seconds."""
        return int(time.time())


class _Ids(IdFactory):
    def new_request_id(self) -> str:
        """Generate a new request identifier matching legacy formatting."""
        return new_request_id()

    def new_fbr_id(self) -> str:
        """Generate a billing record identifier."""
        return f"fbr_{uuid.uuid4()}"


_api_key_context: contextvars.ContextVar[str] = contextvars.ContextVar(
    "gateway2_api_key_id",
    default="anonymous",
)


class RequestScopedAuth(AuthContext):
    """Auth context backed by a context variable."""

    def api_key_id(self) -> str:
        return _api_key_context.get()


def set_request_api_key(api_key_id: str) -> contextvars.Token[str]:
    """Bind the active API key identifier for the current request context."""
    return _api_key_context.set(api_key_id)


def reset_request_api_key(token: contextvars.Token[str]) -> None:
    """Restore the previously active API key identifier."""
    _api_key_context.reset(token)


class _ErrorsFactory:
    def rate_limited(self, *, scope: str, phase: str | None = None) -> RateLimitedError:
        """Construct a rate-limited error for the provided scope."""
        return RateLimitedError(scope=scope, phase=phase)

    def no_capacity(self, *, model: str) -> NoCapacityError:
        """Construct a no-capacity error for the canonical model."""
        return NoCapacityError(model=model)


class _AdapterInvoker(ProviderInvoker):
    def __init__(self, adapter: ProviderAdapter) -> None:
        self._adapter = adapter

    def complete(
        self,
        model: str,
        prompt: str,
        *,
        temperature: float,
        api_key_override: str | None,
        system: str | None,
        context: list[dict[str, str]],
        responses_payload: Mapping[str, Any],
        org_limits: OrgLimits,
    ) -> tuple[str, int, int, Mapping[str, object]]:
        """Invoke the provider adapter synchronously."""
        params: dict[str, Any] = {
            "temperature": temperature,
            "api_key_override": api_key_override,
            "system": system,
            "context": context,
            "responses_payload": responses_payload,
            "org_limits": {
                "prompt_tokens": org_limits.prompt_tokens,
                "completion_tokens": org_limits.completion_tokens,
                "prompt_multiplier": org_limits.prompt_multiplier,
            },
        }
        text, pt, ct, metadata = self._adapter.complete(model, prompt, **params)
        return text, pt, ct, dict(metadata)

    def stream_complete(
        self,
        model: str,
        prompt: str,
        *,
        temperature: float,
        api_key_override: str | None,
        system: str | None,
        context: list[dict[str, str]],
        responses_payload: Mapping[str, Any],
        org_limits: OrgLimits,
        enable_judge_stream: bool = False,
    ) -> Iterable[str]:
        """Invoke the provider adapter in streaming mode."""
        params: dict[str, Any] = {
            "temperature": temperature,
            "api_key_override": api_key_override,
            "system": system,
            "context": context,
            "responses_payload": responses_payload,
            "org_limits": {
                "prompt_tokens": org_limits.prompt_tokens,
                "completion_tokens": org_limits.completion_tokens,
                "prompt_multiplier": org_limits.prompt_multiplier,
            },
            "enable_judge_stream": enable_judge_stream,
        }
        return self._adapter.stream_complete(model, prompt, **params)


def _build_alias_catalog() -> AliasCatalog:
    """Build the alias catalog that mirrors legacy gateway behavior."""

    mapping: dict[str, dict[str, str]] = {}
    for alias in LEGACY_ALIAS_MAP.values():
        provider_models = {}
        if alias.provider_models:
            provider_models.update(alias.provider_models)
        mapping.setdefault(alias.canonical_model, provider_models)
    return AliasCatalog(mapping)


def build_default_services(
    *,
    gateway_settings: GatewaySettings | None = None,
    legacy_settings: object | None = None,
    options: RuntimeOptions | None = None,
) -> Services:
    """Construct the default gateway service container."""
    runtime_options = options or RuntimeOptions()
    if gateway_settings is not None and legacy_settings is not None:
        raise TypeError("Provide only one of gateway_settings or legacy_settings.")
    if legacy_settings is not None:
        gateway_settings = GatewaySettings(
            postgres_dsn=legacy_settings.stores.postgres_dsn,
            redis_url=legacy_settings.stores.redis_url,
            auto_init=legacy_settings.stores.auto_init,
            default_tpm=legacy_settings.limits.default_tpm,
        )
    effective_settings = gateway_settings or _load_settings(runtime_options)
    gateway_config = effective_settings if effective_settings else GatewaySettings()
    settings = Settings(
        default_tier=gateway_config.default_tier,
        sur_dir=gateway_config.sur_dir,
        enable_judge_stream=runtime_options.enable_judge_stream,
        # Use defaults from Settings dataclass for rbc_ttl_seconds and capability_gating
    )
    registry = _ProviderRegistry()
    secret_store = InMemorySecretStore()
    credential_pool = _CredentialPool(ProviderPool())
    billing_backend = _select_billing_store(effective_settings, options=runtime_options)
    idempotency_backend = _select_idempotency_store(effective_settings, options=runtime_options)
    tpm_backend = _select_tpm_limiter(effective_settings, options=runtime_options)
    concurrency_backend = _select_concurrency_limiter(effective_settings, options=runtime_options)
    replay_backend = _select_replay_store(effective_settings, options=runtime_options)
    billing_store = _BillingStore(billing_backend)
    idempotency_store = _Idempotency(idempotency_backend)
    # Use fetch function that logs when plan lookup is needed (P0 gap until control-plane).
    plan_cache = OrgPlanCache(_default_plan_fetch_fn, default_tier=settings.default_tier)
    plans = _OrgPlanStore(plan_cache)
    pricing = _Pricing()
    rate_limiter = GatewayRateLimiter()
    tpm_limiter = tpm_backend
    concurrency = concurrency_backend
    replay = _Replay(replay_backend)
    sur_store: SurStore = FileSurStore(settings.sur_dir)
    tokenizer = _Tokenizer()
    clock = _Clock()
    ids = _Ids()
    auth = RequestScopedAuth()
    alias = _build_alias_catalog()
    # Enable capability gating by default for policy enforcement (P0 requirement).
    admission = AdmissionController(_Signer(), capability_gating=settings.capability_gating)
    router = RouterLiteAdapter()
    errors = _ErrorsFactory()

    adapters = default_adapters()
    provider_client_map = {
        name: _AdapterInvoker(adapter) for name, adapter in adapters.items()
    }
    executor = Executor(
        providers=provider_client_map,
        enable_judge_stream=settings.enable_judge_stream,
    )

    return Services(
        settings=settings,
        registry=registry,
        admission=admission,
        router=router,
        pricing=pricing,
        billing=billing_store,
        idem=idempotency_store,
        secrets=secret_store,
        credentials=credential_pool,
        plans=plans,
        tpm=tpm_limiter,
        rpm=rate_limiter,
        ccy=concurrency,
        replay=replay,
        sur=sur_store,
        tokenizer=tokenizer,
        clock=clock,
        ids=ids,
        auth=auth,
        alias=alias,
        executor=executor,
        errors=errors,
    )


__all__ = [
    "RuntimeConfigError",
    "RuntimeOptions",
    "RequestScopedAuth",
    "build_default_services",
    "reset_request_api_key",
    "set_request_api_key",
]
