from __future__ import annotations

import re
import time
import warnings
from collections.abc import Callable
from datetime import datetime
from functools import wraps
from typing import Any, Literal, TypeVar, overload

import pandas as pd
import pyarrow as pa
from typing_extensions import ParamSpec

from dbnl import api, config
from dbnl.errors import (
    DBNLConflictingProjectError,
    DBNLDuplicateError,
    DBNLInputValidationError,
    DBNLInvalidDataError,
    DBNLNotLoggedInError,
    DBNLResourceNotFoundError,
    DBNLRunError,
    DBNLUnsupportedDataTypeError,
    WaitTimeoutError,
)
from dbnl.logging import logger
from dbnl.sdk import semconv
from dbnl.sdk.models import Filter, LLMModel, Metric, Project, Run
from dbnl.sdk.spans import (
    convert_otlp_traces_data,
    extract_input_from_spans,
    extract_output_from_spans,
    extract_timestamp_from_spans,
    spans_from_otlp,
)
from dbnl.sdk.types import Field, Schema, String, TimestampTZ, from_arrow
from dbnl.sdk.util import generate_app_url


def login(
    *,
    api_token: str | None = None,
    api_url: str | None = None,
    app_url: str | None = None,
    verify: bool = True,
) -> None:
    """
    Setup dbnl SDK to make authenticated requests. After login is run successfully, the dbnl client
    will be able to issue secure and authenticated requests against hosted endpoints of the dbnl service.


    :param api_token: dbnl API token for authentication; token can be found at /tokens page of the dbnl app.
        If None is provided, the environment variable `DBNL_API_TOKEN` will be used by default.
    :param namespace_id: The namespace ID to use for the session.
    :param api_url: The base url of the Distributional API. By default, this is set to localhost:8080/api, for sandbox users.
        For other users, please contact your sys admin. If None is provided, the environment variable `DBNL_API_URL` will be used by default.
    :param app_url: An optional base url of the Distributional app. If this variable is not set, the app url is inferred from the DBNL_API_URL
        variable. For on-prem users, please contact your sys admin if you cannot reach the Distributional UI.
    """
    config.load(
        api_token=api_token,
        api_url=api_url,
        app_url=app_url,
    )
    logger.setLevel(config.log_level())
    if verify:
        api.ensure_valid_token()
        api.maybe_warn_invalid_version()


T = TypeVar("T")
P = ParamSpec("P")


def validate_login(func: Callable[P, T]) -> Callable[P, T]:
    """
    Decorator to validate that the user has logged in before making a request
    """

    @wraps(func)
    def wrapper(*args: P.args, **kwargs: P.kwargs) -> T:
        if not config.loaded():
            raise DBNLNotLoggedInError()
        return func(*args, **kwargs)

    return wrapper


@overload
def log(
    *,
    project_id: str,
    data: pd.DataFrame,
    data_start_time: datetime,
    data_end_time: datetime,
    wait_timeout: float | None = 600,
) -> None: ...


@overload
def log(
    *,
    project_id: str,
    data_start_time: datetime,
    data_end_time: datetime,
    otlp_data: pd.Series,
    otlp_format: Literal["otlp_json", "otlp_proto"] | None = None,
    wait_timeout: float | None = 600,
) -> None: ...


