from __future__ import annotations

from typing import Any, Mapping, TypedDict

from relationalai.services.constants import REASONER_TYPE_LOGIC, REASONER_TYPE_PREDICTIVE, REASONER_TYPE_PRESCRIPTIVE

from .errors import JobClientDisconnectedError, JobInvalidRequestError
from .models import Job

_PUBLIC_REASONER_TYPES: dict[str, tuple[str, str]] = {
    # input (upper) -> (canonical jobs discriminator, canonical public label)
    #
    # Public labels:
    REASONER_TYPE_LOGIC: (REASONER_TYPE_LOGIC, "Logic"),
    REASONER_TYPE_PRESCRIPTIVE: ("SOLVER", "Prescriptive"),
    REASONER_TYPE_PREDICTIVE: ("ML", "Predictive"),
    #
    # Legacy/internal discriminators (accepted for programmatic compatibility):
    "SOLVER": ("SOLVER", "Prescriptive"),
    "ML": ("ML", "Predictive"),
}

_PUBLIC_REASONER_TYPES_EXPECTED = "Logic, Prescriptive, or Predictive"


class JobIncludeDict(TypedDict, total=False):
    """Structured include options for `JobsClient.get/list/wait`.

    Only keys set to True are included.
    """

    details: bool
    events: bool
    artifacts: bool  # LOGIC-only
    problems: bool  # LOGIC-only
    raw: bool


_ALL_EXPANSIONS: set[str] = {"details", "events", "artifacts", "problems", "raw"}
_LOGIC_ONLY_EXPANSIONS: set[str] = {"artifacts", "problems"}
_COMMON_EXPANSIONS: set[str] = _ALL_EXPANSIONS - _LOGIC_ONLY_EXPANSIONS

_CLIENT_DISCONNECTED_ABORT_REASON = "client disconnected"


def raise_on_client_disconnected(*, reasoner_type: str, job: Job) -> None:
    """Raise when a LOGIC transaction was aborted due to client disconnect."""
    if str(reasoner_type).strip().upper() != REASONER_TYPE_LOGIC:
        return
    reason = (job.abort_reason or "").strip().lower()
    if reason == _CLIENT_DISCONNECTED_ABORT_REASON:
        raise JobClientDisconnectedError()


def normalize_include(include: JobIncludeDict | Mapping[str, bool] | None) -> set[str]:
    """Normalize structured `include=` into a lower-cased set of include keys.

    Supports:
    - mapping of include keys to booleans, where only truthy values are enabled
    - ``None``
    """
    if include is None:
        return set()
    if not isinstance(include, Mapping):
        raise JobInvalidRequestError("Job `include` must be a dict-like mapping (e.g. {'details': True}).")

    out: set[str] = set()
    for k, v in include.items():
        if not isinstance(k, str) or not k.strip():
            raise JobInvalidRequestError("Job `include` keys must be non-empty strings.")
        key = k.strip().lower()
        if key not in _ALL_EXPANSIONS:
            allowed = ", ".join(sorted(_ALL_EXPANSIONS))
            raise JobInvalidRequestError(f"Invalid job include key {k!r}. Allowed: {allowed}")
        if not isinstance(v, bool):
            raise JobInvalidRequestError(f"Job `include` values must be bool (key={k!r}).")
        if v:
            out.add(key)
    return out


def validate_include_for_type(reasoner_type: str, include: set[str]) -> None:
    """Raise if `include` contains keys not supported for this reasoner type."""
    t = require_nonempty_str("reasoner_type", reasoner_type).upper()
    if t == REASONER_TYPE_LOGIC:
        return
    bad = include.intersection(_LOGIC_ONLY_EXPANSIONS)
    if bad:
        raise JobInvalidRequestError(
            "Unsupported includes for non-LOGIC jobs: "
            + ", ".join(sorted(bad))
            + ". (artifacts/problems are LOGIC-only)"
        )


def include_options_for_type(reasoner_type: str) -> JobIncludeDict:
    """Return an include dict template showing available keys for a reasoner type.

    All values are False by default; set desired keys to True.
    """
    t = require_nonempty_str("reasoner_type", reasoner_type).upper()
    keys = set(_ALL_EXPANSIONS) if t == REASONER_TYPE_LOGIC else set(_COMMON_EXPANSIONS)
    return {k: False for k in sorted(keys)}  # type: ignore[return-value]


def require_nonempty_str(field: str, value: Any) -> str:
    if not isinstance(value, str) or not value.strip():
        raise JobInvalidRequestError(f"Job `{field}` is required.")
    return value.strip()


def normalize_public_reasoner_type(value: Any, *, field: str = "reasoner_type") -> str:
    """Normalize a reasoner type-like string to a canonical Jobs discriminator.

    Public reasoner type names accepted by the API:
    - Logic -> LOGIC
    - Prescriptive -> SOLVER
    - Predictive -> ML

    For programmatic backwards compatibility, the canonical Jobs gateway discriminators
    (`LOGIC`/`SOLVER`/`ML`) are also accepted.
    """
    raw = require_nonempty_str(field, value)
    k = raw.upper()

    v = _PUBLIC_REASONER_TYPES.get(k)
    if v is not None:
        canonical, _label = v
        return canonical

    raise JobInvalidRequestError(f"Invalid job `{field}` {raw!r}. Expected: {_PUBLIC_REASONER_TYPES_EXPECTED}.")


