#  Copyright 2026 Synnax Labs, Inc.
#
#  Use of this software is governed by the Business Source License included in the file
#  licenses/BSL.txt.
#
#  As of the Change Date specified in that file, in accordance with the Business Source
#  License, use of this software will be governed by the Apache License, Version 2.0,
#  included in the file licenses/APL.txt.

from typing import overload

from freighter import Empty, UnaryClient, send_required
from pydantic import BaseModel, Field

from synnax.ontology.payload import ID, CrudeID, Resource
from synnax.util.normalize import normalize


class RetrieveReq(BaseModel):
    ids: list[ID]
    children: bool = False
    parents: bool = False
    include_schema: bool = False
    exclude_field_data: bool = False
    term: str | None = None
    limit: int | None = None
    offset: int | None = None
    types: list[str] | None = None


class AddChildrenReq(BaseModel):
    id: ID
    children: list[ID]


class RemoveChildrenReq(BaseModel):
    id: ID
    children: list[ID]


class MoveChildrenReq(BaseModel):
    from_: ID = Field(alias="from")
    to: ID
    children: list[ID]


class RetrieveRes(BaseModel):
    resources: list[Resource]


class Client:
    _client: UnaryClient

    def __init__(self, client: UnaryClient) -> None:
        self._client = client

    @overload
    def retrieve(
        self,
        id: CrudeID,
        *,
        children: bool = False,
        parents: bool = False,
        include_schema: bool = False,
        exclude_field_data: bool = False,
    ) -> Resource: ...

    @overload
    def retrieve(
        self,
        id: list[CrudeID],
        *,
        children: bool = False,
        parents: bool = False,
        include_schema: bool = False,
        exclude_field_data: bool = False,
    ) -> list[Resource]: ...

    def retrieve(
        self,
        id: CrudeID | list[CrudeID] | None = None,
        *,
        children: bool = False,
        parents: bool = False,
        include_schema: bool = False,
        exclude_field_data: bool = False,
    ) -> Resource | list[Resource]:
        is_single = False
        if not isinstance(id, list):
            if id is None:
                id = []
            else:
                id = [id]
                is_single = True
        req = RetrieveReq(
            ids=[ID(i) for i in id],
            children=children,
            parents=parents,
            include_schema=include_schema,
            exclude_field_data=exclude_field_data,
        )
        resources = self.__exec_retrieve(req)
        if is_single:
            return resources[0]
        return resources

    def retrieve_children(
        self,
        id: CrudeID | list[CrudeID],
    ) -> list[Resource]:
        normalized: list[CrudeID] = normalize(id)
        return self.__exec_retrieve(
            RetrieveReq(ids=[ID(i) for i in normalized], children=True)
        )

    def __exec_retrieve(self, req: RetrieveReq) -> list[Resource]:
        return send_required(
            self._client, "/ontology/retrieve", req, RetrieveRes
        ).resources

    def retrieve_parents(
        self,
        id: CrudeID | list[CrudeID],
    ) -> list[Resource]:
        normalized: list[CrudeID] = normalize(id)
        return self.__exec_retrieve(
            RetrieveReq(ids=[ID(i) for i in normalized], parents=True)
        )

    def move_children(self, from_: CrudeID, to: CrudeID, *children: CrudeID) -> None:

        send_required(
            self._client,
            "/ontology/move-children",
            MoveChildrenReq.model_validate(
                {"from": ID(from_), "to": ID(to), "children": [ID(i) for i in children]}
            ),
            Empty,
        )

    def remove_children(self, id: CrudeID, *children: CrudeID) -> None:
        send_required(
            self._client,
            "/ontology/remove-children",
            RemoveChildrenReq(id=ID(id), children=[ID(i) for i in children]),
            Empty,
        )

    def add_children(self, id: CrudeID, *children: CrudeID) -> None:
        send_required(
            self._client,
            "/ontology/add-children",
            AddChildrenReq(id=ID(id), children=[ID(i) for i in children]),
            Empty,
        )