@validate_login
def log(
    *,
    project_id: str,
    data: pd.DataFrame | None = None,
    data_start_time: datetime,
    data_end_time: datetime,
    wait_timeout: float | None = 600,
    otlp_data: pd.Series | None = None,
    otlp_format: Literal["otlp_json", "otlp_proto"] | None = None,
) -> None:
    """
    Log data for a date range to a project.

    :param project: The Project id to send the logs to.
    :param data: Pandas DataFrame with the data. See the DBNL Semantic Convetion.
    :param data_start_time: Data start date.
    :param data_end_time: Data end time.
    :param wait_timeout: If set, the function will block for up to `wait_timeout` seconds until the
        data is done processing, defaults to 10 minutes.

    :raises DBNLNotLoggedInError: dbnl SDK is not logged in
    :raises DBNLInputValidationError: Input does not conform to expected format

    #### Examples:

    .. code-block:: python

        from datetime import datetime, UTC

        import dbnl
        import pandas as pd

        dbnl.login()

        proj = dbnl.get_or_create_project(name="test")
        test_data = pd.DataFrame(
            {"timestamp": [datetime(2025, 1, 1, 11, 39, 53, tzinfo=UTC)]}
            {"input": ["Hello"]}
            {"output": ["World"]}
        )

        run = dbnl.log(
            project=proj,
            column_data=test_data,
            data_start_time=datetime(2025, 1, 1, tzinfo=UTC),
            data_end_time=datetime(2025, 1, 2, tzinfo=UTC),
        )
    """
    if data is None and otlp_data is None:
        raise DBNLInputValidationError("At least one of data and otlp_data must be provided")
    if data is not None and otlp_data is not None:
        raise DBNLInputValidationError("Only one of data and otlp_data can be provided")

    if data is not None:
        _log_data(project_id, data, data_start_time, data_end_time, wait_timeout)

    if otlp_data is not None:
        _log_otlp_data(project_id, otlp_data, data_start_time, data_end_time, otlp_format, wait_timeout)


def _log_data(
    project_id: str,
    data: pd.DataFrame,
    data_start_time: datetime,
    data_end_time: datetime,
    wait_timeout: float | None = 600,
) -> None:
    data = _convert_traces_data(data)
    required_columns = [semconv.Columns.TIMESTAMP, semconv.Columns.INPUT, semconv.Columns.OUTPUT]
    for c in required_columns:
        if c.value not in data.columns:
            raise DBNLInvalidDataError(f"DataFrame is missing required '{c.value}' column")
    schema = _get_schema_from_dataframe(data)
    data = _cast_dataframe(data, schema)
    _check_timestamp(data, data_start_time, data_end_time)
    run = _create_run(project_id, data_start_time, data_end_time)
    api.post_results(run.id, data)
    if "spans" in data.columns:
        api.post_spans(run.id, data["spans"].list.flatten().struct.explode())
    _close_run(run.id, wait_timeout)


def _log_otlp_data(
    project_id: str,
    data: pd.Series,
    data_start_time: datetime,
    data_end_time: datetime,
    format: Literal["otlp_json", "otlp_proto"] | None = None,
    wait_timeout: float | None = 600,
) -> None:
    spans = spans_from_otlp(data, format)
    run = _create_run(project_id, data_start_time, data_end_time)
    api.post_spans(run.id, spans)
    _close_run(run.id, wait_timeout)


def _get_schema_from_dataframe(data: pd.DataFrame) -> Schema:
    fields = []
    for f in pa.Schema.from_pandas(data):
        field = semconv.field(f.name)
        if field is not None:
            # If field is in semconv, force semconv type.
            fields.append(field)
        else:
            # Otherwise, infer type.
            try:
                fields.append(Field(f.name, from_arrow(f.type)))
            except ValueError as e:
                raise DBNLUnsupportedDataTypeError(f"Unsupported field '{f.name}' with type '{f.type}'") from e
    return Schema(fields=fields)


def _check_timestamp(data: pd.DataFrame, data_start_time: datetime, data_end_time: datetime) -> None:
    timestamp_column = semconv.Columns.TIMESTAMP.value
    if timestamp_column not in data.columns:
        raise DBNLInvalidDataError(f"DataFrame is missing required '{timestamp_column}' column")
    if data.empty:
        return
    if data[timestamp_column].isnull().any():
        raise DBNLInvalidDataError(f"DataFrame contains null values in '{timestamp_column}' column")
    if data[timestamp_column].min() < data_start_time:
        raise DBNLInvalidDataError(
            f"DataFrame contains timestamps before data_start_time: {data[timestamp_column].min()} < {data_start_time.isoformat()}"
        )
    if data[timestamp_column].max() >= data_end_time:
        raise DBNLInvalidDataError(
            f"DataFrame contains timestamps after or equal to data_end_time: {data[timestamp_column].max()} >= {data_end_time.isoformat()}"
        )


