import re
import shlex
from typing import List, Tuple, Optional

from .ir import (
    Pipeline,
    SearchStage,
    FieldsStage,
    StatsStage,
    SortStage,
    LimitStage,
    FilterStage,
    AggFunc,
    Filter,
    BoolExpr,
    and_filters,
    make_text_filter,
)

COND_RE = re.compile(r"'([^']+)'\s*(>=|<=|=|!=|>|<)\s*(\"[^\"]*\"|[^\s|]+)")
KV_LOOSE_RE = re.compile(r"([A-Za-z_][\w.]*)=(\"[^\"]*\"|[^\s|)]+)")
CMP_LOOSE_RE = re.compile(r"([A-Za-z_][\w.]*)\s*(>=|<=|=|!=|>|<)\s*([^\s|)]+)")
FUNC_RE = re.compile(r"(\w+)\(([^)]*)\)\s*(?:as\s+(?:'([^']+)'|\"([^\"]+)\"|([\w.]+)))?")
IN_FUNC_RE = re.compile(r"in\(\s*([A-Za-z_][\w.]*)\s*,\s*([^)]+)\)")
REGEX_FUNC_RE = re.compile(r"regex\(\s*([A-Za-z_][\w.]*)\s*,\s*(\"[^\"]*\"|[^)]+)\)")


def _split_stages(text: str) -> List[str]:
    parts = []
    buf = []
    i = 0
    in_quote = False
    while i < len(text):
        ch = text[i]
        if ch == '"':
            in_quote = not in_quote
            buf.append(ch)
        elif ch == '|' and not in_quote:
            parts.append(''.join(buf).strip())
            buf = []
        else:
            buf.append(ch)
        i += 1
    if buf:
        parts.append(''.join(buf).strip())
    return parts

# Boolean tokenizer and parser
from typing import Union as _Union

def _strip_quotes(v: str) -> str:
    if v is None:
        return v
    if (v.startswith('"') and v.endswith('"')) or (v.startswith("'") and v.endswith("'")):
        return v[1:-1]
    return v

def _tokenize_bool(body: str) -> List[_Union[Filter, str]]:
    tokens: List[_Union[Filter, str]] = []
    s = body.strip()
    pos = 0
    while pos < len(s):
        # skip whitespace
        m = re.match(r"\s+", s[pos:])
        if m:
            pos += m.end()
            continue
        # parentheses
        if s[pos] == '(':
            tokens.append('LPAREN')
            pos += 1
            continue
        if s[pos] == ')':
            tokens.append('RPAREN')
            pos += 1
            continue
        # connectors
        m = re.match(r"(?i)\b(AND|OR|NOT)\b", s[pos:])
        if m:
            tokens.append(m.group(1).upper())
            pos += m.end()
            continue
        # in(field, ...)
        m = IN_FUNC_RE.match(s[pos:])
        if m:
            field = m.group(1)
            values_raw = m.group(2)
            vals = []
            for v in [x.strip() for x in values_raw.split(',') if x.strip()]:
                vals.append(_strip_quotes(v))
            tokens.append(Filter(op='in', field=field, value=vals))
            pos += m.end()
            continue
        # regex(field, pattern)
        m = REGEX_FUNC_RE.match(s[pos:])
        if m:
            field = m.group(1)
            pattern = _strip_quotes(m.group(2))
            tokens.append(Filter(op='regex', field=field, value=pattern))
            pos += m.end()
            continue
        # quoted field comparison
        m = COND_RE.match(s[pos:])
        if m:
            field = m.group(1)
            op = m.group(2)
            value = _strip_quotes(m.group(3))
            op_map = {'=': 'eq', '!=': 'neq', '>': 'gt', '>=': 'gte', '<': 'lt', '<=': 'lte'}
            tokens.append(Filter(op=op_map.get(op, op), field=field, value=value))
            pos += m.end()
            continue
        # loose kv
        m = KV_LOOSE_RE.match(s[pos:])
        if m:
            field = m.group(1)
            value = _strip_quotes(m.group(2))
            tokens.append(Filter(op='eq', field=field, value=value))
            pos += m.end()
            continue
        # loose cmp
        m = CMP_LOOSE_RE.match(s[pos:])
        if m:
            field = m.group(1)
            op = m.group(2)
            value = _strip_quotes(m.group(3))
            op_map = {'=': 'eq', '!=': 'neq', '>': 'gt', '>=': 'gte', '<': 'lt', '<=': 'lte'}
            tokens.append(Filter(op=op_map.get(op, op), field=field, value=value))
            pos += m.end()
            continue
        # quoted text token
        m = re.match(r'"[^"\\]*(?:\\.[^"\\]*)*"', s[pos:])
        if m:
            val = _strip_quotes(m.group(0))
            tokens.append(make_text_filter(val))
            pos += m.end()
            continue
        # bareword as text filter
        m = re.match(r"[^\s|()]+", s[pos:])
        if m:
            tokens.append(make_text_filter(m.group(0)))
            pos += m.end()
            continue
        # fallback advance to avoid infinite loop
        pos += 1
    # insert implicit ANDs
    out: List[_Union[Filter, str]] = []
    prev_atom = False
    for t in tokens:
        is_atom = isinstance(t, Filter) or t == 'LPAREN'
        if prev_atom and (is_atom or t == 'NOT'):
            out.append('AND')
        out.append(t)
        prev_atom = isinstance(t, Filter) or t == 'RPAREN'
    return out

