"""Reasoner configuration sections for :func:`relationalai.config.Config`.

Most projects configure reasoner defaults in a YAML file (typically
``raiconfig.yaml``) under the ``reasoners`` section (for example
``reasoners.logic.size`` and ``reasoners.logic.auto_suspend_mins``). This
module defines the Pydantic models that validate those settings, along with the
allowed reasoner size strings for each deployment platform.

You can also configure everything in code by passing dictionaries to
the :func:`~relationalai.config.create_config` factory function (useful for programmatic
setup or for overrides in tests).

Start with :class:`~relationalai.config.config_reasoners_fields.ReasonersConfig`
(the top-level ``reasoners`` section) for configuring base reasoner API settings.
Shared fields live in :class:`~relationalai.config.config_reasoners_fields.ReasonerConfig`,
and logic reasoner-specific options live in
:class:`~relationalai.config.config_reasoners_fields.LogicReasonerConfig`.
"""

from __future__ import annotations

from typing import Any, Literal

from pydantic import BaseModel, ConfigDict, Field

from ..util.docutils import include_in_docs


# include this module in documentation
__include_in_docs__ = True


# Engine size type aliases
# Internal development sizes
InternalEngineSize = Literal["XS", "S", "M", "L"]
"""Internal-only reasoner size values.

These sizes should not be used by end users and are typically hidden unless
:attr:`~relationalai.config.config_reasoners_fields.ReasonersConfig.show_all_sizes` is enabled.
"""

# AWS Snowflake sizes
AWSEngineSize = Literal["HIGHMEM_X64_S", "HIGHMEM_X64_M", "HIGHMEM_X64_L", "GPU_NV_S"]
"""Valid reasoner size values on AWS-backed deployments.

Use one of these strings for :attr:`~relationalai.config.config_reasoners_fields.ReasonerConfig.size`.
"""

# Azure sizes
AzureEngineSize = Literal["HIGHMEM_X64_S", "HIGHMEM_X64_M", "HIGHMEM_X64_SL"]
"""Valid reasoner size values on Azure-backed deployments.

Use one of these strings for :attr:`~relationalai.config.config_reasoners_fields.ReasonerConfig.size`.
"""

# All valid engine sizes across platforms
EngineSize = InternalEngineSize | AWSEngineSize | AzureEngineSize
"""All supported reasoner size values.

This is the type accepted by :attr:`~relationalai.config.config_reasoners_fields.ReasonerConfig.size`.
Internal/dev sizes are included but are typically hidden unless
:attr:`~relationalai.config.config_reasoners_fields.ReasonersConfig.show_all_sizes` is enabled.
"""

# LQP Semantics Version
# When the logic engine introduces breaking changes, this version is bumped
# [2026-01-09] Version "1" opts into hard validation errors from the engine
LqpSemanticsVersion = Literal["1"]


@include_in_docs
class ReasonerLogicLqpConfig(BaseModel):
    """Configure LQP-specific behavior for the logic reasoner.

    Most users should leave this unset. Only set it if a RelationalAI release
    note or support instruction asks you to opt into a specific semantics
    version.

    Attributes
    ----------
    semantics_version : {"1"}, optional
        Semantics version to opt into.
    """

    semantics_version: LqpSemanticsVersion | None = Field(
        default=None,
        description="LQP semantics version for breaking changes (e.g., '1' for hard validation errors)"
    )


