from __future__ import annotations

from dataclasses import dataclass

from ...metamodel import builtins as bt
from ...metamodel.metamodel import (
    Aggregate,
    Exists,
    Logical,
    Lookup,
    Match,
    Model as MMModel,
    Not,
    ScalarType,
    Task,
    Union,
    UnionType,
    Update,
    Var,
)
from ...metamodel.metamodel_analyzer import Analyzer as MMAnalyzer, Normalize
from ...metamodel.typer import Typer
from .coalesce import SQL2CoalescedArtifacts, SQL2Coalescer
from .graph import SQL2GraphArtifacts, SQL2SingleUpdateGraphBuilder
from .model_analysis import ModelAnalysis, determine_tables
from .normalize import SQL2NormalizedArtifacts, SQL2Normalizer
from .task_semantics import bindable_var_ids, task_var_ids
from .table_shape import SQL2TableShapeSidecar, build_sql2_table_shape_sidecar


@dataclass(frozen=True)
class SQL2FragmentProjectionSpec:
    relation_id: int
    relation_name: str
    coalesce_group_id: int
    relation_key_field_name: str
    relation_key_type_id: int | None
    value_field_names: tuple[str, ...]
    key_var_columns: tuple[tuple[int, str, bool], ...]
    optional: bool


@dataclass(frozen=True)
class SQL2FragmentProjectionOutputColumn:
    relation_id: int
    field_name: str
    output_alias: str


@dataclass(frozen=True)
class SQL2FragmentProjectionJoinGroup:
    group_id: int
    relation_ids: tuple[int, ...]
    join_var_ids: tuple[int, ...]
    output_anchor_join_type: str


@dataclass(frozen=True)
class SQL2FragmentProjectionSourceGroup:
    group_id: int
    join_group_id: int
    relation_ids: tuple[int, ...]
    output_anchor_join_type: str


@dataclass(frozen=True)
class SQL2FragmentProjectionAssemblyPlan:
    anchor_relation_id: int | None
    join_groups: tuple[SQL2FragmentProjectionJoinGroup, ...]
    source_groups: tuple[SQL2FragmentProjectionSourceGroup, ...]
    output_columns: tuple[SQL2FragmentProjectionOutputColumn, ...]


@dataclass(frozen=True)
class SQL2FragmentPlan:
    query_mode: str
    projection_specs: tuple[SQL2FragmentProjectionSpec, ...]
    projection_assembly: SQL2FragmentProjectionAssemblyPlan
    projection_relation_task_var_ids_by_relation_id: tuple[tuple[int, tuple[int, ...]], ...]
    projection_anchor_relation_ids: tuple[int, ...]
    projection_anchor_correlation_var_ids: tuple[int, ...]
    output_anchor_projection_relation_ids: tuple[int, ...]
    output_anchor_projection_correlation_var_ids: tuple[int, ...]
    fallback_relation_id: int | None
    requires_row_identity_dedupe: bool
    output_relation_id: int | None
    output_relation_key_field_name: str | None
    output_relation_key_type_id: int | None
    output_relation_key_var_id: int | None


@dataclass(frozen=True)
class SQL2Analysis:
    model_analysis: ModelAnalysis
    table_shape: SQL2TableShapeSidecar
    graph: SQL2GraphArtifacts
    coalesced: SQL2CoalescedArtifacts
    normalized: SQL2NormalizedArtifacts
    model_key_preserving_table_names: tuple[str, ...] = ()
    fragment_plan: SQL2FragmentPlan | None = None