def normalize_public_reasoner_type_label(value: Any, *, field: str = "reasoner_type") -> str:
    """Return a stable public-facing label (Logic/Prescriptive/Predictive)."""
    raw = require_nonempty_str(field, value)
    k = raw.upper()
    v = _PUBLIC_REASONER_TYPES.get(k)
    if v is None:
        # Reuse the canonical validator + error messages.
        _ = normalize_public_reasoner_type(raw, field=field)
        return raw  # unreachable
    _canonical, label = v
    return label


def normalize_optional_nonempty_str(field: str, value: Any) -> str | None:
    if value is None:
        return None
    if not isinstance(value, str) or not value.strip():
        raise JobInvalidRequestError(f"Job `{field}` must be a non-empty string when provided.")
    return value.strip()


def require_payload_dict(value: Any) -> dict[str, Any]:
    if not isinstance(value, dict):
        raise JobInvalidRequestError("Job `payload` must be a dict.")
    return value


def normalize_abort_detached_payload(*, reasoner_type: str, payload: dict[str, Any], enable: bool) -> dict[str, Any]:
    """Return a copy of `payload` with a normalized `abort_detached` flag.

    Policy:
    - Only LOGIC transactions are allowed to send `abort_detached`.
    - Only blocking/polling call-sites should enable it.
    """
    out = dict(payload)
    # Always strip user-provided value unless explicitly enabled for LOGIC.
    out.pop("abort_detached", None)
    if enable and str(reasoner_type).strip().upper() == REASONER_TYPE_LOGIC:
        out["abort_detached"] = True
    return out


# ---------------------------------------------------------------------------
# Payload helpers (best-effort convenience).
# ---------------------------------------------------------------------------


def compose_payload(
    *,
    reasoner_type: str,
    database: str | None = None,
    raw_code: str,
    inputs: dict[str, Any] | None = None,
    readonly: bool = True,
    nowait_durable: bool = False,
    language: str = "rel",
    headers: dict[str, Any] | None = None,
    timeout_mins: int | None = None,
    # Included for parity with job payloads; ignored by LOGIC create in v1.
    bypass_index: bool = False,
    gi_setup_skipped: bool = False,
) -> dict[str, Any]:
    """Build a best-effort payload for `client.jobs.create(...)` for any type.

    Routing:
    - `reasoner_type == "Logic"`: payload is shaped for **transaction** creation.
    - otherwise: payload is shaped for **job** creation (backend/job-type specific).
    """
    rt = normalize_public_reasoner_type(reasoner_type, field="reasoner_type")
    in_d = inputs or {}

    # Base keys that are useful for job payloads and debugging.
    out: dict[str, Any] = {
        "raw_code": str(raw_code),
        "inputs": in_d,
        "readonly": bool(readonly),
        "nowait_durable": bool(nowait_durable),
        "language": str(language),
        "bypass_index": bool(bypass_index),
        "gi_setup_skipped": bool(gi_setup_skipped),
    }
    if database is not None and str(database).strip():
        out["database"] = str(database)
    if headers is not None:
        out["headers"] = dict(headers)
    if timeout_mins is not None:
        out["timeout_mins"] = int(timeout_mins)

    if rt == REASONER_TYPE_LOGIC:
        if database is None or not str(database).strip():
            raise ValueError("database is required for reasoner_type='Logic'.")
        # Stable transaction-create keys expected by our SQL/DA wiring.
        out["dbname"] = str(database)
        out["query"] = str(raw_code)
        out["v1_inputs"] = in_d

    return out


def make_txn_payload(
    *,
    dbname: str,
    query: str,
    v1_inputs: dict[str, Any] | None = None,
    readonly: bool = True,
    nowait_durable: bool = False,
    language: str = "rel",
    headers: dict[str, Any] | None = None,
    timeout_mins: int | None = None,
) -> dict[str, Any]:
    """Build a payload for `client.jobs.create(\"Logic\", ...)`.

    Prefer `compose_payload(reasoner_type=..., ...)` for a single helper that works
    for both LOGIC transactions and non-LOGIC jobs.
    """
    return compose_payload(
        reasoner_type="Logic",
        database=dbname,
        raw_code=query,
        inputs=v1_inputs,
        readonly=readonly,
        nowait_durable=nowait_durable,
        language=language,
        headers=headers,
        timeout_mins=timeout_mins,
    )


def make_job_payload(
    *,
    database: str | None = None,
    raw_code: str,
    inputs: dict[str, Any] | None = None,
    readonly: bool = True,
    nowait_durable: bool = False,
    language: str = "rel",
    headers: dict[str, Any] | None = None,
    timeout_mins: int | None = None,
    bypass_index: bool = False,
    gi_setup_skipped: bool = False,
) -> dict[str, Any]:
    """Build a best-effort non-LOGIC jobs payload.

    This is a convenience wrapper when the caller doesn't want to specify
    `reasoner_type`. The returned payload is suitable for non-LOGIC jobs; callers
    should still pass the actual `reasoner_type` to `client.jobs.create(...)`.
    """
    return compose_payload(
        reasoner_type="Prescriptive",  # any non-Logic type yields the same job-shaped payload
        database=database,
        raw_code=raw_code,
        inputs=inputs,
        readonly=readonly,
        nowait_durable=nowait_durable,
        language=language,
        headers=headers,
        timeout_mins=timeout_mins,
        bypass_index=bypass_index,
        gi_setup_skipped=gi_setup_skipped,
    )


__all__ = [
    "JobIncludeDict",
    "compose_payload",
    "include_options_for_type",
    "normalize_public_reasoner_type",
    "normalize_public_reasoner_type_label",
    "make_job_payload",
    "make_txn_payload",
    "normalize_include",
    "normalize_optional_nonempty_str",
    "require_nonempty_str",
    "require_payload_dict",
    "validate_include_for_type",
]