def _convert_traces_data(data: pd.DataFrame) -> pd.DataFrame:
    if not "traces_data" in data.columns:
        return data
    if "spans" in data.columns:
        warnings.warn("DataFrame contains both `traces_data` and `spans` columns. `spans` column will be used.")
        return data.drop("traces_data", axis=1)
    spans = convert_otlp_traces_data(data["traces_data"])
    if len(spans) != len(data):
        raise DBNLInvalidDataError(
            f"Expected `traces_data` column to contain exactly one trace per row, found {len(spans)} traces across {len(data)} rows"
        )
    data["spans"] = spans
    if not "timestamp" in data.columns:
        data["timestamp"] = data["spans"].map(extract_timestamp_from_spans).astype(TimestampTZ().to_pandas())
    if not "input" in data.columns:
        data["input"] = data["spans"].map(extract_input_from_spans).astype(String().to_pandas())
    if not "output" in data.columns:
        data["output"] = data["spans"].map(extract_output_from_spans).astype(String().to_pandas())
    return data.drop("traces_data", axis=1)


def _cast_dataframe(data: pd.DataFrame, schema: Schema) -> pd.DataFrame:
    try:
        return data.astype(schema.to_pandas())
    except ValueError as e:
        m = re.search("Error while type casting for column '(\w+)'", str(e))
        if m is not None and len(m.groups()) == 1:
            col = m.groups()[0]
            f = schema.field(col)
            if f is not None:
                raise DBNLInvalidDataError(f"Expected '{col}' with type '{f.type}': {str(e)}") from e
        raise DBNLInvalidDataError(str(e)) from e


def _create_run(
    project_id: str,
    data_start_time: datetime | None,
    data_end_time: datetime | None,
) -> Run:
    resp_dict = api.create_run(
        project_id=project_id,
        data_start_time=data_start_time,
        data_end_time=data_end_time,
    )
    return Run.from_dict(resp_dict)


def _get_run(run_id: str) -> Run:
    resp_dict = api.get_run(run_id=run_id)
    return Run.from_dict(resp_dict)


def _close_run(
    run_id: str,
    wait_timeout: float | None = 600,
) -> Run:
    resp_dict = api.close_run(run_id=run_id)
    run = Run.from_dict(resp_dict)
    logger.info(
        "Initiated run close. View results at: %s",
        generate_app_url(f"ns/{run.namespace_id}/projects/{run.project_id}/status"),
    )
    if wait_timeout is not None:
        return _wait_for_run_close(run_id=run_id, timeout=wait_timeout)
    return run


def _wait_for_run_close(run_id: str, timeout: float = 600, polling_interval: float = 5) -> Run:
    logger.info("Waiting for run to close...")
    start = time.time()
    run = _get_run(run_id=run_id)
    while run.completed_at is None and (time.time() - start) < timeout:
        time.sleep(polling_interval)
        run = _get_run(run_id=run_id)
    if run.completed_at is None:
        project_status_url = generate_app_url(f"ns/{run.namespace_id}/projects/{run.project_id}/status")
        raise WaitTimeoutError(
            f"Run [{run.data_start_time}-{run.data_end_time}] hasn't finished after waiting for {timeout} seconds."
            f" See {project_status_url} to check on run status and consider extending or removing `wait_timeout` to prevent this exception from being thrown."
        )
    if run.status != "closed":
        raise DBNLRunError(f"Run {run} ended in '{run.status}' state: {run.failure}")
    logger.info("Run closed successfully")
    return run