class SQL2ModelAnalyzer:
    """Model analyzer for the clean-room SQL2 pipeline."""

    def __init__(self):
        self.typer = Typer(enforce=True)
        self.normalizer = Normalize()
        self.mm_analyzer = MMAnalyzer()
        self.graph_builder = SQL2SingleUpdateGraphBuilder()
        self.coalescer = SQL2Coalescer()
        self.model_normalizer = SQL2Normalizer()
        self.model_analysis: ModelAnalysis | None = None
        self.sql2_analysis: SQL2Analysis | None = None

    def normalize_model(self, model: MMModel) -> MMModel:
        model = self.typer.infer_model(model)
        normalized_root = self.normalizer.normalize(model.root)
        model = model.mut(root=normalized_root)
        self.mm_analyzer.analyze(model)
        return model

    def normalize_fragment(self, fragment: Task) -> Task:
        fragment = self.typer.infer_query(fragment)
        return self.normalizer.normalize(fragment)

    def analyze(self, model: MMModel) -> SQL2Analysis:
        model = self.normalize_model(model)
        analysis = ModelAnalysis(model)
        analysis.analysis_root = model.root
        determine_tables(analysis)
        table_shape = build_sql2_table_shape_sidecar(analysis)
        graph = self.graph_builder.build(analysis)
        coalesced = self.coalescer.coalesce(graph=graph, table_shape=table_shape)
        normalized = self.model_normalizer.normalize(
            coalesced,
            model_analysis=analysis,
            table_shape=table_shape,
        )
        model_key_preserving_table_names = self._model_key_preserving_table_names(
            normalized=normalized,
            table_shape=table_shape,
        )

        sql2_analysis = SQL2Analysis(
            model_analysis=analysis,
            table_shape=table_shape,
            graph=graph,
            coalesced=coalesced,
            normalized=normalized,
            model_key_preserving_table_names=model_key_preserving_table_names,
            fragment_plan=None,
        )
        self.model_analysis = analysis
        self.sql2_analysis = sql2_analysis
        return sql2_analysis

    def analyze_fragment(self, fragment: Task) -> SQL2Analysis:
        if self.model_analysis is None or self.sql2_analysis is None:
            raise ValueError("SQL2 fragment analysis requires a previously analyzed model.")
        requested_projection_relation_names = [
            update.relation.name
            for update in self._collect_updates(fragment)
            if not self._is_output_marker_update(update)
        ]
        normalized_fragment = self.normalize_fragment(fragment)
        normalized_fragment = self._prune_redundant_fragment_populations(normalized_fragment)
        query_analysis = self.model_analysis.for_query()
        query_analysis.analysis_root = normalized_fragment

        # Reuse model table-shape/source mappings; query-local synthetic relations
        # are handled via graph/coalescing fallback metadata.
        table_shape = self.sql2_analysis.table_shape
        graph = self.graph_builder.build(query_analysis)
        coalesced = self.coalescer.coalesce(graph=graph, table_shape=table_shape)
        normalized = self.model_normalizer.normalize(
            coalesced,
            model_analysis=query_analysis,
            table_shape=table_shape,
            model_key_preserving_table_names=self.sql2_analysis.model_key_preserving_table_names,
        )
        updates_with_optional = self._collect_updates_with_optional_context(normalized_fragment)
        normalized_updates = [update for update, _optional in updates_with_optional]
        optional_relation_ids = {
            update.relation.id
            for update, optional in updates_with_optional
            if optional
        }
        projection_updates = self._resolve_projection_updates(
            requested_projection_relation_names=requested_projection_relation_names,
            normalized_updates=normalized_updates,
        )
        output_update = next(
            (update for update in normalized_updates if self._is_output_marker_update(update)),
            None,
        )
        relation_ids_with_channels = self._relation_ids_with_channels(normalized)
        coalesce_group_id_by_relation_id = self._projection_coalesce_group_id_by_relation_id(
            projection_updates=projection_updates,
            normalized=normalized,
        )
        projection_specs = [
            SQL2FragmentProjectionSpec(
                relation_id=update.relation.id,
                relation_name=update.relation.name,
                coalesce_group_id=coalesce_group_id_by_relation_id.get(update.relation.id, 0),
                relation_key_field_name=update.relation.fields[0].name if update.relation.fields else "_id",
                relation_key_type_id=(
                    getattr(update.relation.fields[0].type, "id", None)
                    if update.relation.fields
                    else None
                ),
                value_field_names=self._value_field_names(update),
                key_var_columns=self._projection_key_var_columns(update),
                optional=update.relation.id in optional_relation_ids,
            )
            for update in projection_updates
        ]
        projection_relation_task_var_ids_by_relation_id = self._projection_relation_task_var_ids_by_relation_id(
            normalized=normalized,
            relation_ids={spec.relation_id for spec in projection_specs},
        )
        requires_row_identity_dedupe = self._fragment_requires_row_identity_dedupe(
            projection_relation_ids={update.relation.id for update in projection_updates},
            normalized=normalized,
        )
        projection_assembly = self._build_projection_assembly_plan(
            projection_specs,
        )
        output_relation_id = output_update.relation.id if output_update is not None else None
        projection_anchor_relation_ids = self._projection_relation_ids_for_anchor(
            projection_specs=projection_specs,
            relation_ids_with_channels=relation_ids_with_channels,
            required_key_type_id=None,
        )
        projection_anchor_correlation_var_ids = self._projection_correlation_var_ids(
            normalized=normalized,
            projection_specs=tuple(
                spec for spec in projection_specs if spec.relation_id in set(projection_anchor_relation_ids)
            ),
        )
        output_anchor_projection_relation_ids = self._projection_relation_ids_for_anchor(
            projection_specs=projection_specs,
            relation_ids_with_channels=relation_ids_with_channels,
            required_key_type_id=(
                getattr(output_update.relation.fields[0].type, "id", None)
                if output_update is not None and output_update.relation.fields
                else None
            ),
        )
        output_anchor_projection_correlation_var_ids = self._projection_correlation_var_ids(
            normalized=normalized,
            projection_specs=tuple(
                spec for spec in projection_specs if spec.relation_id in set(output_anchor_projection_relation_ids)
            ),
        )
        fallback_relation_id = (
            projection_anchor_relation_ids[-1]
            if projection_anchor_relation_ids
            else None
        )
        use_output_anchor = self._fragment_use_output_anchor(
            normalized=normalized,
            projection_specs=projection_specs,
            projection_assembly=projection_assembly,
            requires_row_identity_dedupe=requires_row_identity_dedupe,
            output_relation_id=output_relation_id,
        )
        query_mode = self._fragment_query_mode(
            use_output_anchor=use_output_anchor,
            projection_anchor_relation_ids=projection_anchor_relation_ids,
            output_relation_id=output_relation_id,
            fallback_relation_id=fallback_relation_id,
        )
        fragment_plan = SQL2FragmentPlan(
            query_mode=query_mode,
            projection_specs=tuple(projection_specs),
            projection_assembly=projection_assembly,
            projection_relation_task_var_ids_by_relation_id=tuple(
                sorted(projection_relation_task_var_ids_by_relation_id.items())
            ),
            projection_anchor_relation_ids=projection_anchor_relation_ids,
            projection_anchor_correlation_var_ids=projection_anchor_correlation_var_ids,
            output_anchor_projection_relation_ids=output_anchor_projection_relation_ids,
            output_anchor_projection_correlation_var_ids=output_anchor_projection_correlation_var_ids,
            fallback_relation_id=fallback_relation_id,
            requires_row_identity_dedupe=requires_row_identity_dedupe,
            output_relation_id=output_relation_id,
            output_relation_key_field_name=(
                output_update.relation.fields[0].name
                if output_update is not None and output_update.relation.fields
                else None
            ),
            output_relation_key_type_id=(
                getattr(output_update.relation.fields[0].type, "id", None)
                if output_update is not None and output_update.relation.fields
                else None
            ),
            output_relation_key_var_id=(
                output_update.args[0].id
                if output_update is not None and output_update.args and isinstance(output_update.args[0], Var)
                else None
            ),
        )
        return SQL2Analysis(
            model_analysis=query_analysis,
            table_shape=table_shape,
            graph=graph,
            coalesced=coalesced,
            normalized=normalized,
            model_key_preserving_table_names=self.sql2_analysis.model_key_preserving_table_names,
            fragment_plan=fragment_plan,
        )

    def _prune_redundant_fragment_populations(self, task: Task) -> Task:
        if isinstance(task, Logical):
            rewritten_body = tuple(
                self._prune_redundant_fragment_populations(child)
                for child in task.body
            )
            return task.mut(body=self._prune_redundant_population_lookups_in_body(rewritten_body))
        if isinstance(task, Aggregate):
            return task.mut(body=self._prune_redundant_fragment_populations(task.body))
        if isinstance(task, Exists):
            return task.mut(task=self._prune_redundant_fragment_populations(task.task))
        if isinstance(task, Not):
            return task.mut(task=self._prune_redundant_fragment_populations(task.task))
        if isinstance(task, Union):
            return task.mut(
                tasks=tuple(
                    self._prune_redundant_fragment_populations(branch)
                    for branch in task.tasks
                )
            )
        if isinstance(task, Match):
            return task.mut(
                tasks=tuple(
                    self._prune_redundant_fragment_populations(branch)
                    for branch in task.tasks
                )
            )
        return task

    def _prune_redundant_population_lookups_in_body(
        self,
        body: tuple[Task, ...],
    ) -> tuple[Task, ...]:
        rewritten: list[Task] = []
        for ix, child in enumerate(body):
            if not self._population_lookup_redundantly_narrowed(
                lookup=child,
                later_tasks=body[ix + 1 :],
            ):
                rewritten.append(child)
        return tuple(rewritten)

    def _population_lookup_redundantly_narrowed(
        self,
        *,
        lookup: Task,
        later_tasks: tuple[Task, ...],
    ) -> bool:
        if not isinstance(lookup, Lookup):
            return False
        if not isinstance(lookup.relation, ScalarType):
            return False
        if len(lookup.args) != 1 or not isinstance(lookup.args[0], Var):
            return False
        var = lookup.args[0]
        if not self._type_strictly_narrows_population(var.type, lookup.relation):
            return False
        for later_task in later_tasks:
            if not isinstance(later_task, Lookup):
                continue
            if isinstance(later_task.relation, ScalarType):
                continue
            if any(arg == var for arg in later_task.args):
                return True
        return False

    def _type_strictly_narrows_population(self, type_, population_type: ScalarType) -> bool:
        if type_ == population_type:
            return False
        if bt.extends(type_, population_type):
            return True
        if isinstance(type_, UnionType) and type_.types:
            return all(t == population_type or bt.extends(t, population_type) for t in type_.types) and any(
                t != population_type for t in type_.types
            )
        return False

    def _projection_correlation_var_ids(
        self,
        *,
        normalized: SQL2NormalizedArtifacts,
        projection_specs: tuple[SQL2FragmentProjectionSpec, ...],
    ) -> tuple[int, ...]:
        relation_ids = {spec.relation_id for spec in projection_specs}
        if not relation_ids:
            return ()
        var_ids_by_relation: dict[int, set[int]] = {relation_id: set() for relation_id in relation_ids}
        for group in normalized.groups:
            for channel in group.channels:
                if channel.target_relation_id not in relation_ids:
                    continue
                var_ids_by_relation[channel.target_relation_id].update(
                    task_var_ids(channel.rule)
                )

        non_empty_sets = [var_ids for var_ids in var_ids_by_relation.values() if var_ids]
        if not non_empty_sets:
            shared_var_ids: set[int] = set()
        elif len(non_empty_sets) == 1:
            shared_var_ids = set(non_empty_sets[0])
        else:
            shared_var_ids = set.intersection(*non_empty_sets)

        for spec in projection_specs:
            for var_id, _field_name, _is_primary in spec.key_var_columns:
                shared_var_ids.add(var_id)
        return tuple(sorted(shared_var_ids))


    def _fragment_query_mode(
        self,
        *,
        use_output_anchor: bool,
        projection_anchor_relation_ids: tuple[int, ...],
        output_relation_id: int | None,
        fallback_relation_id: int | None,
    ) -> str:
        if use_output_anchor:
            return "output_anchor"
        if projection_anchor_relation_ids:
            return "projection_anchor"
        if output_relation_id is not None:
            return "output_anchor"
        if fallback_relation_id is not None:
            return "fallback_relation"
        return "none"

    def _fragment_use_output_anchor(
        self,
        *,
        normalized: SQL2NormalizedArtifacts,
        projection_specs: list[SQL2FragmentProjectionSpec],
        projection_assembly: SQL2FragmentProjectionAssemblyPlan,
        requires_row_identity_dedupe: bool,
        output_relation_id: int | None,
    ) -> bool:
        if output_relation_id is None or not requires_row_identity_dedupe:
            return False
        optional_by_relation_id = {spec.relation_id: spec.optional for spec in projection_specs}
        needs_output_anchor = False
        for group in projection_assembly.join_groups:
            group_optional_flags = [
                optional_by_relation_id[relation_id]
                for relation_id in group.relation_ids
                if relation_id in optional_by_relation_id
            ]
            if not group_optional_flags:
                continue
            if any(group_optional_flags) and not any(not optional for optional in group_optional_flags):
                needs_output_anchor = True
                break
        if not needs_output_anchor:
            return False
        direct_projection_safe_by_relation_id = dict(
            normalized.relation_direct_projection_safe_by_relation_id
        )
        if not direct_projection_safe_by_relation_id.get(output_relation_id, False):
            return False
        output_anchor_safe_by_relation_id = dict(
            normalized.relation_output_anchor_safe_by_relation_id
        )
        return output_anchor_safe_by_relation_id.get(output_relation_id, False)

    def _fragment_requires_row_identity_dedupe(
        self,
        *,
        projection_relation_ids: set[int],
        normalized: SQL2NormalizedArtifacts,
    ) -> bool:
        if not projection_relation_ids:
            return True
        direct_projection_safe_by_relation_id = dict(
            normalized.relation_direct_projection_safe_by_relation_id
        )
        return any(
            not direct_projection_safe_by_relation_id.get(relation_id, False)
            for relation_id in projection_relation_ids
        )

    def _projection_coalesce_group_id_by_relation_id(
        self,
        *,
        projection_updates: list[Update],
        normalized: SQL2NormalizedArtifacts,
    ) -> dict[int, int]:
        channels_by_relation_id: dict[int, list[object]] = {}
        for group in normalized.groups:
            for channel in group.channels:
                channels_by_relation_id.setdefault(channel.target_relation_id, []).append(channel)

        relation_signature_by_relation_id: dict[int, tuple[object, ...]] = {}
        for relation_id, channels in channels_by_relation_id.items():
            signature = tuple(
                (
                    self._projection_channel_signature(channel),
                    self._projection_channel_aggregate_signature(channel),
                )
                for channel in channels
            )
            relation_signature_by_relation_id[relation_id] = signature

        group_id_by_signature: dict[tuple[object, ...], int] = {}
        out: dict[int, int] = {}
        next_group_id = 0
        for update in projection_updates:
            relation_id = update.relation.id
            signature = (
                getattr(update.relation.fields[0].type, "id", None) if update.relation.fields else None,
                self._projection_key_var_columns(update),
                len(self._value_field_names(update)),
                relation_signature_by_relation_id.get(relation_id, ()),
            )
            group_id = group_id_by_signature.get(signature)
            if group_id is None:
                group_id = next_group_id
                next_group_id += 1
                group_id_by_signature[signature] = group_id
            out[relation_id] = group_id
        return out

    def _projection_relation_ids_for_anchor(
        self,
        *,
        projection_specs: list[SQL2FragmentProjectionSpec],
        relation_ids_with_channels: set[int],
        required_key_type_id: int | None,
    ) -> tuple[int, ...]:
        relation_ids: list[int] = []
        for spec in projection_specs:
            if not spec.value_field_names:
                continue
            if required_key_type_id is not None and spec.relation_key_type_id != required_key_type_id:
                continue
            if spec.relation_id not in relation_ids_with_channels:
                continue
            relation_ids.append(spec.relation_id)
        return tuple(relation_ids)

    def _projection_channel_signature(self, channel) -> tuple[object, ...]:
        union_modes = tuple(
            sorted(mode for _task_id, mode in channel.union_mode_by_task_id)
        )
        return (
            channel.read_table_names,
            union_modes,
        )

    def _relation_ids_with_channels(self, normalized: SQL2NormalizedArtifacts) -> set[int]:
        relation_ids: set[int] = set()
        for group in normalized.groups:
            for channel in group.channels:
                relation_ids.add(channel.target_relation_id)
        return relation_ids

    def _projection_channel_aggregate_signature(self, channel) -> tuple[object, ...]:
        return tuple(
            (
                inst.shape.aggregation_relation_id,
                inst.shape.group_shape,
                inst.shape.projection_shape,
                inst.shape.non_output_arg_shape,
            )
            for inst in channel.aggregate_instances
        )

    def _projection_relation_task_var_ids_by_relation_id(
        self,
        *,
        normalized: SQL2NormalizedArtifacts,
        relation_ids: set[int],
    ) -> dict[int, tuple[int, ...]]:
        channel_var_sets_by_relation_id: dict[int, list[set[int]]] = {
            relation_id: [] for relation_id in relation_ids
        }
        for group in normalized.groups:
            for channel in group.channels:
                relation_id = channel.target_relation_id
                if relation_id not in channel_var_sets_by_relation_id:
                    continue
                channel_var_sets_by_relation_id[relation_id].append(set(bindable_var_ids(channel.rule)))
        var_ids_by_relation_id: dict[int, tuple[int, ...]] = {}
        for relation_id, channel_var_sets in channel_var_sets_by_relation_id.items():
            if not channel_var_sets:
                var_ids_by_relation_id[relation_id] = ()
                continue
            shared_var_ids = set.intersection(*channel_var_sets)
            var_ids_by_relation_id[relation_id] = tuple(sorted(shared_var_ids))
        return {
            relation_id: var_ids
            for relation_id, var_ids in var_ids_by_relation_id.items()
        }

    def _build_projection_assembly_plan(
        self,
        projection_specs: list[SQL2FragmentProjectionSpec],
    ) -> SQL2FragmentProjectionAssemblyPlan:
        if not projection_specs:
            return SQL2FragmentProjectionAssemblyPlan(
                anchor_relation_id=None,
                join_groups=(),
                source_groups=(),
                output_columns=(),
            )
        specs_by_group_id: dict[int, list[SQL2FragmentProjectionSpec]] = {}
        group_order: list[int] = []
        for spec in projection_specs:
            group_specs = specs_by_group_id.get(spec.coalesce_group_id)
            if group_specs is None:
                specs_by_group_id[spec.coalesce_group_id] = [spec]
                group_order.append(spec.coalesce_group_id)
            else:
                group_specs.append(spec)
        join_groups: list[SQL2FragmentProjectionJoinGroup] = []
        for group_id in group_order:
            group_specs = specs_by_group_id[group_id]
            seen_var_ids: set[int] = set()
            join_var_ids: list[int] = []
            for spec in group_specs:
                for var_id, _field_name, _is_primary in spec.key_var_columns:
                    if var_id in seen_var_ids:
                        continue
                    seen_var_ids.add(var_id)
                    join_var_ids.append(var_id)
            join_groups.append(
                SQL2FragmentProjectionJoinGroup(
                    group_id=group_id,
                    relation_ids=tuple(spec.relation_id for spec in group_specs),
                    join_var_ids=tuple(join_var_ids),
                    output_anchor_join_type="INNER"
                    if any(not spec.optional for spec in group_specs)
                    else "LEFT",
                )
            )
        source_groups: list[SQL2FragmentProjectionSourceGroup] = []
        next_source_group_id = 0
        for join_group in join_groups:
            if not join_group.relation_ids:
                continue
            source_groups.append(
                SQL2FragmentProjectionSourceGroup(
                    group_id=next_source_group_id,
                    join_group_id=join_group.group_id,
                    relation_ids=join_group.relation_ids,
                    output_anchor_join_type=join_group.output_anchor_join_type,
                )
            )
            next_source_group_id += 1
        output_columns = tuple(
            SQL2FragmentProjectionOutputColumn(
                relation_id=spec.relation_id,
                field_name=field_name,
                output_alias=spec.relation_name,
            )
            for spec in projection_specs
            for field_name in spec.value_field_names
        )
        anchor_relation_id = (
            join_groups[0].relation_ids[0]
            if join_groups and join_groups[0].relation_ids
            else None
        )
        return SQL2FragmentProjectionAssemblyPlan(
            anchor_relation_id=anchor_relation_id,
            join_groups=tuple(join_groups),
            source_groups=tuple(source_groups),
            output_columns=output_columns,
        )

    def _model_key_preserving_table_names(
        self,
        *,
        normalized: SQL2NormalizedArtifacts,
        table_shape: SQL2TableShapeSidecar,
    ) -> tuple[str, ...]:
        key_preserving_by_relation_id = dict(normalized.relation_key_preserving_by_relation_id)
        table_names = sorted(
            table.table_name
            for table in table_shape.tables
            if key_preserving_by_relation_id.get(table.owner_relation_id, False)
        )
        return tuple(table_names)

    def _collect_updates(self, task: Task) -> list[Update]:
        updates: list[Update] = []

        def visit(node: Task) -> None:
            if isinstance(node, Update):
                updates.append(node)
                return
            if isinstance(node, Logical):
                for child in node.body:
                    visit(child)
                return
            if isinstance(node, Exists):
                visit(node.task)
                return
            if isinstance(node, Not):
                visit(node.task)
                return
            if isinstance(node, Aggregate):
                visit(node.body)
                return
            if isinstance(node, Union):
                for child in node.tasks:
                    visit(child)
                return
            if isinstance(node, Match):
                for child in node.tasks:
                    visit(child)

        visit(task)
        return updates

    def _collect_updates_with_optional_context(self, task: Task) -> list[tuple[Update, bool]]:
        updates: list[tuple[Update, bool]] = []

        def visit(node: Task, optional_context: bool) -> None:
            if isinstance(node, Update):
                updates.append((node, optional_context))
                return
            if isinstance(node, Logical):
                next_optional = optional_context or node.optional
                for child in node.body:
                    visit(child, next_optional)
                return
            if isinstance(node, Exists):
                visit(node.task, optional_context)
                return
            if isinstance(node, Not):
                visit(node.task, optional_context)
                return
            if isinstance(node, Aggregate):
                visit(node.body, optional_context)
                return
            if isinstance(node, Union):
                for child in node.tasks:
                    visit(child, optional_context)
                return
            if isinstance(node, Match):
                for child in node.tasks:
                    visit(child, optional_context)

        visit(task, False)
        return updates

    def _resolve_projection_updates(
        self,
        *,
        requested_projection_relation_names: list[str],
        normalized_updates: list[Update],
    ) -> list[Update]:
        if not requested_projection_relation_names:
            return []
        remaining = [update for update in normalized_updates if not self._is_output_marker_update(update)]
        resolved: list[Update] = []
        for relation_name in requested_projection_relation_names:
            match_ix = next(
                (ix for ix, update in enumerate(remaining) if update.relation.name == relation_name),
                None,
            )
            if match_ix is None:
                continue
            resolved.append(remaining.pop(match_ix))
        return resolved

    def _is_output_marker_update(self, update: Update) -> bool:
        return update.relation.name.startswith("Output") and len(update.relation.fields) == 1

    def _value_field_names(self, update: Update) -> tuple[str, ...]:
        relation = update.relation
        names = tuple(
            field.name
            for ix, field in enumerate(relation.fields)
            if ix != 0 and not field.input
        )
        if names:
            if len(names) > 1:
                return (names[-1],)
            return names
        if len(relation.fields) >= 2:
            return (relation.fields[-1].name,)
        return ()

    def _projection_key_var_columns(self, update: Update) -> tuple[tuple[int, str, bool], ...]:
        if self._projection_relation_uses_output_only_key(update):
            if not update.args or not isinstance(update.args[0], Var) or not update.relation.fields:
                return ()
            return ((update.args[0].id, update.relation.fields[0].name, True),)
        key_arg_count = max(0, len(update.args) - 1)
        out: list[tuple[int, str, bool]] = []
        for ix, (field, arg) in enumerate(zip(update.relation.fields, update.args)):
            if ix >= key_arg_count or not isinstance(arg, Var):
                continue
            out.append((arg.id, field.name, ix == 0))
        return tuple(out)

    def _projection_relation_uses_output_only_key(self, update: Update) -> bool:
        relation = update.relation
        if len(relation.fields) <= 2:
            return False
        if not relation.name.startswith("v") or not relation.name[1:].isdigit():
            return False
        first_field_type_name = getattr(relation.fields[0].type, "name", "")
        return isinstance(first_field_type_name, str) and first_field_type_name.startswith("Output")


SQLModelAnalyzer = SQL2ModelAnalyzer

__all__ = [
    "SQL2Analysis",
    "SQL2FragmentPlan",
    "SQL2FragmentProjectionAssemblyPlan",
    "SQL2FragmentProjectionJoinGroup",
    "SQL2FragmentProjectionOutputColumn",
    "SQL2FragmentProjectionSourceGroup",
    "SQL2FragmentProjectionSpec",
    "SQL2ModelAnalyzer",
    "SQLModelAnalyzer",
]
