from __future__ import annotations

import logging
from dataclasses import asdict
from typing import Mapping, Optional

from keiro.gateway.infra.budgets import resolve_budget_limits
from keiro.protocol import parse_responses_tool_calls

from ..domain import budget
from ..domain.errors import NoCapacityError
from ..domain.model import (
    ChatResponse,
    Choice,
    Message,
    RouterInfo,
    StreamingResponse,
)
from ..domain.resolution import resolve_model_metadata
from .container import Services
from .pipeline import Context, PipelineResult, Step
from .types import OrgLimits

LOG = logging.getLogger(__name__)


_FANOUT_OVERRIDES: dict[str, int] = {
    "takumi/eb1": 3,
    "takumi/eb1-a": 3,
    "takumi/eb1-pro": 4,
    "takumi/eb1-codex": 4,
    "takumi/eb1-codex-high": 4,
}


def _normalize_model(model: str) -> str:
    candidate = (model or "").strip().lower()
    if candidate.startswith("takumi/"):
        return candidate
    return f"takumi/{candidate}" if candidate.startswith("eb1") else candidate


def _fanout_multiplier(model: str) -> int:
    normalized = _normalize_model(model)
    return _FANOUT_OVERRIDES.get(normalized, 1)


def step_idempotency(services: Services, ctx: Context) -> Optional[PipelineResult]:
    """Normalize idempotent requests by reusing the original request id.

    Args:
        services: Pipeline service container.
        ctx: Mutable pipeline context for the active request.

    Returns:
        ``None``; the step currently never short-circuits the pipeline.
    """
    if not ctx.req.idempotency_key:
        return None
    claimed = services.idem.claim(ctx.req.org_id, ctx.req.idempotency_key, ctx.request_id)
    if not claimed:
        existing = services.idem.get_request_id(ctx.req.org_id, ctx.req.idempotency_key)
        if existing:
            ctx.request_id = existing
    return None


def step_model_resolution(services: Services, ctx: Context) -> Optional[PipelineResult]:
    """Resolve the canonical model, provider hint, and alias overrides.

    Args:
        services: Pipeline service container.
        ctx: Mutable pipeline context for the active request.

    Returns:
        ``None``.
    """
    provider_hint, canonical_model, overrides = resolve_model_metadata(ctx.req.model)
    ctx.canonical_model = canonical_model
    ctx.provider_hint = None if provider_hint == "unknown" else provider_hint
    ctx.alias_overrides = overrides
    return None


def step_plan_and_budgets(services: Services, ctx: Context) -> Optional[PipelineResult]:
    """Materialize the billing plan and enforce prompt/completion budgets.

    Args:
        services: Pipeline service container.
        ctx: Mutable pipeline context for the active request.

    Returns:
        ``None`` when limits allow the request.
    """
    plan = services.plans.get(ctx.req.org_id, fallback_tier=services.settings.default_tier)
    ctx.plan = plan
    ctx.price_book_version = str(plan.get("price_book_version", "p0-0"))
    limits = resolve_budget_limits(plan, provider=ctx.provider_hint or "unknown")
    ctx.prompt_budget = limits.prompt_tokens
    ctx.completion_budget = limits.completion_tokens
    ctx.prompt_estimate = services.tokenizer.estimate_messages(ctx.prompt_messages)
    fanout = _fanout_multiplier(ctx.canonical_model)
    ctx.prompt_multiplier = fanout
    scaled_prompt_estimate = ctx.prompt_estimate * fanout
    LOG.info(
        "plan_and_budgets",
        extra={
            "request_id": ctx.request_id,
            "prompt_budget": ctx.prompt_budget,
            "completion_budget": ctx.completion_budget,
            "prompt_estimate": ctx.prompt_estimate,
            "prompt_multiplier": fanout,
            "prompt_estimate_scaled": scaled_prompt_estimate,
        },
    )
    limit_plan = dict(plan)
    limit_plan.setdefault("prompt_token_budget", ctx.prompt_budget)
    try:
        budget.enforce_prompt(scaled_prompt_estimate, limit_plan)
        budget.enforce_requested_completion(
            ctx.req.responses.max_output_tokens,
            ctx.completion_budget,
        )
    except budget.BudgetExceededError as exc:  # pragma: no cover
        raise services.errors.rate_limited(scope="budget") from exc
    return None