@overload
def get_project(
    *,
    project_id: str,
) -> Project:
    """
    Retrieve a Project by id.

    :param id: The id for the existing Project.

    :raises DBNLNotLoggedInError: dbnl SDK is not logged in
    :raises DBNLProjectNotFoundError: Project with the given name does not exist.

    :return: Project

    #### Examples:

    .. code-block:: python

        import dbnl

        dbnl.login()

        proj_1 = dbnl.create_project(name="test_p1")
        proj_2 = dbnl.get_project(project_id=proj_1.id)

        # Calling get_project will yield same Project object
        assert proj_1.id == proj_2.id

        # DBNLProjectNotFoundError: A dbnl Project with id not_exist does not exist
        proj_3 = dbnl.get_project(project_id="not_exist")
    """


@overload
def get_project(
    *,
    name: str,
) -> Project:
    """
    Retrieve a Project by name.

    :param name: The name for the existing Project.

    :raises DBNLNotLoggedInError: dbnl SDK is not logged in
    :raises DBNLProjectNameNotFoundError: Project with the given name does not exist.

    :return: Project

    #### Examples:

    .. code-block:: python

        import dbnl

        dbnl.login()

        proj_1 = dbnl.create_project(name="test_p1")
        proj_2 = dbnl.get_project_by_name(name="test_p1")

        # Calling get_project_by_name will yield same Project object
        assert proj_1.id == proj_2.id

        # DBNLProjectNameNotFoundError: A dbnl Project with name not_exist does not exist
        proj_3 = dbnl.get_project_by_name(name="not_exist")
    """


@validate_login
def get_project(
    *,
    project_id: str | None = None,
    name: str | None = None,
) -> Project:
    if project_id is not None and name is not None:
        raise DBNLInputValidationError("Only one of `project_id` and `name` can be provided")

    if name is not None:
        try:
            resp_dict = api.get_project_by_name(name=name)
        except DBNLResourceNotFoundError:
            resp_dict = api.get_project_by_name(name=name, archived=True)
        return Project.from_dict(resp_dict)

    if project_id is not None:
        resp_dict = api.get_project(project_id=project_id)
        return Project.from_dict(resp_dict)

    raise DBNLInputValidationError("One of `project_id` and `name` must be provided")


@validate_login
def create_project(
    *,
    name: str,
    description: str | None = None,
    default_llm_model_id: str | None = None,
    default_llm_model_name: str | None = None,
    template: Literal["default"] | None = "default",
) -> Project:
    """
    Create a new Project

    :param name: Name for the Project
    :param description: Description for the Project, defaults to None. Description is limited to 255 characters.

    :raises DBNLNotLoggedInError: dbnl SDK is not logged in
    :raises DBNLAPIValidationError: dbnl API failed to validate the request
    :raises DBNLConflictingProjectError: Project with the same name already exists

    :return: Project

    #### Examples:

    .. code-block:: python

        import dbnl

        dbnl.login()

        proj_1 = dbnl.create_project(name="test_p1")

        # DBNLConflictingProjectError: A Project with name test_p1 already exists.
        proj_2 = dbnl.create_project(name="test_p1")
    """
    if default_llm_model_id is not None and default_llm_model_name is not None:
        raise DBNLInputValidationError("Only one of llm_model_id and llm_model_name can be provided")

    if default_llm_model_name is not None:
        default_llm_model_id = get_llm_model(name=default_llm_model_name).id

    try:
        resp_dict = api.create_project(
            name=name,
            description=description,
            schedule="daily",
            default_llm_model_id=default_llm_model_id,
            template=template,
        )
    except DBNLDuplicateError:
        raise DBNLConflictingProjectError(name)

    project = Project.from_dict(resp_dict)

    logger.info(
        "View Project %s at: %s",
        project.name,
        generate_app_url(f"ns/{project.namespace_id}/projects/{project.id}"),
    )

    return project


