from __future__ import annotations

import json
from typing import Callable

from pydantic import BaseModel

from palimpzest.constants import AggFunc, Cardinality
from palimpzest.core.data import context, dataset
from palimpzest.core.elements.filters import Filter
from palimpzest.core.elements.groupbysig import GroupBySig
from palimpzest.core.lib.schemas import Average, Count, Max, Min, Sum
from palimpzest.utils.hash_helpers import hash_for_id


class LogicalOperator:
    """
    A logical operator is an operator that operates on Sets.

    Right now it can be one of:
    - BaseScan (scans data from a root Dataset)
    - ContextScan (loads the context for a root Dataset)
    - FilteredScan (scans input Set and applies filter)
    - ConvertScan (scans input Set and converts it to new Schema)
    - LimitScan (scans up to N records from a Set)
    - GroupByAggregate (applies a group by on the Set)
    - Aggregate (applies an aggregation on the Set)
    - TopKScan (fetches documents from a provided input for a given query)
    - Map (applies a function to each record in the Set without adding any new columns)
    - ComputeOperator (executes a computation described in natural language)
    - SearchOperator (executes a search query on the input Context)

    Every logical operator must declare the get_logical_id_params() and get_logical_op_params() methods,
    which return dictionaries of parameters that are used to compute the logical op id and to implement
    the logical operator (respectively).
    """

    def __init__(
        self,
        output_schema: type[BaseModel],
        input_schema: type[BaseModel] | None = None,
        depends_on: list[str] | None = None,
    ):
        # TODO: can we eliminate input_schema?
        self.output_schema = output_schema
        self.input_schema = input_schema
        self.depends_on = [] if depends_on is None else sorted(depends_on)
        self.logical_op_id: str | None = None
        self.unique_logical_op_id: str | None = None

        # compute the fields generated by this logical operator
        input_field_names = list(self.input_schema.model_fields) if self.input_schema is not None else []
        self.generated_fields = sorted(
            [field_name for field_name in self.output_schema.model_fields if field_name not in input_field_names]
        )

    def __str__(self) -> str:
        raise NotImplementedError("Abstract method")

    def __eq__(self, other) -> bool:
        all_id_params_match = all(value == getattr(other, key) for key, value in self.get_logical_id_params().items())
        return isinstance(other, self.__class__) and all_id_params_match

    def copy(self) -> LogicalOperator:
        logical_op_copy = self.__class__(**self.get_logical_op_params())
        logical_op_copy.logical_op_id = self.logical_op_id
        logical_op_copy.unique_logical_op_id = self.unique_logical_op_id
        return logical_op_copy

    def logical_op_name(self) -> str:
        """Name of the logical operator."""
        return str(self.__class__.__name__)

    def get_unique_logical_op_id(self) -> str:
        """
        Get the unique logical operator id for this logical operator.
        """
        return self.unique_logical_op_id

    def set_unique_logical_op_id(self, unique_logical_op_id: str) -> None:
        """
        Set the unique logical operator id for this logical operator.
        This is used to uniquely identify the logical operator in the query plan.
        """
        self.unique_logical_op_id = unique_logical_op_id

    def get_logical_id_params(self) -> dict:
        """
        Returns a dictionary mapping of logical operator parameters which are relevant
        for computing the logical operator id.

        NOTE: Should be overriden by subclasses to include class-specific parameters.
        NOTE: input_schema and output_schema are not included in the id params because
              they depend on how the Optimizer orders operations.
        """
        # TODO: should we use `generated_fields` after getting rid of them in PhysicalOperator?
        return {"generated_fields": self.generated_fields}

    def get_logical_op_params(self) -> dict:
        """
        Returns a dictionary mapping of logical operator parameters which may be used to
        implement a physical operator associated with this logical operation.

        NOTE: Should be overriden by subclasses to include class-specific parameters.
        """
        return {
            "input_schema": self.input_schema,
            "output_schema": self.output_schema,
            "depends_on": self.depends_on,
        }

    def get_logical_op_id(self):
        """
        TODO: turn this into a property?

        NOTE: We do not call this in the __init__() method as subclasses may set parameters
              returned by self.get_logical_op_params() after they call to super().__init__().
        """
        # return self.logical_op_id if we've computed it before
        if self.logical_op_id is not None:
            return self.logical_op_id

        # get op name and op parameters which are relevant for computing the id
        logical_op_name = self.logical_op_name()
        logical_id_params = self.get_logical_id_params()
        logical_id_params = {k: str(v) for k, v in logical_id_params.items()}

        # compute, set, and return the op_id
        hash_str = json.dumps({"logical_op_name": logical_op_name, **logical_id_params}, sort_keys=True)
        self.logical_op_id = hash_for_id(hash_str)

        return self.logical_op_id

    def get_generated_fields(self) -> list[str]:
        """Returns the names of the fields generated by this logical operator."""
        return self.generated_fields

    def __hash__(self):
        if not self.logical_op_id:
            raise ValueError("logical_op_id not set, unable to hash")
        return int(self.logical_op_id, 16)


