import dataclasses
from collections import defaultdict
from collections.abc import Iterable
from functools import partial
from itertools import islice
from typing import TYPE_CHECKING

from nagra.exceptions import UnresolvedFK, ValidationError
from nagra.utils import logger
from nagra.transaction import ExecMany


if TYPE_CHECKING:
    from pandas import DataFrame
    from polars import LazyFrame


class WriterMixin:
    """
    Utility class that provide common methods for Update and Upsert
    """

    def __init__(self):
        self.groups, self.resolve_stm = self.prepare()

    def prepare(self):
        """
        Organise columns in groups and prepare statement to
        resolve fk based on columns expressions
        """
        groups = defaultdict(list)
        for col in self.columns:
            if "." in col:
                head, tail = col.split(".", 1)
                groups[head].append(tail)
            else:
                groups[col] = None

        resolve_stm = {}
        for col, to_select in groups.items():
            if not to_select:
                continue
            cond = ["(= %s {})" % c for c in to_select]
            ftable = self.table.schema.get(self.table.foreign_keys[col])
            select = ftable.select(ftable.primary_key).where(*cond)
            resolve_stm[col] = select.stm()
        return groups, resolve_stm

    def execute(self, *values):
        ids = self.executemany([values])
        if ids:
            return ids[0]

    def executemany(self, records: Iterable[tuple]) -> list:
        # Transform list of records into a dataframe-like dict
        value_df = dict(zip(self.columns, zip(*records)))
        if not value_df:
            return []
        arg_df = {}
        for col, to_select in self.groups.items():
            if to_select:
                values = list(zip(*(value_df[f"{col}.{s}"] for s in to_select)))
                # Try to instanciate lru cache
                cache_key = (self.resolve_stm[col], str(self.lenient))
                lru = self.trn.get_fk_cache(cache_key, fn=partial(self._resolve, col))
                if lru is not None:
                    arg_df[col] = lru.run(values)
                else:
                    arg_df[col] = self._resolve(col, values)
            else:
                arg_df[col] = value_df[col]

        # Build arg iterable
        args = self._exec_args(arg_df)
        # Work by chunks
        stm = self.stm()
        ids = []
        returning = self.table.primary_key is not None
        while True:
            chunk = list(islice(args, 1000))
            if not chunk:
                break
            match self.trn.flavor:
                case "sqlite" | "mssql":
                    for item in chunk:
                        cursor = self.trn.execute(stm, item)
                        if returning:
                            new_id = cursor.fetchone()
                            ids.append(new_id[0] if new_id else None)
                        cursor.close()

                case "postgresql":
                    cursor = self.trn.executemany(stm, chunk, returning)
                    if returning:
                        ids.extend(r and r[0] for r in cursor)

        # If conditions are present, enforce those
        if self._check:
            self.validate(ids)
        return ids

    def validate(self, ids: list[int]):
        iter_ids = iter(ids)
        pk = self.table.primary_key
        while True:
            chunk = list(islice(iter_ids, 1000))
            if not chunk:
                return
            cond = self._check + [f"(in {pk} %s)" % (" {}" * len(chunk))]
            select = self.table.select("(count *)").where(*cond)
            (count,) = select.execute(*chunk).fetchone()
            if count != len(chunk):
                msg = f"Validation failed! Condition is: {self._check} )"
                raise ValidationError(msg)

    def _resolve(self, col, values):
        # XXX Detect situation where more than on result is found for
        # a given value (we could also enforce that we only resolve
        # columns with unique constraints) ?
        stm = self.resolve_stm[col]
        exm = ExecMany(stm, values, trn=self.trn)
        for res, vals in zip(exm, values):
            if res is not None:
                yield res[0]
            elif any(v is None for v in vals):
                # One of the values is not given
                yield None
            elif self.lenient is True or col in self.lenient:
                msg = "Value '%s' not found for foreign key column '%s' of table %s"
                logger.info(msg, vals, col, self.table)
                yield None
            else:
                raise UnresolvedFK(
                    f"Unable to resolve '{vals}' (for foreign key "
                    f"{col} of table {self.table.name})"
                )

    def __call__(self, records):
        return self.executemany(records)

    def from_pandas(self, df: "DataFrame"):
        if df.empty:
            return self.executemany([])

        # Convert non-basic types to string
        is_copy = False
        for col in self.columns:
            if df[col].dtype in ("int", "float", "bool", "str"):
                continue
            if df[col].dtype == "object":
                # bytes is not a dedicated type, we rely on the first
                # value and hope the column type is consistent
                if isinstance(df[col][0], bytes):
                    continue
            if not is_copy:
                df = df.copy()
                is_copy = True
            df[col] = df[col].astype(str)

        rows = df[self.columns].values
        return self.executemany(rows)

    def from_polars(self, df: "LazyFrame", batch: bool = False):
        """
        Write data from a polars LazyFrame. Set `batch` to True
        to enable streaming through the collect_batches() method
        """
        from polars import Struct, col

        # Ignore extra columns
        df.select(self.columns)
        schema = df.collect_schema()
        # Convert non-basic types to string
        for name, dtype in schema.items():
            if dtype == Struct:
                df = df.with_columns(col(name).struct.json_encode())

        if batch:
            chunks = df.collect_batches()
        else:
            chunks = [df.collect()]

        res = []
        for chunk in chunks:
            rows = chunk.iter_rows()
            res += self.executemany(rows)
        return res

    def from_dict(self, records):
        # Create select object in order to generate the same column names
        select = self.table.select(*self.columns)
        field_names = [f.name for f in dataclasses.fields(select.to_dataclass())]

        # Extract dict values - allows for field names or dotted column format
        f_or_c = list(zip(field_names, self.columns))
        rows = (
            tuple(getter(record, field, col) for col, field in f_or_c)
            for record in records
        )
        return self.executemany(rows)


def getter(record, field, col):
    """
    Utility function to extract value as a column name (with
    potentially dots) or a field name from record
    """
    if field in record:
        return record[field]
    if col in record:
        return record[col]
    raise KeyError(f"KeyError: neither {field} or {col} found")


def _slicer(chunk_size=10_000):
    start = 0
    while True:
        stop = start + chunk_size
        yield start, stop
        start = stop