def _build_bool_expr(tokens: List[_Union[Filter, str]]) -> Optional[BoolExpr]:
    if not tokens:
        return None
    prec = {'OR': 1, 'AND': 2, 'NOT': 3}
    output: List[_Union[Filter, BoolExpr, str]] = []
    ops: List[str] = []
    for t in tokens:
        if isinstance(t, Filter):
            output.append(t)
        elif t in ('AND', 'OR', 'NOT'):
            while ops and ops[-1] != 'LPAREN' and prec.get(ops[-1], 0) >= prec[t]:
                output.append(ops.pop())
            ops.append(t)
        elif t == 'LPAREN':
            ops.append(t)
        elif t == 'RPAREN':
            while ops and ops[-1] != 'LPAREN':
                output.append(ops.pop())
            if ops and ops[-1] == 'LPAREN':
                ops.pop()
    while ops:
        output.append(ops.pop())
    # Build AST from RPN
    stack: List[_Union[Filter, BoolExpr]] = []
    def _merge(op: str, left: _Union[Filter, BoolExpr], right: Optional[_Union[Filter, BoolExpr]] = None) -> BoolExpr:
        if op == 'NOT':
            child = left
            if isinstance(child, BoolExpr) and child.op == 'not':
                return child
            return BoolExpr(op='not', children=[child])
        # binary
        children: List[_Union[Filter, BoolExpr]] = []
        for ch in (left, right):
            if isinstance(ch, BoolExpr) and ch.op == op.lower():
                children.extend(ch.children)
            else:
                children.append(ch)
        return BoolExpr(op=op.lower(), children=children) 
    for item in output:
        if isinstance(item, Filter):
            stack.append(item)
        elif item in ('AND', 'OR'):
            r = stack.pop() if stack else None
            l = stack.pop() if stack else None
            if l is None: l = make_text_filter('')
            if r is None: r = make_text_filter('')
            stack.append(_merge(item, l, r))
        elif item == 'NOT':
            a = stack.pop() if stack else make_text_filter('')
            stack.append(_merge('NOT', a))
    if not stack:
        return None
    top = stack[-1]
    return top if isinstance(top, BoolExpr) else BoolExpr(op='and', children=[top])