@validate_login
def get_or_create_project(
    *,
    name: str,
    description: str | None = None,
    default_llm_model_id: str | None = None,
    default_llm_model_name: str | None = None,
    template: Literal["default"] | None = "default",
) -> Project:
    """
    Get the Project with the specified name or create a new one if it does not exist

    :param name: Name for the Project
    :param description: Description for the Project, defaults to None

    :raises DBNLNotLoggedInError: dbnl SDK is not logged in
    :raises DBNLAPIValidationError: dbnl API failed to validate the request

    :return: Newly created or matching existing Project

    #### Examples:

    .. code-block:: python

        import dbnl

        dbnl.login()

        proj_1 = dbnl.create_project(name="test_p1")
        proj_2 = dbnl.get_or_create_project(name="test_p1")

        # Calling get_or_create_project will yield same Project object
        assert proj_1.id == proj_2.id
    """

    try:
        return get_project(name=name)
    except DBNLResourceNotFoundError:
        try:
            return create_project(
                name=name,
                description=description,
                default_llm_model_id=default_llm_model_id,
                default_llm_model_name=default_llm_model_name,
                template=template,
            )
        except DBNLConflictingProjectError:
            return get_project(name=name)


@overload
def get_or_create_metric(
    *,
    project_id: str,
    name: str,
    table: Literal["spans", "traces", "sessions"] = "traces",
    expression: str,
    description: str | None = None,
    greater_is_better: bool | None = None,
) -> Metric:
    """
    Get a Metric by name, or create it if it does not exist.

    :param project_id: The Project ID to get the Metric for
    :param name: Name of the Metric to get
    :param table: Table to get the Metric for
    :param expression: Expression string e.g. `length(traces.input)`
    :param description: Optional description of what computation the metric is performing
    :param greater_is_better: Flag indicating whether greater values are semantically 'better' than lesser values
    """


@overload
def get_or_create_metric(
    *,
    project: Project,
    name: str,
    table: Literal["spans", "traces", "sessions"] = "traces",
    expression: str,
    description: str | None = None,
    greater_is_better: bool | None = None,
) -> Metric:
    """
    Get a Metric by name, or create it if it does not exist.

    :param project: The Project to get the Metric for
    :param name: Name of the Metric to get
    :param table: Table to get the Metric for
    :param expression: Expression string e.g. `length(traces.input)`
    :param description: Optional description of what computation the metric is performing
    :param greater_is_better: Flag indicating whether greater values are semantically 'better' than lesser values
    """


@validate_login
def get_or_create_metric(
    *,
    project: Project | None = None,
    project_id: str | None = None,
    name: str,
    table: Literal["spans", "traces", "sessions"] = "traces",
    expression: str,
    description: str | None = None,
    greater_is_better: bool | None = None,
) -> Metric:
    if project is not None and project_id is not None:
        raise DBNLInputValidationError("Only one of `project` and `project_id` can be provided")

    if project is not None:
        warnings.warn("`project` is deprecated. Use `project_id` instead.")
        project_id = project.id
    if project_id is None:
        raise DBNLInputValidationError("One of `project` or `project_id` must be provided")

    try:
        return get_metric(name=name)
    except DBNLResourceNotFoundError:
        return create_metric(
            project_id=project_id,
            name=name,
            table=table,
            expression=expression,
            description=description,
            greater_is_better=greater_is_better,
        )


@overload
def create_metric(
    *,
    project_id: str,
    name: str,
    table: Literal["spans", "traces", "sessions"] = "traces",
    expression: str,
    description: str | None = None,
    greater_is_better: bool | None = None,
) -> Metric:
    """
    Create a new Metric

    :param project_id: The Project ID to create the Metric for
    :param name: Name for the Metric
    :param table: Table to create the Metric for
    :param expression: Expression string e.g. `length(traces.input)`
    :param description: Optional description of what computation the metric is performing
    :param greater_is_better: Flag indicating whether greater values are semantically 'better' than lesser values

    :raises DBNLNotLoggedInError: dbnl SDK is not logged in
    :raises DBNLInputValidationError: Input does not conform to expected format

    :return: Created Metric
    """


