"""
TiDB Query Plan data structures and operator types.
Based on the Golang implementation from plan-change-capture.
"""

from abc import ABC, abstractmethod
from dataclasses import dataclass
from enum import Enum
from typing import List, Optional


class PlanVersion(str, Enum):
    V2 = "v2"
    V3 = "v3"
    V4 = "v4"
    UNKNOWN = "unknown"


class OpType(Enum):
    UNKNOWN = "unknown"
    HASH_JOIN = "hash_join"
    INDEX_JOIN = "index_join"
    MERGE_JOIN = "merge_join"
    HASH_AGG = "hash_agg"
    STREAM_AGG = "stream_agg"
    SELECTION = "selection"
    PROJECTION = "projection"
    TABLE_READER = "table_reader"
    TABLE_SCAN = "table_scan"
    INDEX_READER = "index_reader"
    INDEX_SCAN = "index_scan"
    INDEX_LOOKUP = "index_lookup"
    POINT_GET = "point_get"
    MAX_ONE_ROW = "max_one_row"
    APPLY = "apply"
    LIMIT = "limit"
    SORT = "sort"
    TOP_N = "top_n"
    TABLE_DUAL = "table_dual"
    TABLE_FULL_SCAN = "table_full_scan"
    TABLE_RANGE_SCAN = "table_range_scan"
    TABLE_ROW_ID_SCAN = "table_row_id_scan"
    SELECT_LOCK = "select_lock"
    SHOW = "show"
    MEM_TABLE_SCAN = "mem_table_scan"
    INDEX_RANGE_SCAN = "index_range_scan"
    INDEX_FULL_SCAN = "index_full_scan"
    BATCH_POINT_GET = "batch_point_get"
    UNION_SCAN = "union_scan"
    DELETE = "delete"
    UPDATE = "update"
    INSERT = "insert"  # NOTE: Not implemented in Golang reference


class JoinType(Enum):
    UNKNOWN = "unknown"
    INNER = "inner"
    LEFT_OUTER = "left_outer"
    RIGHT_OUTER = "right_outer"
    SEMI = "semi"
    ANTI_SEMI = "anti_semi"
    LEFT_OUTER_SEMI = "left_outer_semi"
    ANTI_LEFT_OUTER_SEMI = "anti_left_outer_semi"


class TaskType(Enum):
    ROOT = "root"
    TIKV = "tikv"
    TIFLASH = "tiflash"


def op_type_is_data_source(op_type: OpType) -> bool:
    """Check if operator type is a data source."""
    return op_type in {
        OpType.TABLE_READER,
        OpType.INDEX_READER,
        OpType.INDEX_LOOKUP,
        OpType.POINT_GET,
    }


def op_type_is_join(op_type: OpType) -> bool:
    """Check if operator type is a join."""
    return op_type in {OpType.INDEX_JOIN, OpType.MERGE_JOIN, OpType.HASH_JOIN}


class Operator(ABC):
    """Abstract base class for all operators."""

    @abstractmethod
    def get_id(self) -> str:
        pass

    @abstractmethod
    def get_type(self) -> OpType:
        pass

    @abstractmethod
    def get_est_row(self) -> float:
        pass

    @abstractmethod
    def get_task(self) -> TaskType:
        pass

    @abstractmethod
    def format(self, indent: int = 0) -> str:
        pass

    @abstractmethod
    def get_children(self) -> List["Operator"]:
        pass

    @abstractmethod
    def set_child(self, index: int, child: "Operator"):
        pass


@dataclass
class BaseOp(Operator):
    """Base operator implementation."""

    id: str
    op_type: OpType
    est_row: float
    task: TaskType
    children: List[Operator]

    def get_id(self) -> str:
        return self.id

    def get_type(self) -> OpType:
        return self.op_type

    def get_est_row(self) -> float:
        return self.est_row

    def get_task(self) -> TaskType:
        return self.task

    def format(self, indent: int = 0) -> str:
        result = " " * indent + self.id + "\n"
        for child in self.children:
            result += child.format(indent + 4)
        return result

    def get_children(self) -> List[Operator]:
        return self.children

    def set_child(self, index: int, child: Operator):
        self.children[index] = child