def step_limits_precharge(services: Services, ctx: Context) -> Optional[PipelineResult]:
    """Consume preliminary RPM/TPM/concurrency budgets prior to execution.

    Args:
        services: Pipeline service container.
        ctx: Mutable pipeline context for the active request.

    Returns:
        ``None`` when limits are available.
    """
    if not services.rpm.check(ctx.req.org_id, cost=1):
        raise services.errors.rate_limited(scope="rpm")
    precharge_tokens = ctx.prompt_estimate * max(1, ctx.prompt_multiplier)
    if not services.tpm.consume_input(ctx.req.org_id, precharge_tokens):
        raise services.errors.rate_limited(scope="tpm", phase="prompt_precharge")
    if not services.ccy.acquire(ctx.req.org_id):
        raise services.errors.rate_limited(scope="concurrency")
    ctx.finalizers.append(lambda: services.ccy.release(ctx.req.org_id))
    return None


def step_admission(services: Services, ctx: Context) -> Optional[PipelineResult]:
    """Run admission control to select admissible providers for the request.

    Args:
        services: Pipeline service container.
        ctx: Mutable pipeline context for the active request.

    Returns:
        ``None`` when admission succeeds.
    """
    capability_map = services.registry.capabilities()

    def _supports_responses(provider: str) -> bool:
        """Return whether the provider supports the Responses API."""
        capabilities = capability_map.get(provider)
        if not isinstance(capabilities, dict):
            return False
        return bool(capabilities.get("supports_responses_api"))

    candidates = [
        provider
        for provider in services.registry.providers_for(ctx.canonical_model)
        if _supports_responses(provider)
    ]
    # Do not synthesize candidates when none are registered; fail closed.
    if ctx.provider_hint and ctx.provider_hint in candidates:
        candidates = [ctx.provider_hint, *[p for p in candidates if p != ctx.provider_hint]]
    elif ctx.provider_hint:
        if _supports_responses(ctx.provider_hint) and candidates:
            candidates = [ctx.provider_hint, *candidates]

    if not candidates:
        raise services.errors.no_capacity(model=ctx.canonical_model)

    admission = services.admission.admit(
        org_id=ctx.req.org_id,
        request_id=ctx.request_id,
        model=ctx.canonical_model,
        candidate_providers=candidates,
        policy=ctx.req.policy,
        max_prompt_tokens=ctx.prompt_budget,
        max_completion_tokens=ctx.completion_budget,
        ttl_seconds=services.settings.rbc_ttl_seconds,
        price_book_version=ctx.price_book_version,
        capabilities=capability_map,
        now_s=services.clock.now_s(),
    )
    if not admission.decision.allow or not admission.decision.admissible_providers:
        raise NoCapacityError(model=ctx.canonical_model)
    ctx.admission = admission.decision
    ctx.rbc = admission.rbc
    ctx.rbc_token = admission.token
    return None


def step_route_and_credentials(services: Services, ctx: Context) -> Optional[PipelineResult]:
    """Select a provider via router heuristics and acquire credentials.

    Args:
        services: Pipeline service container.
        ctx: Mutable pipeline context for the active request.

    Returns:
        ``None``.
    """
    assert ctx.admission is not None, "admission step must run before routing"
    chosen, components, sticky = services.router.choose(
        ctx.admission.admissible_providers,
        session_id=ctx.req.session_id,
    )
    ctx.chosen_provider = chosen
    ctx.router_info = RouterInfo(session_affinity=sticky, score_components=dict(components))
    ctx.provider_model = ctx.alias_overrides.get(chosen) or services.alias.resolve_provider_model(
        ctx.canonical_model,
        chosen,
    )
    if chosen == "takumi":
        ctx.api_key_override = None
    else:
        ctx.api_key_override = services.credentials.acquire(ctx.req.org_id, chosen)
    ctx.billing_mode = services.plans.billing_mode(ctx.plan, chosen)
    return None


