from enum import StrEnum
from typing import TYPE_CHECKING, Annotated, Any, Literal

from pydantic import Field, RootModel

from kb_dashboard_core.filters.view import KbnFilter
from kb_dashboard_core.panels.charts.esql.columns.view import KbnESQLColumnTypes
from kb_dashboard_core.panels.charts.lens.columns.view import KbnLensColumnTypes
from kb_dashboard_core.panels.view import KbnBasePanel, KbnBasePanelEmbeddableConfig
from kb_dashboard_core.queries.view import KbnESQLQuery, KbnQuery
from kb_dashboard_core.shared.view import BaseVwModel, KbnReference, OmitIfNone

if TYPE_CHECKING:
    from .datatable.view import KbnDatatableVisualizationState
    from .gauge.view import KbnGaugeVisualizationState
    from .heatmap.view import KbnHeatmapVisualizationState
    from .metric.view import KbnESQLMetricVisualizationState, KbnMetricVisualizationState
    from .mosaic.view import KbnMosaicVisualizationState
    from .pie.view import KbnPieVisualizationState
    from .tagcloud.view import KbnTagcloudVisualizationState
    from .xy.view import KbnXYVisualizationState

    KbnVisualizationStateTypes = (
        KbnPieVisualizationState
        | KbnMetricVisualizationState
        | KbnESQLMetricVisualizationState
        | KbnGaugeVisualizationState
        | KbnHeatmapVisualizationState
        | KbnXYVisualizationState
        | KbnDatatableVisualizationState
        | KbnTagcloudVisualizationState
        | KbnMosaicVisualizationState
    )

# region Form Data Source
## Form Based
# "formBased": {                                                <--- KbnFormBasedDataSourceState
#     "layers": {                                               <--- KbnFormBasedDataSourceStateById
#         "b38fc34e-2df5-41a4-913b-e3e182998ef4": {             <--- KbnFormBasedDataSourceStateLayer
#             "columns": {
#                 "c4ea489f-9a04-4120-b9ec-aaa7a3a2d865": {     <--- dict[str, KbnLensColumnTypes]
#                     ...
#                 },
#             },
#             "columnOrder": [
#                 "44777308-3ea7-4f72-86bf-b4defdf71a0c",
#             ],
#             "sampling": 1,
#             "ignoreGlobalFilters": false,
#             "incompleteColumns": {},
#             "indexPatternId": "27a3148b-d1d4-4455-8acf-e63c94071a5b"
#         }
#     },
#     "currentIndexPatternId": "27a3148b-d1d4-4455-8acf-e63c94071a5b"
# },


class KbnFormBasedDataSourceStateLayer(BaseVwModel):
    """Represents the datasource state for a single layer within a Lens panel in the Kibana JSON structure."""

    columns: dict[str, KbnLensColumnTypes] = Field(default_factory=dict)
    columnOrder: list[str] = Field(default_factory=list)
    incompleteColumns: dict[str, Any] = Field(default_factory=dict)
    sampling: int
    indexPatternId: Annotated[str | None, OmitIfNone()] = None


class KbnFormBasedDataSourceStateLayerById(RootModel[dict[str, KbnFormBasedDataSourceStateLayer]]):
    """Represents a mapping of layer IDs to their corresponding KbnLayerDataSourceState objects."""

    root: dict[str, KbnFormBasedDataSourceStateLayer] = Field(default_factory=dict)


class KbnFormBasedDataSourceState(BaseVwModel):
    """Represents the datasource state for a Lens panel in the Kibana JSON structure."""

    layers: KbnFormBasedDataSourceStateLayerById = Field(default_factory=KbnFormBasedDataSourceStateLayerById)
    currentIndexPatternId: Annotated[str | None, OmitIfNone()] = None


# endregion Form Data Source

# region Text Data Source
# "textBased": {                                                          <--- KbnTextBasedDataSourceState
#     "layers": {                                                         <--- KbnTextBasedDataSourceStateLayerById
#         "42607131-7dd7-4935-a1eb-c9bed5cd302c": {                       <--- KbnTextBasedDataSourceStateLayer
#             "query": {                                                  <--- KbnESQLQuery
#                 "esql": "FROM metrics-* | ..."
#             },
#             "columns": [                                                <--- list[KbnESQLColumnTypes]
#                 {
#                     "columnId": "0e4107a7-db0d-4a80-a2e6-f528bc71e9f3",
#                     "fieldName": "agent.name"
#                 },
#                 {
#                     "columnId": "ae78ca1b-bad2-47b3-aaaf-ecd8aa05761c",
#                     "fieldName": "count(*)"
#                 }
#             ],
#             "allColumns": [                                             <--- list[KbnESQLColumnTypes]
#                 ...
#             ]
#         }
#     },
#     "indexPatternRefs": [                                            <--- list[KbnIndexPatternRef]
#         {
#             "id": ".....",
#             "title": "metrics-*",
#             "timeField": "@timestamp"
#         }
#     ]
# }


