from __future__ import annotations

from dataclasses import dataclass
from typing import Iterable

from ...metamodel.metamodel import Aggregate, Construct, Exists, Logical, Lookup, Match, Not, Task, Union, Value, Var


@dataclass(frozen=True, slots=True)
class TaskSemantics:
    task_var_ids: frozenset[int]
    produced_var_ids: frozenset[int]
    bindable_var_ids: frozenset[int]


def task_var_ids(task: Task) -> set[int]:
    return set(analyze_task_semantics(task).task_var_ids)


def produced_var_ids(task: Task) -> set[int]:
    return set(analyze_task_semantics(task).produced_var_ids)


def bindable_var_ids(task: Task) -> set[int]:
    return set(analyze_task_semantics(task).bindable_var_ids)


def analyze_task_semantics(task: Task) -> TaskSemantics:
    return _analyze(task=task, memo={})


def _analyze(*, task: Task, memo: dict[int, TaskSemantics]) -> TaskSemantics:
    key = id(task)
    existing = memo.get(key)
    if existing is not None:
        return existing

    semantics = _compute(task=task, memo=memo)
    memo[key] = semantics
    return semantics


def _compute(*, task: Task, memo: dict[int, TaskSemantics]) -> TaskSemantics:
    if isinstance(task, Lookup):
        vars_all = _value_var_ids(task.args)
        if not task.relation.fields:
            produced = set(vars_all)
        else:
            produced: set[int] = set()
            for field, arg in zip(task.relation.fields, task.args):
                if field.input:
                    continue
                produced.update(_value_var_ids((arg,)))
        return TaskSemantics(
            task_var_ids=frozenset(vars_all),
            produced_var_ids=frozenset(produced),
            bindable_var_ids=frozenset(vars_all),
        )

    if isinstance(task, Construct):
        vars_all = {task.id_var.id}
        vars_all.update(_value_var_ids(task.values))
        return TaskSemantics(
            task_var_ids=frozenset(vars_all),
            produced_var_ids=frozenset((task.id_var.id,)),
            bindable_var_ids=frozenset((task.id_var.id,)),
        )

    if isinstance(task, Aggregate):
        body = _analyze(task=task.body, memo=memo)
        vars_all = set(body.task_var_ids)
        vars_all.update(value.id for value in task.group if isinstance(value, Var))
        vars_all.update(value.id for value in task.projection if isinstance(value, Var))
        vars_all.update(value.id for value in task.args if isinstance(value, Var))
        out = set()
        if task.args and isinstance(task.args[-1], Var):
            out.add(task.args[-1].id)
        return TaskSemantics(
            task_var_ids=frozenset(vars_all),
            produced_var_ids=frozenset(out),
            bindable_var_ids=frozenset(out),
        )

    if isinstance(task, Union):
        vars_all = {value.id for value in task.outputs if isinstance(value, Var)}
        for child in task.tasks:
            vars_all.update(_analyze(task=child, memo=memo).task_var_ids)
        out = {value.id for value in task.outputs if isinstance(value, Var)}
        return TaskSemantics(
            task_var_ids=frozenset(vars_all),
            produced_var_ids=frozenset(out),
            bindable_var_ids=frozenset(out),
        )

    if isinstance(task, Match):
        vars_all = {value.id for value in task.outputs if isinstance(value, Var)}
        for child in task.tasks:
            vars_all.update(_analyze(task=child, memo=memo).task_var_ids)
        out = {value.id for value in task.outputs if isinstance(value, Var)}
        return TaskSemantics(
            task_var_ids=frozenset(vars_all),
            produced_var_ids=frozenset(out),
            bindable_var_ids=frozenset(out),
        )

    if isinstance(task, Logical):
        vars_all: set[int] = set()
        produced: set[int] = set()
        bindable: set[int] = set()
        for child in task.body:
            child_sem = _analyze(task=child, memo=memo)
            vars_all.update(child_sem.task_var_ids)
            produced.update(child_sem.produced_var_ids)
            bindable.update(child_sem.bindable_var_ids)
        return TaskSemantics(
            task_var_ids=frozenset(vars_all),
            produced_var_ids=frozenset(produced),
            bindable_var_ids=frozenset(bindable),
        )

    if isinstance(task, Not):
        child = _analyze(task=task.task, memo=memo)
        return TaskSemantics(
            task_var_ids=child.task_var_ids,
            produced_var_ids=frozenset(),
            bindable_var_ids=frozenset(),
        )

    if isinstance(task, Exists):
        child = _analyze(task=task.task, memo=memo)
        vars_all = set(child.task_var_ids)
        vars_all.update(value.id for value in task.vars if isinstance(value, Var))
        return TaskSemantics(
            task_var_ids=frozenset(vars_all),
            produced_var_ids=frozenset(),
            bindable_var_ids=frozenset(),
        )

    return TaskSemantics(
        task_var_ids=frozenset(),
        produced_var_ids=frozenset(),
        bindable_var_ids=frozenset(),
    )


def _value_var_ids(values: Iterable[Value]) -> set[int]:
    out: set[int] = set()
    for value in values:
        if isinstance(value, Var):
            out.add(value.id)
        elif isinstance(value, (tuple, list)):
            out.update(_value_var_ids(value))
    return out
