from __future__ import annotations

from collections import defaultdict
from dataclasses import dataclass, field
from enum import Enum

from relationalai.util.structures import OrderedSet

from ...metamodel.builtins import is_entity_type
from ...metamodel.metamodel import Data, Field, Model as MMModel, Relation, ScalarType, Table, Task


class TableKind(Enum):
    CONCEPT = 1
    DATA = 2
    RELATIONSHIP = 3
    EXTERNAL = 4
    OUTPUT = 5


@dataclass
class TableInfo:
    relation: Relation = field()
    columns: list["ColumnInfo"] = field(default_factory=list)
    kind: TableKind = field(default=TableKind.CONCEPT)

    @property
    def name(self) -> str:
        return self.relation.name

    def __hash__(self) -> int:
        return hash(self.relation)


@dataclass
class ColumnInfo:
    table: TableInfo = field()
    column: Relation | Field = field()

    @property
    def is_id(self) -> bool:
        return self.table.relation == self.column

    @property
    def name(self) -> str:
        if self.is_id:
            return "_id"
        return self.column.name

    def __hash__(self) -> int:
        return hash((self.table.relation.id, self.column.id))


@dataclass
class ModelAnalysis:
    model: MMModel = field()
    analysis_root: Task | None = None
    writes: dict[ColumnInfo, OrderedSet[object]] = field(
        default_factory=lambda: defaultdict(OrderedSet[object])
    )
    table_info: dict[Relation, TableInfo] = field(default_factory=dict)
    column_info: dict[Relation | Field, ColumnInfo] = field(default_factory=dict)
    tables_with_writes: set[TableInfo] = field(default_factory=set)
    unwritten_logical_tables: set[TableInfo] = field(default_factory=set)

    def get_column_info(self, relation: Relation | Field) -> ColumnInfo | None:
        found = self.column_info.get(relation, None)
        if not found:
            self._capture_relation(relation if isinstance(relation, Relation) else relation._relation)
            found = self.column_info.get(relation, None)
        return found

    def get_table_info(self, relation: Relation) -> TableInfo | None:
        found = self.table_info.get(relation, None)
        if not found and relation.fields:
            self._capture_relation(relation)
            found = self.table_info.get(relation, None)
        return found

    def is_property(self, relation: Relation) -> bool:
        if self.model._analysis.is_property(relation):
            return True
        if relation.fields and isinstance(relation.fields[0].type, Table) and relation in relation.fields[0].type.columns:
            return True
        return False

    def _capture_relation(self, relation: Relation) -> None:
        info = self.table_info
        col_info = self.column_info

        if _is_function(relation) or relation in col_info or relation in info:
            return

        if isinstance(relation, Data):
            info[relation] = TableInfo(relation, kind=TableKind.DATA)
            info[relation].columns.append(ColumnInfo(info[relation], relation))
            col_info[relation] = info[relation].columns[-1]
        elif isinstance(relation, Table):
            kind = TableKind.EXTERNAL if not relation.uri.startswith("dataframe://") else TableKind.OUTPUT
            info[relation] = TableInfo(relation, kind=kind)
            info[relation].columns.append(ColumnInfo(info[relation], relation))
            col_info[relation] = info[relation].columns[-1]
            for table_col in relation.columns:
                self._capture_relation(table_col)
        elif isinstance(relation, ScalarType):
            if not is_entity_type(relation):
                return
            info[relation] = TableInfo(relation)
            info[relation].columns.append(ColumnInfo(info[relation], relation))
            col_info[relation] = info[relation].columns[-1]
        elif (
            self.is_property(relation)
            and (len(relation.fields) == 2 or isinstance(relation.fields[0].type, Table))
            and is_entity_type(relation.fields[0].type)
        ):
            root = relation.fields[0].type
            info[root].columns.append(ColumnInfo(info[root], relation))
            col_info[relation] = info[root].columns[-1]
        else:
            info[relation] = TableInfo(relation, kind=TableKind.RELATIONSHIP)
            for field in relation.fields:
                column = ColumnInfo(info[relation], field)
                info[relation].columns.append(column)
                col_info[field] = column

    def for_query(self) -> ModelAnalysis:
        return ModelAnalysis(
            model=self.model,
            analysis_root=self.analysis_root,
            table_info=self.table_info.copy(),
            column_info=self.column_info.copy(),
        )


def determine_tables(analysis: ModelAnalysis) -> None:
    for relation in analysis.model.types:
        analysis._capture_relation(relation)
    for relation in analysis.model.relations:
        analysis._capture_relation(relation)


def _is_function(relation: Relation) -> bool:
    return any(field.input for field in relation.fields)


__all__ = [
    "ColumnInfo",
    "ModelAnalysis",
    "TableInfo",
    "TableKind",
    "determine_tables",
]
