"""Solver engine lifecycle and job execution.

Provides the Solver and Provider classes for managing solver engines
and executing optimization jobs on RelationalAI's solver infrastructure.
"""

from __future__ import annotations

import json
import logging
import textwrap
from dataclasses import dataclass
from typing import Any, Optional, cast

import relationalai
from relationalai import debugging
from relationalai.clients.direct_access_client import DirectAccessClient
from relationalai.clients.resources.snowflake import (
    APP_NAME,
    DirectAccessResources,
    Resources as SnowflakeResources,
)
from relationalai.clients.util import poll_with_specified_overhead
from relationalai.errors import QueryTimeoutExceededException, ResponseStatusException
from relationalai.tools.cli_controls import Spinner
from relationalai.tools.constants import DEFAULT_QUERY_TIMEOUT_MINS

logger = logging.getLogger(__name__)

ENGINE_TYPE_SOLVER = "SOLVER"
# TODO (dba) The ERP still uses `worker` instead of `engine`. Change
# this once we fix this in the ERP.
WORKER_ERRORS = [
    "worker is suspended",
    "create/resume",
    "worker not found",
    "no workers found",
    "worker was deleted",
]
ENGINE_ERRORS = [
    "engine is suspended",
    "create/resume",
    "engine not found",
    "no engines found",
    "engine was deleted",
]
ENGINE_NOT_READY_MSGS = [
    "worker is in pending",
    "worker is provisioning",
    "worker is not ready to accept jobs",
    "resume in progress",
]

JOB_ABORT_REASON_TIMEOUT = "job timeout"


@dataclass
class PollingState:
    job_id: str
    continuation_token: str
    is_done: bool
    log_to_console: bool


