from __future__ import annotations

from typing import Annotated
from typing import Any
from typing import Literal
from typing import TypeAlias

from amsdal_utils.models.enums import TransactionType
from pydantic import BaseModel
from pydantic import Field

from amsdal_ml.mcp_server.server_model_explorer.services.types import LiteralValue
from amsdal_ml.mcp_server.server_model_explorer.services.types import TransactionArgumentMode
from amsdal_ml.mcp_server.server_model_explorer.services.types import TransactionParameterKind
from amsdal_ml.ml_retrievers.query_creator import CreateAnalysis
from amsdal_ml.ml_retrievers.query_retriever import RetrieveAnalysis


class PlannedArgument(BaseModel):
    name: str
    mode: TransactionArgumentMode
    literal_value: LiteralValue = None
    query: str | None = None
    explicit_object_id: str | None = None
    note: str | None = None


class PlannedArguments(BaseModel):
    items: list[PlannedArgument]


class TransactionParameterSpec(BaseModel):
    name: str
    annotation: str
    kind: TransactionParameterKind
    model_name: str | None = None
    nullable: bool = False


class TransactionSpec(BaseModel):
    name: str
    description: str
    scope: TransactionType
    parameters: list[TransactionParameterSpec]


class TransactionArgBase(BaseModel):
    name: str
    kind: TransactionParameterKind
    model_name: str | None = None
    error: str | None = None


class LiteralArgAnalysis(TransactionArgBase):
    mode: Literal[TransactionArgumentMode.LITERAL] = TransactionArgumentMode.LITERAL
    value: Any | None = None


class ExistingRefArgAnalysis(TransactionArgBase):
    mode: Literal[TransactionArgumentMode.EXISTING_REF] = TransactionArgumentMode.EXISTING_REF
    query: str
    existing_analysis: RetrieveAnalysis | None = None


class NewObjectArgAnalysis(TransactionArgBase):
    mode: Literal[TransactionArgumentMode.NEW_OBJECT] = TransactionArgumentMode.NEW_OBJECT
    query: str
    create_analysis: CreateAnalysis | None = None


class ExplicitIdArgAnalysis(TransactionArgBase):
    mode: Literal[TransactionArgumentMode.EXPLICIT_ID] = TransactionArgumentMode.EXPLICIT_ID
    value: str


class ErrorArgAnalysis(TransactionArgBase):
    mode: Literal['error'] = 'error'


TransactionArgAnalysis: TypeAlias = Annotated[
    LiteralArgAnalysis | ExistingRefArgAnalysis | NewObjectArgAnalysis | ExplicitIdArgAnalysis | ErrorArgAnalysis,
    Field(discriminator='mode'),
]


class TransactionRunAnalysis(BaseModel):
    transaction_name: str
    args: list[TransactionArgAnalysis]
    is_executable: bool
    summary: str | None = None