@overload
def create_metric(
    *,
    project: Project,
    name: str,
    table: Literal["spans", "traces", "sessions"] = "traces",
    expression: str,
    description: str | None = None,
    greater_is_better: bool | None = None,
) -> Metric:
    """
    Create a new Metric

    :param project: The Project to create the Metric for
    :param name: Name for the Metric
    :param table: Table to create the Metric for
    :param expression: Expression string e.g. `length(traces.input)`
    :param description: Optional description of what computation the metric is performing
    :param greater_is_better: Flag indicating whether greater values are semantically 'better' than lesser values

    :raises DBNLNotLoggedInError: dbnl SDK is not logged in
    :raises DBNLInputValidationError: Input does not conform to expected format

    :return: Created Metric
    """


@validate_login
def create_metric(
    *,
    project: Project | None = None,
    project_id: str | None = None,
    name: str,
    table: Literal["spans", "traces", "sessions"] = "traces",
    expression: str,
    description: str | None = None,
    greater_is_better: bool | None = None,
) -> Metric:
    """
    Create a new Metric

    :param project: The Project to create the Metric for
    :param name: Name for the Metric
    :param table: Table to create the Metric for
    :param expression: Expression string e.g. `length(traces.input)`
    :param description: Optional description of what computation the metric is performing
    :param greater_is_better: Flag indicating whether greater values are semantically 'better' than lesser values

    :raises DBNLNotLoggedInError: dbnl SDK is not logged in
    :raises DBNLInputValidationError: Input does not conform to expected format

    :return: Created Metric
    """
    if project is not None and project_id is not None:
        raise DBNLInputValidationError("Only one of `project` and `project_id` can be provided")

    if project is not None:
        warnings.warn("`project` is deprecated. Use `project_id` instead.")
        project_id = project.id
    if project_id is None:
        raise DBNLInputValidationError("One of `project` or `project_id` must be provided")

    resp_dict = api.create_metric(
        project_id=project_id,
        name=name,
        table=table,
        expression=expression,
        description=description,
        greater_is_better=greater_is_better,
    )
    return Metric.from_dict(resp_dict)


@overload
def get_metric(
    *,
    metric_id: str,
) -> Metric:
    """
    Get a Metric by ID

    :param metric_id: ID of the metric to get

    :raises DBNLNotLoggedInError: dbnl SDK is not logged in
    :raises DBNLAPIValidationError: dbnl API failed to validate the request

    :return: The requested metric
    """


@overload
def get_metric(
    *,
    name: str,
) -> Metric:
    """
    Get a Metric by name

    :param name: Name of the metric to get

    :raises DBNLNotLoggedInError: dbnl SDK is not logged in
    :raises DBNLAPIValidationError: dbnl API failed to validate the request

    :return: The requested metric
    """


@validate_login
def get_metric(
    *,
    metric_id: str | None = None,
    name: str | None = None,
) -> Metric:
    if metric_id is not None and name is not None:
        raise DBNLInputValidationError("Only one of `metric_id` and `name` can be provided")

    if metric_id is not None:
        resp_dict = api.get_metric(metric_id=metric_id)
        return Metric.from_dict(resp_dict)

    if name is not None:
        resp_dict = api.get_metric_by_name(name=name)
        return Metric.from_dict(resp_dict)

    raise DBNLInputValidationError("One of `metric_id` and `name` must be provided")


@validate_login
def update_metric(
    *,
    metric_id: str,
    name: str | None = None,
    expression: str | None = None,
    description: str | None = None,
    greater_is_better: bool | None = None,
) -> Metric:
    resp_dict = api.update_metric(
        metric_id=metric_id,
        name=name,
        expression=expression,
        description=description,
        greater_is_better=greater_is_better,
    )
    return Metric.from_dict(resp_dict)