class Solver:
    def __init__(
        self,
        solver_name: str,
        engine_name: str | None = None,
        engine_size: str | None = None,
        auto_suspend_mins: int | None = None,
        resources: SnowflakeResources | None = None,
    ):
        self.provider = Provider(resources=resources)
        self.solver_name = solver_name.lower()

        self.rai_config = self.provider.resources.config
        settings: dict[str, Any] = {}
        if "experimental" in self.rai_config:
            exp_config = self.rai_config.get("experimental", {})
            if isinstance(exp_config, dict):
                if "solvers" in exp_config:
                    settings = exp_config["solvers"].copy()

        # Engine configuration fields are not necessary for the solver
        # settings so we `pop` them from the settings object. Default
        # size and auto_suspend_mins are set in the `Provider` methods.
        engine_name = engine_name or settings.pop("engine", None)
        if not engine_name:
            engine_name = self.provider.resources.get_user_based_engine_name()
        self.engine_name = engine_name

        self.engine_size = engine_size or settings.pop("engine_size", None)
        self.engine_auto_suspend_mins = auto_suspend_mins or settings.pop(
            "auto_suspend_mins", None
        )

        # Set default CSV store setting if not already configured
        if "store" not in settings:
            settings["store"] = {}
        if "csv" not in settings["store"]:
            settings["store"]["csv"] = {}
        if "enabled" not in settings["store"]["csv"]:
            settings["store"]["csv"]["enabled"] = True

        # The settings are used when creating a solver engine, they
        # may configure each individual solver.
        self.engine_settings = settings

        # Optimistically set the engine object to a `READY` engine to
        # avoid checking the engine status on each execution.
        self.engine: Optional[dict[str, Any]] = {"name": engine_name, "state": "READY"}

        return None

    # --------------------------------------------------
    # Helper
    # --------------------------------------------------
    def _auto_create_solver_async(self):
        name = self.engine_name
        auto_suspend_mins = self.engine_auto_suspend_mins
        size = self.engine_size
        settings = self.engine_settings
        with Spinner(
            "Checking solver status",
            leading_newline=True,
        ) as spinner:
            engine = None
            engines = [e for e in self.provider.list_solvers() if e["name"] == name]
            if len(engines) > 1:
                raise RuntimeError(
                    f"Expected at most 1 engine named '{name}', but found {len(engines)}."
                )
            if len(engines) != 0:
                engine = engines[0]

            if engine:
                # TODO (dba) Logic engines support altering the
                # auto_suspend_mins setting. Currently, we don't have
                # this capability for solver engines, so users need to
                # recreate or use another engine. For both the size
                # and Gurobi configuration the user anyways has to
                # create a new one as this configuration must happen
                # when the engine is created.
                settings_cannot_be_altered_msg = "The configuration of a solver engine happens when the engine is created and _cannot_ be changed. You either need to specify a new engine name or delete the current engine.\n\nSee `solvers.Provider().delete_solver()`."
                # Make sure that the solver requested is enabled
                # on the engine.
                if self.solver_name not in engine["solvers"]:
                    raise Exception(
                        f"Solver `{self.solver_name}` is not enabled on `{name}`.\n\n"
                        + settings_cannot_be_altered_msg
                    )

                # Make sure size and auto_suspend_mins settings match
                # what the user requests.
                if size is not None and size != engine["size"]:
                    raise Exception(
                        f"Engine `{name}` has size setting of `{engine['size']}` but size is requested to be `{size}`.\n\n"
                        + settings_cannot_be_altered_msg
                    )

                if (
                    auto_suspend_mins is not None
                    and auto_suspend_mins != engine["auto_suspend_mins"]
                ):
                    raise Exception(
                        f"Engine `{name}` has auto_suspend_mins setting of `{engine['auto_suspend_mins']}` but auto_suspend_mins is requested to be `{auto_suspend_mins}`.\n\n"
                        + settings_cannot_be_altered_msg
                    )

                if engine["state"] == "PENDING":
                    spinner.update_messages(
                        {
                            "finished_message": f"Solver {name} is starting",
                        }
                    )
                    pass
                elif engine["state"] == "SUSPENDED":
                    spinner.update_messages(
                        {
                            "finished_message": f"Resuming solver {name}",
                        }
                    )
                    self.provider.resume_solver_async(name)
                elif engine["state"] == "READY":
                    spinner.update_messages(
                        {
                            "finished_message": f"Solver {name} is ready",
                        }
                    )
                    pass
                else:
                    spinner.update_messages(
                        {
                            "message": f"Restarting solver {name}",
                        }
                    )
                    self.provider.delete_solver(name)
                    engine = None
            if not engine:
                # Validate Gurobi config.
                if self.solver_name == "gurobi":
                    is_gurobi_configured = False
                    gurobi_config = settings.get("gurobi", {})
                    if all(
                        k in gurobi_config
                        for k in ["license_secret_name", "external_access_integration"]
                    ):
                        is_gurobi_configured = True
                    if not is_gurobi_configured:
                        raise Exception(
                            "Gurobi is not properly configured. You need to provide both `license_secret_name` and `external_access_integration` in its configuration, see https://docs.relational.ai/build/reasoners/prescriptive/solver-backends/gurobi/#usage"
                        )
                self.provider.create_solver_async(
                    name,
                    settings=settings,
                    size=size,
                    auto_suspend_mins=auto_suspend_mins,
                )
                engine = self.provider.get_solver(name)
                spinner.update_messages(
                    {
                        "finished_message": f"Starting solver {name}...",
                    }
                )

            self.engine = engine

    def _exec_job(
        self, payload, log_to_console=True, query_timeout_mins: Optional[int] = None
    ):
        if self.engine is None:
            raise Exception("Engine not initialized.")

        with debugging.span("job") as job_span:
            # Retry logic. If creating a job fails with an engine
            # related error we will create/resume/... the engine and
            # retry.
            try:
                job_id = self.provider.create_job_async(
                    self.engine["name"], payload, query_timeout_mins=query_timeout_mins
                )
            except Exception as e:
                err_message = str(e).lower()
                if isinstance(e, ResponseStatusException):
                    try:
                        err_message = e.response.json().get("message", "")
                    except (ValueError, AttributeError):
                        err_message = e.response.text if hasattr(e.response, 'text') else str(e)
                if any(
                    kw in err_message.lower()
                    for kw in ENGINE_ERRORS + WORKER_ERRORS + ENGINE_NOT_READY_MSGS
                ):
                    self._auto_create_solver_async()
                    # Wait until the engine is ready.
                    poll_with_specified_overhead(lambda: self._is_solver_ready(), 0.1)
                    job_id = self.provider.create_job_async(
                        self.engine["name"],
                        payload,
                        query_timeout_mins=query_timeout_mins,
                    )
                else:
                    raise

            job_span["job_id"] = job_id
            debugging.event(
                "job_created",
                job_span,
                job_id=job_id,
                engine_name=self.engine["name"],
                job_type=ENGINE_TYPE_SOLVER,
            )
            if not isinstance(job_id, str):
                job_id = ""
            polling_state = PollingState(job_id, "", False, log_to_console)

            try:
                with debugging.span("wait", job_id=job_id):
                    poll_with_specified_overhead(
                        lambda: self._check_job_status(polling_state), 0.1
                    )
            except KeyboardInterrupt as e:
                print(f"Canceling job {job_id}")
                self.provider.cancel_job(job_id)
                raise e

            return job_id

    def _is_solver_ready(self):
        if self.engine is None:
            raise Exception("Engine not initialized.")

        result = self.provider.get_solver(self.engine["name"])

        if result is None:
            raise Exception("No engine available.")

        self.engine = result
        state = result["state"]
        if state != "READY" and state != "PENDING":
            # Might have suspended or otherwise gone. Recreate.
            self._auto_create_solver_async()
        return state == "READY"

    def _check_job_status(self, state):
        response = self.provider.get_job(state.job_id)
        if not response:
            raise RuntimeError(f"No results from get_job('{state.job_id}')")

        status: str = response["state"]
        abort_reason = response.get("abort_reason", "")

        self._print_solver_logs(state)
        if status == "CANCELED" and abort_reason == JOB_ABORT_REASON_TIMEOUT:
            config_file_path = getattr(self.rai_config, "file_path", None)
            timeout_ms = int(response.get("timeout_ms", 0))
            timeout_mins = (
                timeout_ms // 60000
                if timeout_ms > 0
                else int(
                    self.rai_config.get(
                        "query_timeout_mins", DEFAULT_QUERY_TIMEOUT_MINS
                    )
                    or DEFAULT_QUERY_TIMEOUT_MINS
                )
            )
            raise QueryTimeoutExceededException(
                timeout_mins=timeout_mins,
                query_id=state.job_id,
                config_file_path=config_file_path,
            )

        if status == "FAILED":
            raise RuntimeError(
                f"Solver job '{state.job_id}' failed."
            )
        if status == "CANCELED":
            raise RuntimeError(
                f"Solver job '{state.job_id}' was canceled. "
                f"Reason: {abort_reason or 'unknown'}"
            )

        return status == "COMPLETED"

    def _print_solver_logs(self, state: PollingState):
        if state.is_done:
            return

        resp = self.provider.get_job_events(state.job_id, state.continuation_token)

        # Print solver logs to stdout.
        for event in resp["events"]:
            if event["type"] == "LogMessage":
                if state.log_to_console:
                    print(event["event"]["message"])
            else:
                continue

        state.continuation_token = resp["continuation_token"]
        if state.continuation_token == "":
            state.is_done = True


