from __future__ import annotations

from pydantic import BaseModel, Field

from sir import EdgeKind, SemanticNode, SirSnapshot


class TransactionWriteSet(BaseModel):
    routine: SemanticNode
    writes: list[SemanticNode] = Field(default_factory=list)


class TransactionTargetWriters(BaseModel):
    target: SemanticNode
    routines: list[SemanticNode] = Field(default_factory=list)


def transaction_write_sets(snapshot: SirSnapshot) -> list[TransactionWriteSet]:
    nodes = {node.lineage_id: node for node in snapshot.nodes}
    result: dict[str, TransactionWriteSet] = {}
    for edge in snapshot.edges:
        if edge.kind != EdgeKind.WRITES:
            continue
        routine = nodes.get(edge.source_lineage)
        target = nodes.get(edge.target_lineage)
        if routine is None or target is None:
            continue
        result.setdefault(routine.lineage_id, TransactionWriteSet(routine=routine)).writes.append(target)
    write_sets = list(result.values())
    write_sets.sort(key=lambda item: item.routine.qualified_name)
    return write_sets


def transaction_targets(snapshot: SirSnapshot) -> list[TransactionTargetWriters]:
    result: dict[str, TransactionTargetWriters] = {}
    for write_set in transaction_write_sets(snapshot):
        for target in write_set.writes:
            result.setdefault(target.lineage_id, TransactionTargetWriters(target=target)).routines.append(write_set.routine)
    targets = list(result.values())
    targets.sort(key=lambda item: item.target.qualified_name)
    return targets


def routines_touching_target(snapshot: SirSnapshot, target_name: str) -> list[SemanticNode]:
    wanted = target_name.casefold()
    routines: list[SemanticNode] = []
    for write_set in transaction_write_sets(snapshot):
        if any(
            target.name.casefold() == wanted or target.qualified_name.casefold() == wanted
            for target in write_set.writes
        ):
            routines.append(write_set.routine)
    return routines


__all__ = [
    "TransactionTargetWriters",
    "TransactionWriteSet",
    "routines_touching_target",
    "transaction_targets",
    "transaction_write_sets",
]