@validate_login
def delete_metric(
    *,
    metric_id: str,
) -> None:
    """
    Delete a Metric by ID

    :param metric_id: ID of the metric to delete

    :raises DBNLNotLoggedInError: dbnl SDK is not logged in
    :raises DBNLAPIValidationError: dbnl API failed to validate the request

    :return: None
    """
    api.delete_metric(metric_id=metric_id)


@validate_login
def get_or_create_llm_model(
    *,
    name: str,
    description: str | None = None,
    type: Literal["completion", "embedding"] | None = "completion",
    provider: str,
    model: str,
    params: dict[str, Any] = {},
) -> LLMModel:
    """
    Get an LLM Model by name, or create it if it does not exist.

    :param name: Model name
    :param description: Model description, defaults to None
    :param type: Model type (e.g. completion or embedding), defaults to "completion"
    :param provider: Model provider (e.g. openai, bedrock, etc.)
    :param model: Model (e.g. gpt-4, gpt-3.5-turbo, etc.)
    :param params: Model provider parameters (e.g. api key), defaults to {}

    :return: Model
    """
    try:
        return get_llm_model(name=name)
    except DBNLResourceNotFoundError:
        return create_llm_model(
            name=name,
            description=description,
            type=type,
            provider=provider,
            model=model,
            params=params,
        )


@validate_login
def create_llm_model(
    *,
    name: str,
    description: str | None = None,
    type: Literal["completion", "embedding"] | None = "completion",
    provider: str,
    model: str,
    params: dict[str, Any] = {},
) -> LLMModel:
    """
    Create an LLM Model.

    :param name: Model name
    :param description: Model description, defaults to None
    :param type: Model type (e.g. completion or embedding), defaults to "completion"
    :param provider: Model provider (e.g. openai, bedrock, etc.)
    :param model: Model (e.g. gpt-4, gpt-3.5-turbo, etc.)
    :param params: Model provider parameters (e.g. api key), defaults to {}

    :return: LLM Model
    """
    resp = api.create_llm_model(
        name=name,
        description=description,
        type=type,
        provider=provider,
        model=model,
        params=params,
    )
    return LLMModel.from_dict(resp)


@overload
def get_llm_model(
    *,
    llm_model_id: str,
) -> LLMModel:
    """
    Get an LLM Model by id.

    :param llm_model_id: Model id

    :return: LLM Model if found
    """


@overload
def get_llm_model(
    *,
    name: str,
) -> LLMModel:
    """
    Get an LLM Model by name

    :param name: LLM Model name

    :return: LLM Model if found
    """


@validate_login
def get_llm_model(
    *,
    llm_model_id: str | None = None,
    name: str | None = None,
) -> LLMModel:
    if llm_model_id is not None and name is not None:
        raise DBNLInputValidationError("Only one of `llm_model_id` and `name` can be provided")

    if llm_model_id is not None:
        resp = api.get_llm_model(llm_model_id=llm_model_id)
        return LLMModel.from_dict(resp)

    if name is not None:
        resp = api.get_llm_model_by_name(name=name)
        return LLMModel.from_dict(resp)

    raise DBNLInputValidationError("One of `llm_model_id` and `name` must be provided")


@validate_login
def update_llm_model(
    *,
    llm_model_id: str,
    name: str | None = None,
    description: str | None = None,
    model: str | None = None,
    params: dict[str, Any] | None = None,
) -> LLMModel:
    """
    Update an LLM Model by id.

    :param llm_model_id: Model id
    :param name: Model name
    :param description: Model description, defaults to None
    :param model: Model (e.g. gpt-4, gpt-3.5-turbo, etc.)
    :param params: Model provider parameters (e.g. api key), defaults to {}

    :return: Updated LLM Model
    """
    resp = api.update_llm_model(
        llm_model_id=llm_model_id,
        name=name,
        description=description,
        model=model,
        params=params,
    )
    return LLMModel.from_dict(resp)