class Aggregate(LogicalOperator):
    """
    Aggregate is a logical operator that applies an aggregation to the input set and yields a single result.
    This is a base class that has to be further specialized to implement specific aggregation functions.
    """

    def __init__(
        self,
        agg_func: AggFunc | None = None,
        agg_str: str | None = None,
        *args,
        **kwargs,
    ):
        assert agg_func is not None or agg_str is not None, "Either agg_func or agg_str must be provided"
        if kwargs.get("output_schema") is None:
            if agg_func == AggFunc.COUNT:
                kwargs["output_schema"] = Count
            elif agg_func == AggFunc.AVERAGE:
                kwargs["output_schema"] = Average
            elif agg_func == AggFunc.SUM:
                kwargs["output_schema"] = Sum
            elif agg_func == AggFunc.MIN:
                kwargs["output_schema"] = Min
            elif agg_func == AggFunc.MAX:
                kwargs["output_schema"] = Max
            else:
                raise ValueError(f"Unsupported aggregation function: {agg_func}")

        super().__init__(*args, **kwargs)
        self.agg_func = agg_func
        self.agg_str = agg_str

    def __str__(self):
        desc = f"function: {str(self.agg_func.value)}" if self.agg_func else f"agg: {self.agg_str}"
        return f"{self.__class__.__name__}({desc})"

    def get_logical_id_params(self) -> dict:
        logical_id_params = super().get_logical_id_params()
        logical_id_params = {
            "agg_func": self.agg_func,
            "agg_str": self.agg_str,
            **logical_id_params,
        }

        return logical_id_params

    def get_logical_op_params(self) -> dict:
        logical_op_params = super().get_logical_op_params()
        logical_op_params = {
            "agg_func": self.agg_func,
            "agg_str": self.agg_str,
            **logical_op_params,
        }

        return logical_op_params


class BaseScan(LogicalOperator):
    """A BaseScan is a logical operator that represents a scan of a particular root Dataset."""

    def __init__(self, datasource: dataset.Dataset, output_schema: type[BaseModel], *args, **kwargs):
        super().__init__(*args, output_schema=output_schema, **kwargs)
        self.datasource = datasource

    def __str__(self):
        return f"BaseScan({self.datasource},{self.output_schema})"

    def __eq__(self, other) -> bool:
        return (
            isinstance(other, BaseScan)
            and self.input_schema == other.input_schema
            and self.output_schema == other.output_schema
            and self.datasource == other.datasource
        )

    def get_logical_id_params(self) -> dict:
        logical_id_params = super().get_logical_id_params()
        logical_id_params = {
            "id": self.datasource.id,
            **logical_id_params,
        }

        return logical_id_params

    def get_logical_op_params(self) -> dict:
        logical_op_params = super().get_logical_op_params()
        logical_op_params = {"datasource": self.datasource, **logical_op_params}

        return logical_op_params


class ContextScan(LogicalOperator):
    """A ContextScan is a logical operator that loads the context for a particular root Dataset."""

    def __init__(self, context: context.Context, output_schema: type[BaseModel], *args, **kwargs):
        super().__init__(*args, output_schema=output_schema, **kwargs)
        self.context = context

    def __str__(self):
        return f"ContextScan({self.context},{self.output_schema})"

    def __eq__(self, other) -> bool:
        return (
            isinstance(other, ContextScan)
            and self.context.id == other.context.id
        )

    def get_logical_id_params(self) -> dict:
        logical_id_params = super().get_logical_id_params()
        logical_id_params = {
            "id": self.context.id,
            **logical_id_params,
        }

        return logical_id_params

    def get_logical_op_params(self) -> dict:
        logical_op_params = super().get_logical_op_params()
        logical_op_params = {"context": self.context, **logical_op_params}

        return logical_op_params


