from dataclasses import dataclass, field
from typing import Any, Dict, List, Optional, Tuple, Union

# ========== Core IR (Query) ==========

@dataclass
class TimeRange:
    start: Optional[str] = None
    end: Optional[str] = None

@dataclass
class Source:
    repo: Optional[str] = None

@dataclass
class Filter:
    op: str  # eq, neq, gt, gte, lt, lte, in, like, regex, exists, text
    field: Optional[str]
    value: Any

@dataclass
class BoolExpr:
    op: str  # and, or, not
    children: List[Union['Filter', 'BoolExpr']] = field(default_factory=list)

@dataclass
class Projection:
    fields: List[str] = field(default_factory=list)

@dataclass
class AggFunc:
    name: str  # count, sum, avg, max, min, last, first, custom
    field: Optional[str] = None
    alias: Optional[str] = None

@dataclass
class Aggregation:
    funcs: List[AggFunc] = field(default_factory=list)
    by: List[str] = field(default_factory=list)

@dataclass
class Sort:
    by: List[Tuple[str, str]] = field(default_factory=list)  # (field, 'asc'|'desc')

@dataclass
class Limit:
    size: int = 100

@dataclass
class QueryIR:
    source: Source = field(default_factory=Source)
    time_range: Optional[TimeRange] = None
    where: Optional[BoolExpr] = None
    projection: Optional[Projection] = None
    aggregation: Optional[Aggregation] = None
    sort: Optional[Sort] = None
    limit: Optional[Limit] = None

# ========== Pipeline Stages ==========

@dataclass
class Stage:
    type: str
    params: Dict[str, Any] = field(default_factory=dict)

@dataclass
class SearchStage(Stage):
    type: str = "search2"
    start: Optional[str] = None
    end: Optional[str] = None
    repo: Optional[str] = None
    query: Optional[BoolExpr] = None

@dataclass
class FieldsStage(Stage):
    type: str = "fields"
    fields: List[str] = field(default_factory=list)

@dataclass
class FilterStage(Stage):
    type: str = "filter"
    expr: Optional[BoolExpr] = None

@dataclass
class StatsStage(Stage):
    type: str = "stats"
    funcs: List[AggFunc] = field(default_factory=list)
    by: List[str] = field(default_factory=list)

@dataclass
class SortStage(Stage):
    type: str = "sort"
    by: List[Tuple[str, str]] = field(default_factory=list)

@dataclass
class LimitStage(Stage):
    type: str = "limit"
    size: int = 100

@dataclass
class GenericStage(Stage):
    type: str = "generic"

@dataclass
class Pipeline:
    stages: List[Stage] = field(default_factory=list)

# ========== Helpers ==========

def make_text_filter(text: str) -> Filter:
    return Filter(op="text", field=None, value=text)


def and_filters(filters: List[Union[Filter, BoolExpr]]) -> Optional[BoolExpr]:
    if not filters:
        return None
    # Flatten nested ANDs
    flat: List[Union[Filter, BoolExpr]] = []
    for f in filters:
        if isinstance(f, BoolExpr) and f.op == "and":
            flat.extend(f.children)
        else:
            flat.append(f)
    return BoolExpr(op="and", children=flat)