@validate_login
def delete_llm_model(
    *,
    llm_model_id: str,
) -> None:
    """
    Delete an LLM Model by id.

    :param llm_model_id: LLM Model id

    :return: LLM Model if found
    """
    api.delete_llm_model(llm_model_id=llm_model_id)


@validate_login
def get_or_create_filter(
    *,
    project_id: str,
    name: str,
    table: Literal["spans", "traces", "sessions"],
    description: str | None = None,
    conditions: list[api.FilterCondition] | None = None,
    expression: str | None = None,
) -> Filter:
    """
    Get a Filter by name, or create it if it does not exist.

    :param project_id: The Project ID to get the Filter for
    :param name: Name of the Filter to get
    :param table: Table to get the Filter for
    :param description: Optional description of the Filter
    :param conditions: Conditions for the Filter
    :param expression: Expression string e.g. `length(traces.input) > 10`

    :return: Filter
    """
    try:
        return get_filter(name=name)
    except DBNLResourceNotFoundError:
        return create_filter(
            project_id=project_id,
            name=name,
            table=table,
            description=description,
            conditions=conditions or [],
            expression=expression,
        )


@validate_login
def create_filter(
    project_id: str,
    name: str,
    table: Literal["spans", "traces", "sessions"],
    description: str | None = None,
    conditions: list[api.FilterCondition] | None = None,
    expression: str | None = None,
) -> Filter:
    """
    Create a new Filter

    :param project_id: The Project ID to create the Filter for
    :param name: Name for the Filter
    :param table: Table to create the Filter for
    :param description: Optional description of the Filter
    :param conditions: Conditions for the Filter
    :param expression: Expression string e.g. `length(traces.input) > 10`

    :return: Created Filter
    """
    resp_dict = api.create_filter(
        project_id=project_id,
        name=name,
        table=table,
        description=description,
        conditions=conditions or [],
        expression=expression,
    )
    return Filter.from_dict(resp_dict)


@overload
def get_filter(
    *,
    name: str,
) -> Filter:
    """
    Get a Filter by name

    :param name: Filter name

    :return: Filter
    """


@overload
def get_filter(
    *,
    filter_id: str,
) -> Filter:
    """
    Get a Filter by id

    :param filter_id: Filter id

    :return: Filter
    """


@validate_login
def get_filter(
    *,
    filter_id: str | None = None,
    name: str | None = None,
) -> Filter:
    if filter_id is not None and name is not None:
        raise DBNLInputValidationError("Only one of `filter_id` and `name` can be provided")

    if filter_id is not None:
        resp_dict = api.get_filter(filter_id=filter_id)
        return Filter.from_dict(resp_dict)

    if name is not None:
        resp_dict = api.get_filter_by_name(name=name)
        return Filter.from_dict(resp_dict)

    raise DBNLInputValidationError("One of `filter_id` and `name` must be provided")


@validate_login
def update_filter(
    *,
    filter_id: str,
    name: str | None = None,
    description: str | None = None,
    conditions: list[api.FilterCondition] | None = None,
    expression: str | None = None,
) -> Filter:
    """
    Update a Filter by id

    :param filter_id: Filter id
    :param name: Filter name
    :param description: Filter description
    :param conditions: Filter conditions
    :param expression: Filter expression
    :return: Updated Filter
    """
    resp_dict = api.update_filter(
        filter_id=filter_id,
        name=name,
        description=description,
        conditions=conditions,
        expression=expression,
    )
    return Filter.from_dict(resp_dict)


@validate_login
def delete_filter(
    *,
    filter_id: str,
) -> None:
    """
    Delete a Filter by id

    :param filter_id: Filter id

    :return: None
    """
    api.delete_filter(filter_id=filter_id)