def step_billing_seed(services: Services, ctx: Context) -> Optional[PipelineResult]:
    """Seed the billing record for downstream usage tracking.

    Args:
        services: Pipeline service container.
        ctx: Mutable pipeline context for the active request.

    Returns:
        ``None``.
    """
    fbr_id = services.ids.new_fbr_id()
    api_key_id = services.auth.api_key_id()
    record = services.billing.seed(
        fbr_id,
        ctx.req.org_id,
        ctx.req.idempotency_key or ctx.request_id,
        ctx.price_book_version,
        billing_mode=ctx.billing_mode,
        api_key_id=api_key_id,
    )
    ctx.fbr_id = getattr(record, "id", fbr_id) if record is not None else fbr_id
    ctx.api_key_id = api_key_id
    return None


def step_execute(services: Services, ctx: Context) -> Optional[PipelineResult]:
    """Execute the provider call and emit either streaming or buffered output.

    Args:
        services: Pipeline service container.
        ctx: Mutable pipeline context for the active request.

    Returns:
        A streaming response when streaming is enabled; otherwise ``None`` and
        the context is populated with a buffered response.
    """

    assert ctx.chosen_provider is not None, "routing step must choose a provider"
    assert ctx.fbr_id is not None, "billing seed must run before execution"
    exec_result = services.executor.execute(
        org_id=ctx.req.org_id,
        request_id=ctx.request_id,
        provider=ctx.chosen_provider,
        model=ctx.provider_model or ctx.canonical_model,
        responses_request=ctx.req.responses,
        prompt_messages=ctx.prompt_messages,
        stream=ctx.req.stream,
        temperature=ctx.temperature_hint,
        api_key_override=ctx.api_key_override,
        limits=OrgLimits(
            prompt_tokens=ctx.prompt_budget,
            completion_tokens=ctx.completion_budget,
            prompt_multiplier=ctx.prompt_multiplier,
        ),
        plan=ctx.plan,
    )
    ctx.provider_metadata = dict(exec_result.provider_metadata)
    responses_payload = (
        dict(exec_result.responses_payload)
        if isinstance(exec_result.responses_payload, Mapping)
        else None
    )
    provider_model_label = _provider_model_label(
        ctx.chosen_provider,
        ctx.provider_model or ctx.canonical_model,
    )

    if ctx.req.stream and exec_result.stream is not None:
        assert ctx.admission is not None, "admission step must set ctx.admission"
        assert ctx.rbc is not None, "admission step must set ctx.rbc"
        return StreamingResponse(
            id=ctx.request_id,
            org_id=ctx.req.org_id,
            model=ctx.canonical_model,
            provider=ctx.chosen_provider,
            provider_model=provider_model_label,
            router=ctx.router_info,
            admission=ctx.admission,
            rbc=ctx.rbc,
            rbc_token=ctx.rbc_token or "",
            stream=exec_result.stream,
            created=services.clock.now_s(),
            fbr_id=ctx.fbr_id,
            billing_mode=ctx.billing_mode,
            price_book_version=ctx.price_book_version,
            plan=dict(ctx.plan),
            idempotency_key=ctx.req.idempotency_key,
            api_key_id=ctx.api_key_id,
            prompt_estimate=ctx.prompt_estimate * max(1, ctx.prompt_multiplier),
            completion_budget=ctx.completion_budget,
            provider_metadata=ctx.provider_metadata,
            stream_source=(
                exec_result.stream_state
                if exec_result.stream_state is not None
                else exec_result.stream
            ),
            responses_request=ctx.req.responses,
            api_key_override=ctx.api_key_override,
        )

    usage = exec_result.usage
    if not services.tpm.consume_output(ctx.req.org_id, int(usage.completion_tokens)):
        raise services.errors.rate_limited(scope="tpm", phase="completion_postcharge")

    usage.total_tokens = int(usage.prompt_tokens + usage.completion_tokens)
    usage.cost_usd = float(
        services.pricing.calculate(
            model=ctx.provider_model or ctx.canonical_model,
            prompt_tokens=usage.prompt_tokens,
            completion_tokens=usage.completion_tokens,
            price_book_version=ctx.price_book_version,
            billing_mode=ctx.billing_mode,
        )
    )

    services.billing.record_delta(
        ctx.fbr_id,
        usage.prompt_tokens,
        usage.completion_tokens,
        usage.cost_usd,
    )
    services.billing.finalize(ctx.fbr_id)

    assistant_text = _assistant_message_text(exec_result.text, responses_payload)
    tool_calls = _extract_tool_calls_from_payload(responses_payload)
    finish_reason = "tool_calls" if tool_calls else "stop"

    rbc_payload = {}
    if ctx.rbc is not None:
        rbc_payload = {**asdict(ctx.rbc), "token": ctx.rbc_token}

    return ChatResponse(
        id=ctx.request_id,
        model=ctx.canonical_model,
        provider=ctx.chosen_provider,
        created=services.clock.now_s(),
        choices=[
            Choice(
                index=0,
                message=Message(
                    role="assistant",
                    content=assistant_text,
                    tool_calls=tool_calls or None,
                ),
                finish_reason=finish_reason,
            ),
        ],
        usage=usage,
        billing_record_id=ctx.fbr_id,
        router=ctx.router_info,
        raw={
            "rbc": rbc_payload,
            "admission": asdict(ctx.admission) if ctx.admission else {},
            "provider_metadata": ctx.provider_metadata,
            "responses_payload": responses_payload,
            "plan": ctx.plan,
            "idempotency_key": ctx.req.idempotency_key,
            "pricing": {
                "mode": ctx.billing_mode,
                "price_book_version": ctx.price_book_version,
                "cost_usd": usage.cost_usd,
            },
        },
    )