@include_in_docs
class ReasonerConfig(BaseModel):
    """Common configuration fields for a reasoner.

    This model defines the fields shared by all reasoner types. You set these
    fields under one of the reasoner subsections in ``raiconfig.yaml``:

    - ``reasoners.logic`` (logic reasoner, see
      :class:`~relationalai.config.config_reasoners_fields.LogicReasonerConfig`
      for additional logic-specific settings)
    - ``reasoners.predictive`` (predictive reasoner)
    - ``reasoners.prescriptive`` (prescriptive reasoner)

    Attributes
    ----------
    name : str, optional
        Optional name for the reasoner instance.
    size : str
        Compute size (for example, ``"HIGHMEM_X64_S"``).
        See :attr:`~relationalai.config.config_reasoners_fields.EngineSize` for valid size values.
    query_timeout_mins : int, optional
        Abort queries that exceed this duration.
    auto_suspend_mins : int
        Auto-suspend the reasoner after N minutes of inactivity.
    await_storage_vacuum : bool, optional
        Whether to wait for storage vacuum work to finish before suspending. Defaults to ``None``.
    settings : dict, optional
        Backend-specific settings payload.

    Examples
    --------
    Minimal YAML (in ``raiconfig.yaml``) configuring only the logic reasoner:

    ```yaml
    default_connection: sf
    connections:
        sf:
            type: snowflake
            # ...
    reasoners:
        logic:
            size: HIGHMEM_X64_S
            query_timeout_mins: 30
    ```

    Any omitted subsections (for example ``reasoners.predictive``) keep their
    built-in defaults.

    Configure reasoner defaults using a dict:

    >>> from relationalai.config import create_config
    >>> cfg = create_config(reasoners={
    ...     "logic": {
    ...         "size": "HIGHMEM_X64_S",
    ...         "query_timeout_mins": 30,
    ...     }
    ... })

    Configure reasoner defaults using an explicit ReasonerConfig instance:

    >>> from relationalai.config import create_config, ReasonerConfig
    >>> cfg = create_config(
    ...     reasoners={"predictive": ReasonerConfig(size="HIGHMEM_X64_S", query_timeout_mins=30)},
    ... )
    """

    model_config = ConfigDict(extra="allow")

    name: str | None = Field(
        default=None,
        description="Optional name for the reasoner instance"
    )

    size: EngineSize = Field(
        default="HIGHMEM_X64_S",
        description="Compute size (e.g., 'HIGHMEM_X64_S', 'HIGHMEM_X64_M')",
        examples=["HIGHMEM_X64_S", "HIGHMEM_X64_M", "HIGHMEM_X64_L"]
    )

    query_timeout_mins: int | None = Field(
        default=None,
        description="Query timeout in minutes - aborts queries that exceed this duration"
    )

    auto_suspend_mins: int = Field(
        default=60,
        description="Auto-suspend reasoner after N minutes of inactivity"
    )

    await_storage_vacuum: bool | None = Field(
        default=None,
        description="Wait for storage vacuum work to finish before suspending the reasoner",
    )

    settings: dict[str, Any] | None = Field(
        default=None,
        description="Optional reasoner settings payload (backend-dependent)",
    )


@include_in_docs
class LogicReasonerConfig(ReasonerConfig):
    """Configure rule-execution defaults for the logic reasoner.

    Set these fields under ``reasoners.logic`` to control how the logic reasoner
    runs rules (for example, whether to use LQP) and to opt into an LQP
    semantics version when required.

    Attributes
    ----------
    use_lqp : bool
        Whether to use LQP for rule execution.
    incremental_maintenance : str, optional
        Incremental maintenance strategy string.
        Can be set to ``"on"``, ``"auto"``, or ``"off"`` (default).
    emit_constraints : bool
        Whether to emit constraints during execution.
    lqp : ReasonerLogicLqpConfig
        LQP-specific options, such as ``semantics_version``.
        See :class:`~relationalai.config.config_reasoners_fields.ReasonerLogicLqpConfig` for details.

    Examples
    --------
    Minimal YAML (in ``raiconfig.yaml``):

    ```yaml
    default_connection: sf
    connections:
        sf:
            type: snowflake
            # ...
    reasoners:
        logic:
            incremental_maintenance: "on"
    ```

    Configure logic reasoner settings using a dict:

    >>> from relationalai.config import create_config
    >>> cfg = create_config(reasoners={"logic": {"use_lqp": True, "lqp": {"semantics_version": "1"}}})

    Configure logic reasoner settings using an explicit LogicReasonerConfig instance:

    >>> from relationalai.config import create_config, LogicReasonerConfig
    >>> cfg = create_config(
    ...     reasoners={"logic": LogicReasonerConfig(use_lqp=True, lqp={"semantics_version": "1"})},
    ... )
    """

    use_lqp: bool = Field(
        default=True,
        description="Use LQP for reasoner rule execution"
    )

    incremental_maintenance: str | None = Field(
        default="off",
        description="Incremental maintenance strategy"
    )

    emit_constraints: bool = Field(
        default=False,
        description="Emit constraints during execution"
    )

    lqp: ReasonerLogicLqpConfig = Field(
        default_factory=ReasonerLogicLqpConfig,
        description="LQP-specific configuration"
    )