@dataclass
class HashJoinOp(BaseOp):
    """Hash join operator."""

    join_type: JoinType = JoinType.UNKNOWN

    def format(self, indent: int = 0) -> str:
        build_child = self.children[0].get_id() if self.children else ""
        result = " " * indent + f"{self.id} build: {build_child}\n"
        for child in self.children:
            result += child.format(indent + 4)
        return result


@dataclass
class IndexJoinOp(BaseOp):
    """Index join operator."""

    join_type: JoinType = JoinType.UNKNOWN


@dataclass
class MergeJoinOp(BaseOp):
    """Merge join operator."""

    join_type: JoinType = JoinType.UNKNOWN


@dataclass
class TableReaderOp(BaseOp):
    """Table reader operator."""


@dataclass
class TableScanOp(BaseOp):
    """Table scan operator."""

    table: str = ""

    def format(self, indent: int = 0) -> str:
        return " " * indent + f"{self.id}\ttable:{self.table}\n"


@dataclass
class IndexReaderOp(BaseOp):
    """Index reader operator."""


@dataclass
class IndexScanOp(BaseOp):
    """Index scan operator."""

    table: str = ""
    index: str = ""

    def format(self, indent: int = 0) -> str:
        prefix = " " * indent
        table_info = f"table:{self.table}, index:{self.index}"
        return f"{prefix}{self.id}\t{table_info}\n"


@dataclass
class IndexLookupOp(BaseOp):
    """Index lookup operator."""


@dataclass
class IndexRangeScanOp(BaseOp):
    """Index range scan operator."""

    table: str = ""
    index: str = ""


@dataclass
class IndexFullScanOp(BaseOp):
    """Index full scan operator."""

    table: str = ""
    index: str = ""


@dataclass
class SelectionOp(BaseOp):
    """Selection operator."""


@dataclass
class ProjectionOp(BaseOp):
    """Projection operator."""


@dataclass
class PointGetOp(BaseOp):
    """Point get operator."""

    batch: bool = False
    table: str = ""


@dataclass
class BatchPointGetOp(BaseOp):
    """Batch point get operator."""

    table: str = ""


@dataclass
class StreamAggOp(BaseOp):
    """Stream aggregation operator."""


@dataclass
class HashAggOp(BaseOp):
    """Hash aggregation operator."""


@dataclass
class MaxOneRowOp(BaseOp):
    """Max one row operator."""


@dataclass
class ApplyOp(BaseOp):
    """Apply operator."""


@dataclass
class LimitOp(BaseOp):
    """Limit operator."""


@dataclass
class SortOp(BaseOp):
    """Sort operator."""


@dataclass
class TopNOp(BaseOp):
    """TopN operator."""


@dataclass
class TableDualOp(BaseOp):
    """Table dual operator."""


@dataclass
class TableFullScanOp(BaseOp):
    """Table full scan operator."""

    table: str = ""


@dataclass
class TableRangeScanOp(BaseOp):
    """Table range scan operator."""

    table: str = ""


@dataclass
class TableRowIDScanOp(BaseOp):
    """Table RowID scan operator."""

    table: str = ""


@dataclass
class MemTableScanOp(BaseOp):
    """Mem table scan operator."""

    table: str = ""


@dataclass
class SelectLockOp(BaseOp):
    """Select lock operator."""


@dataclass
class ShowOp(BaseOp):
    """Show operator."""


@dataclass
class UnionScanOp(BaseOp):
    """Union scan operator."""


@dataclass
class DeleteOp(BaseOp):
    """Delete operator."""

    table: str = ""


@dataclass
class UpdateOp(BaseOp):
    """Update operator."""

    table: str = ""


@dataclass
class InsertOp(BaseOp):
    """Insert operator. NOTE: Not implemented in Golang reference."""

    table: str = ""


@dataclass
class Plan:
    """Query execution plan."""

    schema: str = ""
    sql: str = ""
    version: PlanVersion = PlanVersion.UNKNOWN
    root: Optional[Operator] = None
    plan_text: str = ""

    def format(self) -> str:
        """Format the plan as a string."""
        if self.root:
            return self.root.format(0)
        return ""