class KbnTextBasedDataSourceStateLayer(BaseVwModel):
    query: KbnESQLQuery
    columns: list[KbnESQLColumnTypes]
    allColumns: list[KbnESQLColumnTypes]
    timeField: str


class KbnTextBasedDataSourceStateLayerById(RootModel[dict[str, KbnTextBasedDataSourceStateLayer]]):
    root: dict[str, KbnTextBasedDataSourceStateLayer] = Field(default_factory=dict)


class KbnTextBasedDataSourceState(BaseVwModel):
    layers: KbnTextBasedDataSourceStateLayerById | None = Field(default=None)


# endregion Text Data Source

# region Index Pattern


class KbnIndexPatternBasedDataSourceStateById(RootModel[dict[str, str]]):
    root: dict[str, str] = Field(default_factory=dict)


class KbnIndexPatternBasedDataSourceState(BaseVwModel):
    """Index Pattern based datasource is not yet implemented."""

    layers: KbnIndexPatternBasedDataSourceStateById = Field(default_factory=KbnIndexPatternBasedDataSourceStateById)


# endregion Index Pattern


# region DataSourceState
class KbnDataSourceState(BaseVwModel):
    """Represents the overall datasource states for a Lens panel in the Kibana JSON structure."""

    formBased: Annotated[KbnFormBasedDataSourceState | None, OmitIfNone()] = None
    """Form-based datasource state. Only included for Lens panels using data views."""

    indexpattern: Annotated[KbnIndexPatternBasedDataSourceState | None, OmitIfNone()] = None
    """Index pattern datasource state (not implemented, rarely used)."""

    textBased: Annotated[KbnTextBasedDataSourceState | None, OmitIfNone()] = None
    """Text-based (ES|QL) datasource state. Only included for ES|QL panels."""


# endregion DataSourceState

# region Visualization

# Note: Base visualization classes (KbnBaseStateVisualization, KbnBaseStateVisualizationLayer,
# KbnLayerColorMapping, etc.) are now in base.py to avoid circular imports


class KbnVisualizationTypeEnum(StrEnum):
    XY = 'lnsXY'
    PIE = 'lnsPie'
    METRIC = 'lnsMetric'
    GAUGE = 'lnsGauge'
    HEATMAP = 'lnsHeatmap'
    DATATABLE = 'lnsDatatable'
    TAGCLOUD = 'lnsTagcloud'


class KbnLensPanelState(BaseVwModel):
    """Represents the 'state' object within a Lens panel in the Kibana JSON structure."""

    visualization: Any  # KbnPieVisualizationState | KbnMetricVisualizationState | KbnXYVisualizationState | KbnTagcloudVisualizationState
    query: KbnQuery | KbnESQLQuery = Field(...)
    filters: list[KbnFilter] = Field(...)
    datasourceStates: KbnDataSourceState = Field(...)
    internalReferences: list[Any] = Field(...)
    adHocDataViews: dict[str, Any] = Field(...)


class KbnLensPanelAttributes(BaseVwModel):
    title: str = ''
    visualizationType: KbnVisualizationTypeEnum
    type: Literal['lens'] = 'lens'
    references: list[KbnReference] = Field(...)
    state: KbnLensPanelState


class KbnLensPanelEmbeddableConfig(KbnBasePanelEmbeddableConfig):
    attributes: KbnLensPanelAttributes

    syncTooltips: bool = Field(
        default=False,
        description='(Optional) Whether to sync tooltips across visualizations. Defaults to False.',
    )

    syncColors: bool = Field(
        default=False,
        description='(Optional) Whether to sync colors across visualizations. Defaults to False.',
    )

    syncCursor: bool = Field(
        default=True,
        description='(Optional) Whether to sync cursor across visualizations. Defaults to True.',
    )

    filters: list[Any] = Field(
        default_factory=list,
        description='(Optional) List of filters applied to the Lens visualization. Defaults to empty list.',
    )

    query: KbnQuery = Field(
        ...,
        description="(Optional) Query object for the Lens visualization. Defaults to empty query with 'kuery' language.",
    )


class KbnLensPanel(KbnBasePanel):
    type: Literal['lens'] = 'lens'
    embeddableConfig: KbnLensPanelEmbeddableConfig
