from datetime import datetime, timedelta
from typing import Callable
from icron import croniter
from bollhav.database import Database
from bollhav.intervals import TZInterval
from bollhav.postgres.columns import PostgresColumn
from bollhav.parquet.columns import ParquetColumn
from bollhav.modes import WriteMode, ModelType
from bollhav.batching import infer_batch_size
from bollhav.sorting import sort_columns


class Model:
    def __init__(
        self,
        name: str,
        source_entity: str,
        table: str = "",
        schema: str = "",
        database: Database | None = None,
        columns: list[PostgresColumn | ParquetColumn] | None = None,
        model_type: ModelType = ModelType.TABLE,
        write_mode: WriteMode = WriteMode.APPEND,
        tags: set[str] = set(),
        cron: str | None = None,
        enabled: bool = True,
        debug: bool = False,
        description: str | None = None,
        target_dsn_env_var: str | None = None,
        source_dsn_env_var: str | None = None,
        source_query: str | None = None,
        column_sorting: Callable | None = sort_columns,
        partitioned_by: str | None = None,
        partitioned_by_index: bool = False,
        begin: datetime | None = None,
        end: datetime | None = None,
        retries: int | None = None,
        lookback: int | None = None,
        name_add_to_tags: bool = True,
        schema_add_to_tags: bool = True,
        model_gets_all_tag: bool = True,
        tz_aware: bool = True,
        **kwargs,
    ):
        if model_type == ModelType.VIEW and write_mode != WriteMode.VIEW:
            raise ValueError("ModelType.VIEW must use WriteMode.VIEW")
        if model_type == ModelType.TABLE and write_mode == WriteMode.VIEW:
            raise ValueError("ModelType.TABLE cannot use WriteMode.VIEW")
        if database is not None and columns is None:
            raise ValueError("columns must be set when database is provided")
        if columns is not None and database is None:
            raise ValueError("database must be set when columns is provided")
        if partitioned_by and columns:
            col_names = {c.name for c in columns}
            if partitioned_by not in col_names:
                raise ValueError(
                    f"Partitioned_by references unknown column: {partitioned_by!r}"
                )
        if tz_aware:
            if begin is not None and (
                begin.tzinfo is None or begin.utcoffset() != timedelta(0)
            ):
                raise ValueError("begin must be UTC-aware")
            if end is not None and (
                end.tzinfo is None or end.utcoffset() != timedelta(0)
            ):
                raise ValueError("end must be UTC-aware")

        self.name = name
        self.source_entity = source_entity
        self.table = table
        self.schema = schema
        self.database = database
        self.columns = columns
        self.model_type = model_type
        self.write_mode = write_mode
        self.tags = set(tags) if tags else set()
        self.name_add_to_tags = name_add_to_tags
        if self.name_add_to_tags:
            tags.add(self.name)
        self.schema_add_to_tags = schema_add_to_tags
        if self.schema_add_to_tags:
            tags.add(self.schema)
        self.every_model_gets_all_tag = model_gets_all_tag
        if self.every_model_gets_all_tag:
            tags.add("all")
        self.cron = cron
        self.enabled = enabled
        self.debug = debug
        self.description = description
        self.target_dsn_env_var = target_dsn_env_var
        self.source_dsn_env_var = source_dsn_env_var
        self.source_query = source_query
        self.column_sorting = column_sorting
        self.partitioned_by = partitioned_by
        self.partitioned_by_index = partitioned_by is not None
        self.begin = begin
        self.end = end
        self.retries = retries
        self.lookback = lookback
        self.tz_aware = tz_aware
        self.batch_size = infer_batch_size(cron) if cron else None
        self.sensitive = (
            any(getattr(c, "sensitive", False) for c in columns) if columns else False
        )
        self.unique_columns = (
            [c for c in columns if getattr(c, "unique", False)] if columns else []
        )

        if write_mode == WriteMode.UPDATE_INSERT and not self.unique_columns:
            raise ValueError(
                "WriteMode.UPDATE_INSERT requires at least one column with unique=True"
            )

        if self.columns and self.column_sorting:
            col_names = [c.name for c in self.columns]
            sorted_names = self.column_sorting(col_names)
            name_to_col = {c.name: c for c in self.columns}
            self.columns = [name_to_col[n] for n in sorted_names]

        for key, val in kwargs.items():
            if callable(val):
                kwargs[key] = val(
                    **{k: v for k, v in kwargs.items() if not callable(v)}
                )
        self.extra = kwargs

    def __repr__(self) -> str:
        return (
            f"Model("
            f"name={self.name!r}, "
            f"source_entity={self.source_entity!r}, "
            f"table={self.table!r}, "
            f"schema={self.schema!r}, "
            f"database={self.database}, "
            f"columns={self.columns!r}, "
            f"model_type={self.model_type}, "
            f"write_mode={self.write_mode}, "
            f"tags={self.tags!r}, "
            f"cron={self.cron!r}, "
            f"enabled={self.enabled}, "
            f"debug={self.debug}, "
            f"description={self.description!r}, "
            f"source_dsn={self.source_dsn_env_var!r}, "
            f"source_query={self.source_query!r}, "
            f"partitioned_by={self.partitioned_by!r}, "
            f"begin={self.begin.isoformat() if self.begin else None!r}, "
            f"end={self.end.isoformat() if self.end else None!r}, "
            f"retries={self.retries!r}, "
            f"lookback={self.lookback!r}, "
            f"sensitive={self.sensitive}, "
            f"unique_columns={self.unique_columns!r}, "
            f"tz_aware={self.tz_aware}, "
            f"extra={self.extra!r})"
        )

    def __eq__(self, other: object) -> bool:
        if not isinstance(other, Model):
            return NotImplemented
        return self.__dict__ == other.__dict__

    def get_batch_intervals(
        self, interval: TZInterval, cron_override: str | None = None
    ) -> list[TZInterval]:
        it = croniter(cron_override or self.cron, interval.since)
        intervals: list[TZInterval] = []
        current = interval.since
        while True:
            tick = it.get_next(datetime)
            if tick >= interval.until:
                break
            intervals.append(TZInterval(current, tick))
            current = tick
        if current < interval.until:
            intervals.append(TZInterval(current, interval.until))
        return intervals