def parse_search2(stage_text: str) -> SearchStage:
    # Remove leading command
    body = stage_text.strip()[len('search2'):].strip()

    start = None
    end = None
    repo = None
    # Extract start/end/repo key-value pairs (allow spaces around '=')
    for key in ['start', 'end', 'repo']:
        m = re.search(rf"\b{key}\s*=\s*\"([^\"]*)\"", body)
        if m:
            val = m.group(1)
            if key == 'start':
                start = val
            elif key == 'end':
                end = val
            elif key == 'repo':
                repo = val
            body = body.replace(m.group(0), '').strip()

    # Parse boolean expression with connectors and parentheses from remaining body
    tokens = _tokenize_bool(body)
    expr = _build_bool_expr(tokens)
    return SearchStage(start=start, end=end, repo=repo, query=expr)


def parse_fields(stage_text: str) -> FieldsStage:
    body = stage_text.strip()[len('fields'):].strip()
    fields = [f.strip() for f in body.split(',') if f.strip()]
    return FieldsStage(fields=fields)


def parse_stats(stage_text: str) -> StatsStage:
    body = stage_text.strip()[len('stats'):].strip()
    by: List[str] = []
    funcs: List[AggFunc] = []

    if ' by ' in body:
        func_part, by_part = body.split(' by ', 1)
        by = [b.strip() for b in by_part.split(',') if b.strip()]
    else:
        func_part = body

    # Multiple functions separated by commas
    for token in [t.strip() for t in func_part.split(',') if t.strip()]:
        m = FUNC_RE.match(token)
        if m:
            name = m.group(1)
            field = m.group(2).strip() or None
            alias = m.group(3) or m.group(4) or m.group(5)
            funcs.append(AggFunc(name=name, field=field, alias=alias))

    return StatsStage(funcs=funcs, by=by)


def parse_sort(stage_text: str) -> SortStage:
    body = stage_text.strip()[len('sort'):].strip()
    size = None
    # optional leading size
    msz = re.match(r"^(\d+)\b\s*(.*)$", body)
    if msz:
        try:
            size = int(msz.group(1))
        except Exception:
            size = None
        body = msz.group(2).strip()
    if body.startswith('by'):
        body = body[2:].strip()
    pairs: List[Tuple[str, str]] = []
    for token in [t.strip() for t in body.split(',') if t.strip()]:
        parts = token.split()
        if not parts:
            continue
        field = parts[0]
        order = 'asc'
        if len(parts) > 1 and parts[1].lower() in ('asc', 'desc'):
            order = parts[1].lower()
        pairs.append((field, order))
    st = SortStage(by=pairs)
    if size is not None:
        st.params['size'] = size
    return st


def parse_limit(stage_text: str) -> LimitStage:
    body = stage_text.strip()[len('limit'):].strip()
    try:
        size = int(body)
    except Exception:
        size = 100
    return LimitStage(size=size)


def parse_filter(stage_text: str) -> FilterStage:
    body = stage_text.strip()[len('filter'):].strip()
    tokens = _tokenize_bool(body)
    expr = _build_bool_expr(tokens)
    fs = FilterStage(expr=expr)
    fs.type = 'where'
    return fs


def parse_where(stage_text: str) -> FilterStage:
    # identical to filter, but keep type as 'where' for rendering
    fs = parse_filter('filter ' + stage_text.strip()[len('where'):].strip())
    fs.type = 'where'
    return fs


def parse_spl(text: str) -> Pipeline:
    stages_text = _split_stages(text.strip())
    pipeline = Pipeline()
    for st in stages_text:
        if not st:
            continue
        head = st.split()[0].lower()
        if head == 'search2':
            pipeline.stages.append(parse_search2(st))
        elif head == 'fields':
            pipeline.stages.append(parse_fields(st))
        elif head == 'stats':
            pipeline.stages.append(parse_stats(st))
        elif head == 'sort':
            pipeline.stages.append(parse_sort(st))
        elif head == 'limit':
            pipeline.stages.append(parse_limit(st))
        elif head == 'filter':
            # Normalize filter to where for SPL
            pipeline.stages.append(parse_where(st))
        elif head == 'where':
            pipeline.stages.append(parse_where(st))
        else:
            # Fallback: keep as generic stage
            from .ir import GenericStage
            pipeline.stages.append(GenericStage(type=head, params={'raw': st}))
    return pipeline