"""Models specific to the nimbus method."""

from typing import Literal

from sqlmodel import JSON, Column, Field, SQLModel

from .generic import SolutionInfo
from .generic_states import SolutionReferenceResponse
from .preference import ReferencePoint


class NIMBUSClassificationRequest(SQLModel):
    """Model of the request to the nimbus method."""

    problem_id: int
    session_id: int | None = Field(default=None)
    parent_state_id: int | None = Field(default=None)

    scalarization_options: dict[str, float | str | bool] | None = Field(sa_column=Column(JSON), default=None)
    solver: str | None = Field(default=None)
    solver_options: dict[str, float | str | bool] | None = Field(sa_column=Column(JSON), default=None)
    preference: ReferencePoint = Field(Column(JSON))

    current_objectives: dict[str, float] = Field(
        sa_column=Column(JSON), description="The objectives used for iteration."
    )
    num_desired: int | None = Field(default=1)


class NIMBUSSaveRequest(SQLModel):
    """Request model for saving solutions from any method's state."""

    problem_id: int
    session_id: int | None = Field(default=None)
    parent_state_id: int | None = Field(default=None)

    solution_info: list[SolutionInfo]

class NIMBUSDeleteSaveRequest(SQLModel):
    """Request model for deletion of a saved solution."""

    state_id : int = Field(description="The ID of the save state.")
    solution_index: int = Field(description="The ID of the solution within the above state.")


class NIMBUSFinalizeRequest(SQLModel):
    """Request model for finalizing the NIMBUS procedure."""

    problem_id: int
    session_id: int | None = Field(default=None)
    parent_state_id: int | None = Field(default=None)

    solution_info: SolutionInfo # the final solution


class NIMBUSClassificationResponse(SQLModel):
    """The response from NIMBUS classification endpoint."""

    response_type: Literal["nimbus.classification"] = "nimbus.classification"

    state_id: int | None = Field(description="The newly created state id")
    previous_preference: ReferencePoint = Field(description="The previous preference used.")
    previous_objectives: dict[str, float] = Field(
        sa_column=Column(JSON), description="The previous solutions objectives used for iteration."
    )
    current_solutions: list[SolutionReferenceResponse] = Field(
        description="The solutions from the current iteration of nimbus."
    )
    saved_solutions: list[SolutionReferenceResponse] = Field(
        description="The best candidate solutions saved by the decision maker."
    )
    all_solutions: list[SolutionReferenceResponse] = Field(
        description="All solutions generated by NIMBUS in all iterations."
    )


class NIMBUSInitializationResponse(SQLModel):
    """The response from NIMBUS classification endpoint."""

    response_type: Literal["nimbus.initialization"] = "nimbus.initialization"

    state_id: int | None = Field(description="The newly created state id")
    current_solutions: list[SolutionReferenceResponse] = Field(
        description="The solutions from the current interation of nimbus."
    )
    saved_solutions: list[SolutionReferenceResponse] = Field(
        description="The best candidate solutions saved by the decision maker."
    )
    all_solutions: list[SolutionReferenceResponse] = Field(
        description="All solutions generated by NIMBUS in all iterations."
    )


class NIMBUSSaveResponse(SQLModel):
    """The response from NIMBUS save endpoint."""

    response_type: Literal["nimbus.save"] = "nimbus.save"

    state_id: int | None = Field(description="The id of the newest state")

class NIMBUSDeleteSaveResponse(SQLModel):
    """Response of NIMBUS save deletion."""

    response_type: str = "nimbus.delete_save"

    message: str | None

class NIMBUSFinalizeResponse(SQLModel):
    """The response from NIMBUS finish endpoint."""

    response_type: Literal["nimbus.finalize"] = "nimbus.finalize"

    state_id: int | None = Field(description="The newly created state id")
    final_solution: SolutionReferenceResponse = Field(
        description="The final solution. We do not need the other current solutions."
    )
    saved_solutions: list[SolutionReferenceResponse] = Field(
        description="The best candidate solutions saved by the decision maker."
    )
    all_solutions: list[SolutionReferenceResponse] = Field(
        description="All solutions generated by NIMBUS in all iterations."
    )


class NIMBUSInitializationRequest(SQLModel):
    """Model of the request to the nimbus method."""

    problem_id: int
    session_id: int | None = Field(default=None)
    parent_state_id: int | None = Field(default=None)

    starting_point: ReferencePoint | SolutionInfo | None = Field(sa_column=Column(JSON), default=None)
    scalarization_options: dict[str, float | str | bool] | None = Field(sa_column=Column(JSON), default=None)
    solver: str | None = Field(default=None)
    solver_options: dict[str, float | str | bool] | None = Field(sa_column=Column(JSON), default=None)


class NIMBUSIntermediateSolutionResponse(SQLModel):
    """The response from NIMBUS classification endpoint."""

    response_type: Literal["nimbus.intermediate"] = "nimbus.intermediate"

    state_id: int | None = Field(description="The newly created state id")
    reference_solution_1: dict[str, float] = Field(
        sa_column=Column(JSON), description="The first solution used when computing intermediate points."
    )
    reference_solution_2: dict[str, float]= Field(
        sa_column=Column(JSON), description="The second solution used when computing intermediate points."
    )
    current_solutions: list[SolutionReferenceResponse] = Field(
        description="The solutions from the current iteration of NIMBUS."
    )
    saved_solutions: list[SolutionReferenceResponse] = Field(
        description="The best candidate solutions saved by the decision maker."
    )
    all_solutions: list[SolutionReferenceResponse] = Field(
        description="All solutions generated by NIMBUS in all iterations."
    )