PIPELINE_STEPS: list[Step] = [
    step_idempotency,
    step_model_resolution,
    step_plan_and_budgets,
    step_limits_precharge,
    step_admission,
    step_route_and_credentials,
    step_billing_seed,
    step_execute,
]


def _provider_model_label(provider: str, model: str) -> str:
    if model.startswith(f"{provider}/"):
        return model
    return f"{provider}/{model}"


def _extract_tool_calls_from_payload(
    payload: Mapping[str, object] | None,
) -> list[dict[str, object]]:
    """Convert Responses API function_call items to chat tool_calls format."""
    if not payload:
        return []
    outputs = payload.get("output")
    if not isinstance(outputs, list):
        return []
    parsed = parse_responses_tool_calls(outputs)
    return [tc.to_chat_dict() for tc in parsed]


def _assistant_message_text(
    response_text: str | None,
    payload: Mapping[str, object] | None,
) -> str:
    """Extract assistant text from the provider result, falling back to payload."""
    if response_text:
        return response_text
    if not payload:
        return ""
    outputs = payload.get("output")
    parts: list[str] = []
    if isinstance(outputs, list):
        for item in outputs:
            if not isinstance(item, Mapping):
                continue
            content = item.get("content")
            if isinstance(content, list):
                for chunk in content:
                    if not isinstance(chunk, Mapping):
                        continue
                    text_value = chunk.get("text")
                    if isinstance(text_value, str):
                        parts.append(text_value)
            elif isinstance(content, str):
                parts.append(content)
    if parts:
        return "".join(parts)
    data_section = payload.get("data")
    if isinstance(data_section, Mapping):
        text_value = data_section.get("text")
        if isinstance(text_value, str):
            return text_value
    return ""