class ConvertScan(LogicalOperator):
    """A ConvertScan is a logical operator that represents a scan of a particular input Dataset, with conversion applied."""

    def __init__(
        self,
        cardinality: Cardinality = Cardinality.ONE_TO_ONE,
        udf: Callable | None = None,
        desc: str | None = None,
        *args,
        **kwargs,
    ):
        super().__init__(*args, **kwargs)
        self.cardinality = cardinality
        self.udf = udf
        self.desc = desc

    def __str__(self):
        return f"ConvertScan({self.input_schema} -> {str(self.output_schema)})"

    def get_logical_id_params(self) -> dict:
        logical_id_params = super().get_logical_id_params()
        logical_id_params = {
            "cardinality": self.cardinality,
            "udf": self.udf,
            "desc": self.desc,
            **logical_id_params,
        }

        return logical_id_params

    def get_logical_op_params(self) -> dict:
        logical_op_params = super().get_logical_op_params()
        logical_op_params = {
            "cardinality": self.cardinality,
            "udf": self.udf,
            "desc": self.desc,
            **logical_op_params,
        }

        return logical_op_params


class Distinct(LogicalOperator):
    def __init__(self, distinct_cols: list[str] | None, *args, **kwargs):
        super().__init__(*args, **kwargs)
        # if distinct_cols is not None, check that all columns are in the input schema
        if distinct_cols is not None:
            for col in distinct_cols:
                assert col in self.input_schema.model_fields, f"Column {col} not found in input schema {self.input_schema} for Distinct operator"

        # store the list of distinct columns, sorted
        self.distinct_cols = (
            sorted([field_name for field_name in self.input_schema.model_fields])
            if distinct_cols is None
            else sorted(distinct_cols)
        )

    def __str__(self):
        return f"Distinct({self.distinct_cols})"

    def get_logical_id_params(self) -> dict:
        logical_id_params = super().get_logical_id_params()
        logical_id_params = {"distinct_cols": self.distinct_cols, **logical_id_params}

        return logical_id_params

    def get_logical_op_params(self) -> dict:
        logical_op_params = super().get_logical_op_params()
        logical_op_params = {
            "distinct_cols": self.distinct_cols,
            **logical_op_params,
        }

        return logical_op_params


class FilteredScan(LogicalOperator):
    """A FilteredScan is a logical operator that represents a scan of a particular input Dataset, with filters applied."""

    def __init__(
        self,
        filter: Filter,
        desc: str | None = None,
        *args,
        **kwargs,
    ):
        super().__init__(*args, **kwargs)
        self.filter = filter
        self.desc = desc

    def __str__(self):
        return f"FilteredScan({str(self.output_schema)}, {str(self.filter)})"

    def get_logical_id_params(self) -> dict:
        logical_id_params = super().get_logical_id_params()
        logical_id_params = {
            "filter": self.filter,
            "desc": self.desc,
            **logical_id_params,
        }

        return logical_id_params

    def get_logical_op_params(self) -> dict:
        logical_op_params = super().get_logical_op_params()
        logical_op_params = {
            "filter": self.filter,
            "desc": self.desc,
            **logical_op_params,
        }

        return logical_op_params


class GroupByAggregate(LogicalOperator):
    def __init__(
        self,
        group_by_sig: GroupBySig,
        *args,
        **kwargs,
    ):
        super().__init__(*args, **kwargs)
        if not self.input_schema:
            raise ValueError("GroupByAggregate requires an input schema")
        (valid, error) = group_by_sig.validate_schema(self.input_schema)
        if not valid:
            raise TypeError(error)
        self.group_by_sig = group_by_sig

    def __str__(self):
        return f"GroupBy({self.group_by_sig.serialize()})"

    def get_logical_id_params(self) -> dict:
        logical_id_params = super().get_logical_id_params()
        logical_id_params = {"group_by_sig": self.group_by_sig, **logical_id_params}

        return logical_id_params

    def get_logical_op_params(self) -> dict:
        logical_op_params = super().get_logical_op_params()
        logical_op_params = {
            "group_by_sig": self.group_by_sig,
            **logical_op_params,
        }

        return logical_op_params


class JoinOp(LogicalOperator):
    def __init__(self, condition: str, on: list[str] | None = None, how: str = "inner", desc: str | None = None, *args, **kwargs):
        super().__init__(*args, **kwargs)
        self.condition = condition
        self.on = on
        self.how = how
        self.desc = desc

    def __str__(self):
        return f"Join(condition={self.condition})" if self.on is None else f"Join(on={self.on}, how={self.how})"

    def get_logical_id_params(self) -> dict:
        logical_id_params = super().get_logical_id_params()
        logical_id_params = {
            "condition": self.condition,
            "on": self.on,
            "how": self.how,
            "desc": self.desc,
            **logical_id_params,
        }

        return logical_id_params

    def get_logical_op_params(self) -> dict:
        logical_op_params = super().get_logical_op_params()
        logical_op_params = {
            "condition": self.condition,
            "on": self.on,
            "how": self.how,
            "desc": self.desc,
            **logical_op_params,
        }

        return logical_op_params