# --------------------------------------------------
# Provider
# --------------------------------------------------
#
# TODO (dba) We use an experimental and unified engine API for
# solvers. Once it is no longer experimental we can remove the
# provider here and use the normal PyRel provider.


class Provider:
    def __init__(self, resources=None):
        if not resources:
            resources = relationalai.Resources()
        if not isinstance(resources, SnowflakeResources):
            raise Exception("Solvers are only supported on SPCS.")

        # Type narrowing: resources is confirmed to be SnowflakeResources
        self.resources: SnowflakeResources = cast(SnowflakeResources, resources)
        self.direct_access_client: Optional[DirectAccessClient] = None

        if isinstance(self.resources, DirectAccessResources):
            self.direct_access_client = self.resources.direct_access_client

    def create_solver(
        self,
        name: str,
        size: str | None = None,
        settings: dict | None = None,
        auto_suspend_mins: int | None = None,
    ):
        if size is None:
            size = "HIGHMEM_X64_S"
        if settings is None:
            settings = {}
        engine_config: dict[str, Any] = {"settings": settings}
        if auto_suspend_mins is not None:
            engine_config["auto_suspend_mins"] = auto_suspend_mins
        self.resources._exec_sql(
            f"CALL {APP_NAME}.experimental.create_engine('{ENGINE_TYPE_SOLVER}', '{name}', '{size}', {engine_config});",
            None,
        )

    def create_solver_async(
        self,
        name: str,
        size: str | None = None,
        settings: dict | None = None,
        auto_suspend_mins: int | None = None,
    ):
        if size is None:
            size = "HIGHMEM_X64_S"

        if self.direct_access_client is not None:
            payload: dict[str, Any] = {
                "name": name,
                "settings": settings,
            }
            if auto_suspend_mins is not None:
                payload["auto_suspend_mins"] = auto_suspend_mins
            if size is not None:
                payload["size"] = size
            response = self.direct_access_client.request(
                "create_engine",
                payload=payload,
                path_params={"engine_type": "solver"},
            )
            if response.status_code != 200:
                raise ResponseStatusException(
                    f"Failed to create engine {name} with size {size}.", response
                )
        else:
            engine_config: dict[str, Any] = {}
            if settings is not None:
                engine_config["settings"] = settings
            if auto_suspend_mins is not None:
                engine_config["auto_suspend_mins"] = auto_suspend_mins
            self.resources._exec_sql(
                f"CALL {APP_NAME}.experimental.create_engine_async('{ENGINE_TYPE_SOLVER}', '{name}', '{size}', {engine_config});",
                None,
            )

    def delete_solver(self, name: str):
        if self.direct_access_client is not None:
            response = self.direct_access_client.request(
                "delete_engine",
                path_params={"engine_type": ENGINE_TYPE_SOLVER, "engine_name": name},
            )
            if response.status_code != 200:
                raise ResponseStatusException("Failed to delete engine.", response)
            return None
        else:
            self.resources._exec_sql(
                f"CALL {APP_NAME}.experimental.delete_engine('{ENGINE_TYPE_SOLVER}', '{name}');",
                None,
            )

    def resume_solver_async(self, name: str):
        if self.direct_access_client is not None:
            response = self.direct_access_client.request(
                "resume_engine",
                path_params={"engine_type": ENGINE_TYPE_SOLVER, "engine_name": name},
            )
            if response.status_code != 200:
                # If the engine is already resuming, treat as a no-op.
                try:
                    msg = response.json().get("message", "")
                except (ValueError, AttributeError):
                    msg = response.text
                if "resume in progress" in msg.lower():
                    logger.debug(
                        "Engine %s resume already in progress, treating as no-op.", name
                    )
                    return None
                raise ResponseStatusException("Failed to resume engine.", response)
            return None
        else:
            try:
                self.resources._exec_sql(
                    f"CALL {APP_NAME}.experimental.resume_engine_async('{ENGINE_TYPE_SOLVER}', '{name}');",
                    None,
                )
            except Exception as e:
                # If the engine is already resuming, treat as a no-op.
                if "resume in progress" in str(e).lower():
                    logger.debug(
                        "Engine %s resume already in progress, treating as no-op.", name
                    )
                    return None
                raise
            return None

    def get_solver(self, name: str):
        if self.direct_access_client is not None:
            response = self.direct_access_client.request(
                "get_engine",
                path_params={"engine_type": ENGINE_TYPE_SOLVER, "engine_name": name},
            )
            if response.status_code != 200:
                raise ResponseStatusException("Failed to get engine.", response)
            solver = response.json()
            if not solver:
                return None
            solver_state = {
                "name": solver["name"],
                "id": solver["id"],
                "size": solver["size"],
                "state": solver["status"],  # callers are expecting 'state'
                "created_by": solver["created_by"],
                "created_on": solver["created_on"],
                "updated_on": solver["updated_on"],
                "version": solver["version"],
                "auto_suspend": solver["auto_suspend_mins"],
                "suspends_at": solver["suspends_at"],
                "solvers": []
                if solver["settings"] == ""
                else [
                    k
                    for (k, v) in json.loads(solver["settings"]).items()
                    if isinstance(v, dict) and v.get("enabled", False)
                ],
            }
            return solver_state
        else:
            results = self.resources._exec_sql(
                f"CALL {APP_NAME}.experimental.get_engine('{ENGINE_TYPE_SOLVER}', '{name}');",
                None,
            )
            return solver_list_to_dicts(results)[0]

    def list_solvers(self, state: str | None = None):
        if self.direct_access_client is not None:
            response = self.direct_access_client.request("list_engines")
            if response.status_code != 200:
                raise ResponseStatusException("Failed to list engines.", response)
            response_content = response.json()
            if not response_content:
                return []
            engines = [
                {
                    "name": engine["name"],
                    "id": engine["id"],
                    "size": engine["size"],
                    "state": engine["status"],  # callers are expecting 'state'
                    "created_by": engine["created_by"],
                    "created_on": engine["created_on"],
                    "updated_on": engine["updated_on"],
                    "auto_suspend_mins": engine["auto_suspend_mins"],
                    "solvers": []
                    if engine["settings"] == ""
                    else [
                        k
                        for (k, v) in json.loads(engine["settings"]).items()
                        if isinstance(v, dict) and v.get("enabled", False)
                    ],
                }
                for engine in response_content.get("engines", [])
                if (state is None or engine.get("status") == state)
                and (engine.get("type") == ENGINE_TYPE_SOLVER)
            ]
            return sorted(engines, key=lambda x: x["name"])
        else:
            where_clause = f"WHERE TYPE='{ENGINE_TYPE_SOLVER}'"
            where_clause = (
                f"{where_clause} AND STATUS = '{state.upper()}'"
                if state
                else where_clause
            )
            statement = f"SELECT NAME,ID,SIZE,STATUS,CREATED_BY,CREATED_ON,UPDATED_ON,AUTO_SUSPEND_MINS,SETTINGS FROM {APP_NAME}.experimental.engines {where_clause};"
            results = self.resources._exec_sql(statement, None)
            return solver_list_to_dicts(results)

    # --------------------------------------------------
    # Job API
    # --------------------------------------------------

    def create_job_async(
        self, engine_name, payload, query_timeout_mins: Optional[int] = None
    ):
        payload_json = json.dumps(payload)

        if (
            query_timeout_mins is None
            and (
                timeout_value := self.resources.config.get(
                    "query_timeout_mins", DEFAULT_QUERY_TIMEOUT_MINS
                )
            )
            is not None
        ):
            query_timeout_mins = int(timeout_value)

        if self.direct_access_client is not None:
            job = {
                "job_type": ENGINE_TYPE_SOLVER,
                "worker_name": engine_name,
                "timeout_mins": query_timeout_mins,
                "payload": payload_json,
            }
            response = self.direct_access_client.request(
                "create_job",
                payload=job,
            )
            if response.status_code != 200:
                raise ResponseStatusException("Failed to create job.", response)
            response_content = response.json()
            return response_content["id"]
        else:
            if query_timeout_mins is not None:
                sql_string = textwrap.dedent(f"""
                CALL {APP_NAME}.experimental.exec_job_async('{ENGINE_TYPE_SOLVER}', '{engine_name}', '{payload_json}', null, {query_timeout_mins})
                """)
            else:
                sql_string = textwrap.dedent(f"""
                CALL {APP_NAME}.experimental.exec_job_async('{ENGINE_TYPE_SOLVER}', '{engine_name}', '{payload_json}')
                """)
            res = self.resources._exec_sql(sql_string, None)
            return res[0]["ID"]

    def list_jobs(self, state=None, limit=None):
        if self.direct_access_client is not None:
            response = self.direct_access_client.request("list_jobs")
            if response.status_code != 200:
                raise ResponseStatusException("Failed to list jobs.", response)
            response_content = response.json()
            if not response_content:
                return []
            jobs = [
                {
                    "id": job["id"],
                    "state": job["state"],
                    "created_by": job["created_by"],
                    "created_on": job["created_on"],
                    "finished_at": job.get("finished_at", None),
                    "duration": job["duration"] if "duration" in job else 0,
                    "solver": json.loads(job["payload"]).get("solver", ""),
                    "engine": job.get("engine_name", job["worker_name"]),
                }
                for job in response_content.get("jobs", [])
                if state is None or job.get("state") == state
            ]
            return sorted(jobs, key=lambda x: x["created_on"], reverse=True)
        else:
            state_clause = f"AND STATE = '{state.upper()}'" if state else ""
            limit_clause = f"LIMIT {limit}" if limit else ""
            results = self.resources._exec_sql(
                f"SELECT ID,STATE,CREATED_BY,CREATED_ON,FINISHED_AT,DURATION,PAYLOAD,ENGINE_NAME FROM {APP_NAME}.experimental.jobs where type='{ENGINE_TYPE_SOLVER}' {state_clause} ORDER BY created_on DESC {limit_clause};",
                None,
            )
            return job_list_to_dicts(results)

    def get_job(self, id: str):
        if self.direct_access_client is not None:
            response = self.direct_access_client.request(
                "get_job", path_params={"job_type": ENGINE_TYPE_SOLVER, "job_id": id}
            )
            if response.status_code != 200:
                raise ResponseStatusException("Failed to get job.", response)
            response_content = response.json()
            return response_content["job"]
        else:
            results = self.resources._exec_sql(
                f"CALL {APP_NAME}.experimental.get_job('{ENGINE_TYPE_SOLVER}', '{id}');",
                None,
            )
            return job_list_to_dicts(results)[0]

    def get_job_events(self, job_id: str, continuation_token: str = ""):
        if self.direct_access_client is not None:
            response = self.direct_access_client.request(
                "get_job_events",
                path_params={
                    "job_type": ENGINE_TYPE_SOLVER,
                    "job_id": job_id,
                    "stream_name": "progress",
                },
                query_params={"continuation_token": continuation_token},
            )
            if response.status_code != 200:
                raise ResponseStatusException("Failed to get job events.", response)
            response_content = response.json()
            if not response_content:
                return {"events": [], "continuation_token": None}
            return response_content
        else:
            results = self.resources._exec_sql(
                f"SELECT {APP_NAME}.experimental.get_job_events('{ENGINE_TYPE_SOLVER}', '{job_id}', '{continuation_token}');",
                None,
            )
            if not results:
                return {"events": [], "continuation_token": None}
            row = results[0][0]
            if not isinstance(row, str):
                row = ""
            return json.loads(row)

    def cancel_job(self, id: str):
        if self.direct_access_client is not None:
            response = self.direct_access_client.request(
                "cancel_job", path_params={"job_type": ENGINE_TYPE_SOLVER, "job_id": id}
            )
            if response.status_code != 200:
                raise ResponseStatusException("Failed to cancel job.", response)
            return None
        else:
            self.resources._exec_sql(
                f"CALL {APP_NAME}.experimental.cancel_job('{ENGINE_TYPE_SOLVER}', '{id}');",
                None,
            )
            return None


