# Copyright 2026 Marimo. All rights reserved.
from __future__ import annotations

from datetime import date, datetime, time, timedelta  # noqa: TC003
from typing import TYPE_CHECKING, Any, Literal

import msgspec

from marimo._types.ids import VariableName
from marimo._utils.msgspec_basestruct import BaseStruct

DataType = Literal[
    "string",
    "boolean",
    "integer",
    "number",
    "date",
    "datetime",
    "time",
    "unknown",
]
# This is the data type based on the source library
# e.g. polars, pandas, numpy, etc.
ExternalDataType = str


class DataTableColumn(BaseStruct):
    """
    Represents a column in a data table.

    Attributes:
        name (str): The name of the column.
        type (DataType): The data type of the column.
        external_type (ExternalDataType): The raw data type of the column.
        sample_values (List[Any]): The sample values of the column.
    """

    name: str
    type: DataType
    external_type: ExternalDataType
    sample_values: list[Any]

    def __post_init__(self) -> None:
        # Sometimes libraries (like pandas, sqlalchemy or ibis) may return column names as objects
        # instead of strings, although their type hints are str
        # Instead of trying to track this down each time, just convert to string
        self.name = str(self.name)


# Local -> Python dataframes
# DuckDB -> DuckDB tables using the global in-memory DuckDB instance
# Connection -> SQL tables using a named data source connection (e.g. SQLAlchemy, or a custom DuckDB connection)
# Catalog -> Data catalog (e.g. iceberg)
DataTableSource = Literal["local", "duckdb", "connection", "catalog"]
DataTableType = Literal["table", "view"]


class DataTable(BaseStruct):
    """
    Represents a data table.

    Attributes:
        source_type (DataTableSource): Type of data source ('local', 'duckdb', 'connection').
        source (str): Can be dialect, or source db name.
        name (str): Name of the data table.
        num_rows (Optional[int]): Total number of rows in the table, if known.
        num_columns (Optional[int]): Total number of columns in the table, if known.
        variable_name (Optional[VariableName]): Variable name referencing this table in code.
        columns (List[DataTableColumn]): List of column definitions and metadata.
        engine (Optional[VariableName]): Database engine or connection handler, if any.
        type (DataTableType): Table type, either 'table' or 'view'. Defaults to 'table'.
        primary_keys (Optional[List[str]]): Column names used as primary keys, if any.
        indexes (Optional[List[str]]): Column names used as indexes, if any.
    """

    source_type: DataTableSource
    source: str
    name: str
    num_rows: int | None
    num_columns: int | None
    variable_name: VariableName | None
    columns: list[DataTableColumn]
    engine: VariableName | None = None
    type: DataTableType = "table"
    primary_keys: list[str] | None = None
    indexes: list[str] | None = None


class Schema(BaseStruct):
    """
    Represents a database schema and its tables.

    A schema may itself contain nested child schemas, e.g. for catalogs with
    hierarchical namespaces such as Iceberg (`top.nested.deep`).

    Attributes:
        name (str): The name of the schema.
        tables (List[DataTable]): Tables in this schema.
        tables_resolved (bool): True when `tables` has been enumerated
            False when table discovery was deferred. Defaults to True
        child_schemas (List[Schema]): Nested child schemas (sub-namespaces).
        child_schemas_resolved (bool): True when `child_schemas` has been
            enumerated. False when discovery was deferred. Defaults to True
    """

    name: str
    tables: list[DataTable]
    tables_resolved: bool = True
    child_schemas: list[Schema] = msgspec.field(default_factory=list)
    child_schemas_resolved: bool = True


class Database(BaseStruct):
    """
    Represents a collection of schemas.

    Attributes:
        name (str): The name of the database
        dialect (str): The dialect of the database
        schemas (List[Schema]): List of schemas in the database.
        schemas_resolved (bool): True when `schemas` has been enumerated.
            False when schema discovery was deferred. Defaults to True
        engine (Optional[VariableName]): Database engine or connection handler, if any.
    """

    name: str
    dialect: str
    schemas: list[Schema]
    schemas_resolved: bool = True
    engine: VariableName | None = None


if TYPE_CHECKING:
    from decimal import Decimal

    NumericLiteral = int | float | Decimal
    TemporalLiteral = date | time | datetime | timedelta
    NonNestedLiteral = NumericLiteral | TemporalLiteral | str | bool | bytes
else:
    # For runtime/msgspec, use Any since msgspec can't handle unions with
    # multiple str-like types (str, datetime, date, time, timedelta)
    NonNestedLiteral = Any


class ColumnStats(BaseStruct):
    """
    Represents stats for a column in a data table.
    """

    total: int | None = None
    nulls: int | None = None
    unique: int | None = None
    min: NonNestedLiteral | None = None
    max: NonNestedLiteral | None = None
    mean: NonNestedLiteral | None = None
    median: NonNestedLiteral | None = None
    std: NonNestedLiteral | None = None
    true: int | None = None
    false: int | None = None
    p5: NonNestedLiteral | None = None
    p25: NonNestedLiteral | None = None
    # p50 is the median
    p75: NonNestedLiteral | None = None
    p95: NonNestedLiteral | None = None


class BinValue(BaseStruct):
    """
    Represents bin values for a column in a data table. This is used for plotting.

    Attributes:
        bin_start (NonNestedLiteral): The start of the bin.
        bin_end (NonNestedLiteral): The end of the bin.
        count (int): The count of values in the bin.
    """

    bin_start: NonNestedLiteral
    bin_end: NonNestedLiteral
    count: int


class ValueCount(BaseStruct):
    """
    Represents a value and its count in a column in a data table.
    Currently used for string columns.

    Attributes:
        value (str): The value.
        count (int): The count of the value.
    """

    value: str
    count: int


class DataSourceConnection(BaseStruct):
    """
    Represents a data source connection.

    Attributes:
        source (str): The source of the data source connection. E.g 'postgres'.
        dialect (str): The dialect of the data source connection. E.g 'postgresql'.
        name (str): The name of the data source connection. E.g 'engine'.
        display_name (str): The display name of the data source connection. E.g 'PostgresQL (engine)'.
        databases (List[Database]): The databases in the data source connection.
        default_database (Optional[str]): The default database in the data source connection.
        default_schema (Optional[str]): The default schema in the data source connection.
    """

    source: str
    dialect: str
    name: str
    display_name: str
    databases: list[Database]
    default_database: str | None = None
    default_schema: str | None = None