@include_in_docs
class PrescriptiveReasonerConfig(ReasonerConfig):
    """Defaults for the prescriptive reasoner."""

    auto_suspend_mins: int = Field(
        default=30,
        description="Auto-suspend reasoner after N minutes of inactivity",
    )


@include_in_docs
class ReasonersConfig(BaseModel):
    """Configure defaults for the reasoners API.

    This section lives under ``reasoners`` in ``raiconfig.yaml`` and controls the
    reasoners backend, polling behavior for long-running operations, and default
    settings for each reasoner type (``logic``, ``predictive``, ``prescriptive``).

    Attributes
    ----------
    backend : {"sql", "direct_access"}
        Backend used for reasoners APIs.
    show_all_sizes : bool
        Show all available reasoner sizes (includes internal/dev sizes).
    poll_initial_delay_s : float
        Initial delay between polls when waiting for operations to complete.
    poll_overhead_rate : float
        Exponential backoff growth rate per poll (for example, ``0.2`` means +20% per poll).
    poll_max_delay_s : float
        Maximum poll delay.
    logic : LogicReasonerConfig
        Defaults for the logic reasoner.
    predictive : ReasonerConfig
        Defaults for the predictive reasoner.
    prescriptive : ReasonerConfig
        Defaults for the prescriptive reasoner.

    Examples
    --------
    Minimal YAML (in ``raiconfig.yaml``):

    ```yaml
    default_connection: sf
    connections:
        sf:
            type: snowflake
            # ...
    reasoners:
        direct_access: true
        poll_max_delay_s: 2.0
        logic:
            size: HIGHMEM_X64_S
            query_timeout_mins: 30
    ```

    Configure reasoners settings using a dict:

    >>> from relationalai.config import create_config
    >>> cfg = create_config(reasoners={
    ...     "poll_max_delay_s": 2.0,
    ...     "logic": {"size": "HIGHMEM_X64_S", "query_timeout_mins": 30},
    ... })

    Configure reasoners settings using an explicit ReasonersConfig instance:

    >>> from relationalai.config import create_config, ReasonersConfig, LogicReasonerConfig
    >>> cfg = create_config(
    ...     reasoners=ReasonersConfig(
    ...         poll_max_delay_s=2.0,
    ...         logic=LogicReasonerConfig(size="HIGHMEM_X64_S", query_timeout_mins=30),
    ...     ),
    ... )
    """

    backend: Literal["sql", "direct_access"] | None = Field(
        default=None,
        description="Legacy compatibility field. Prefer `reasoners.direct_access` or the top-level `direct_access` setting.",
    )

    direct_access_base_url: str | None = Field(
        default=None,
        description="Legacy compatibility field for Direct Access base URL. Newer DA wiring resolves service endpoints dynamically.",
    )

    direct_access: bool | None = Field(
        default=None,
        description="Preferred Direct Access toggle for reasoners. When set, overrides the top-level direct_access default.",
    )

    # Legacy note: there is also `config.engine.show_all_sizes` (engine terminology).
    # The reasoners surface uses the same knob, but names it here to avoid leaking
    # backend "engine" concepts into the reasoners API.
    show_all_sizes: bool = Field(
        default=False,
        description="Show all available reasoner sizes (includes internal/dev sizes).",
    )

    # Polling defaults used by `wait_until_ready(...)`, `ReasonerOperation.wait(...)`,
    # and other "wait until done" helpers. These are *not* retry policy knobs; they
    # control how aggressively we poll backend state during long-running operations.
    poll_initial_delay_s: float = Field(default=0.05, description="Initial poll delay (seconds)")
    poll_overhead_rate: float = Field(default=0.2, description="Exponential backoff growth rate per poll (0.2 => +20%)")
    poll_max_delay_s: float = Field(default=2.0, description="Maximum poll delay (seconds)")

    logic: LogicReasonerConfig = Field(
        default_factory=LogicReasonerConfig,
        description="Logic reasoner configuration"
    )

    predictive: ReasonerConfig = Field(
        default_factory=ReasonerConfig,
        description="Predictive reasoner configuration"
    )

    prescriptive: PrescriptiveReasonerConfig = Field(
        default_factory=PrescriptiveReasonerConfig,
        description="Prescriptive reasoner configuration"
    )