class LimitScan(LogicalOperator):
    def __init__(self, limit: int, *args, **kwargs):
        super().__init__(*args, **kwargs)
        self.limit = limit

    def __str__(self):
        return f"LimitScan({str(self.input_schema)}, {str(self.output_schema)})"

    def get_logical_id_params(self) -> dict:
        logical_id_params = super().get_logical_id_params()
        logical_id_params = {"limit": self.limit, **logical_id_params}

        return logical_id_params

    def get_logical_op_params(self) -> dict:
        logical_op_params = super().get_logical_op_params()
        logical_op_params = {
            "limit": self.limit,
            **logical_op_params,
        }

        return logical_op_params


class Project(LogicalOperator):
    def __init__(self, project_cols: list[str], *args, **kwargs):
        super().__init__(*args, **kwargs)
        self.project_cols = project_cols

    def __str__(self):
        return f"Project({self.input_schema}, {self.project_cols})"

    def get_logical_id_params(self) -> dict:
        logical_id_params = super().get_logical_id_params()
        logical_id_params = {"project_cols": self.project_cols, **logical_id_params}

        return logical_id_params

    def get_logical_op_params(self) -> dict:
        logical_op_params = super().get_logical_op_params()
        logical_op_params = {
            "project_cols": self.project_cols,
            **logical_op_params,
        }

        return logical_op_params


class TopKScan(LogicalOperator):
    """A TopKScan is a logical operator that represents a scan of a particular input Dataset, with a top-k operation applied."""

    def __init__(
        self,
        index,
        search_func,
        search_attr,
        output_attrs,
        k,
        *args,
        **kwargs,
    ):
        super().__init__(*args, **kwargs)
        self.index = index
        self.search_func = search_func
        self.search_attr = search_attr
        self.output_attrs = output_attrs
        self.k = k

    def __str__(self):
        return f"TopKScan({self.input_schema} -> {str(self.output_schema)})"

    def get_logical_id_params(self) -> dict:
        # NOTE: if we allow optimization over index, then we will need to include it in the id params
        # NOTE: ^if we do this, we should probably make a wrapper around the index object to ensure that
        #       it can be serialized as a string properly
        logical_id_params = super().get_logical_id_params()
        logical_id_params = {
            "search_attr": self.search_attr,
            "output_attrs": self.output_attrs,
            "k": self.k,
            **logical_id_params,
        }

        return logical_id_params

    def get_logical_op_params(self) -> dict:
        logical_op_params = super().get_logical_op_params()
        logical_op_params = {
            "index": self.index,
            "search_func": self.search_func,
            "search_attr": self.search_attr,
            "output_attrs": self.output_attrs,
            "k": self.k,
            **logical_op_params,
        }

        return logical_op_params


class ComputeOperator(LogicalOperator):
    """
    A ComputeOperator is a logical operator that performs a computation described in natural language
    on a given Context.
    """

    def __init__(self, context_id: str, instruction: str, *args, **kwargs):
        super().__init__(*args, **kwargs)
        self.context_id = context_id
        self.instruction = instruction

    def __str__(self):
        return f"ComputeOperator(id={self.context_id}, instr={self.instruction:20s})"

    def get_logical_id_params(self) -> dict:
        logical_id_params = super().get_logical_id_params()
        logical_id_params = {
            "context_id": self.context_id,
            "instruction": self.instruction,
            **logical_id_params,
        }

        return logical_id_params

    def get_logical_op_params(self) -> dict:
        logical_op_params = super().get_logical_op_params()
        logical_op_params = {
            "context_id": self.context_id,
            "instruction": self.instruction,
            **logical_op_params,
        }

        return logical_op_params


class SearchOperator(LogicalOperator):
    """
    A SearchOperator is a logical operator that executes a search described in natural language
    on a given Context.
    """

    def __init__(self, context_id: str, search_query: str, *args, **kwargs):
        super().__init__(*args, **kwargs)
        self.context_id = context_id
        self.search_query = search_query

    def __str__(self):
        return f"SearchOperator(id={self.context_id}, search_query={self.search_query:20s})"

    def get_logical_id_params(self) -> dict:
        logical_id_params = super().get_logical_id_params()
        logical_id_params = {
            "context_id": self.context_id,
            "search_query": self.search_query,
            **logical_id_params,
        }

        return logical_id_params

    def get_logical_op_params(self) -> dict:
        logical_op_params = super().get_logical_op_params()
        logical_op_params = {
            "context_id": self.context_id,
            "search_query": self.search_query,
            **logical_op_params,
        }

        return logical_op_params