def solver_list_to_dicts(results):
    if not results:
        return []
    return [
        {
            "name": row["NAME"],
            "id": row["ID"],
            "size": row["SIZE"],
            "state": row["STATUS"],  # callers are expecting 'state'
            "created_by": row["CREATED_BY"],
            "created_on": row["CREATED_ON"],
            "updated_on": row["UPDATED_ON"],
            "auto_suspend_mins": row["AUTO_SUSPEND_MINS"],
            "solvers": []
            if row["SETTINGS"] == ""
            else [
                k
                for (k, v) in json.loads(row["SETTINGS"]).items()
                if isinstance(v, dict) and v.get("enabled", False)
            ],
        }
        for row in results
    ]


def job_list_to_dicts(results):
    if not results:
        return []
    return [
        {
            "id": row["ID"],
            "state": row["STATE"],
            "created_by": row["CREATED_BY"],
            "created_on": row["CREATED_ON"],
            "finished_at": row["FINISHED_AT"],
            "duration": row["DURATION"] if "DURATION" in row else 0,
            "solver": json.loads(row["PAYLOAD"])["solver"],
            "engine": row["ENGINE_NAME"],
            "abort_reason": row["ABORT_REASON"] if "ABORT_REASON" in row else "",
            "timeout_ms": row["TIMEOUT_MS"] if "TIMEOUT_MS" in row else 0,
        }
        for row in results
    ]
