"""Waveform Query Language — SQL-like queries on trace data.

Supports SELECT/WHERE/GROUP BY/ORDER BY/LIMIT and temporal MATCH SEQUENCE.

Usage:
    results = t.query("SELECT index, ts, tb.dut.counter WHERE tb.dut.clk = 1 LIMIT 10")
"""

from __future__ import annotations
import re
from typing import Any, TYPE_CHECKING

if TYPE_CHECKING:
    from ewal.trace import Trace


def query(trace: "Trace", sql: str):
    """Execute a SQL-like query against the trace.

    Supported syntax:
        SELECT col1, col2, ... [AS alias]
        WHERE condition [AND condition ...]
        GROUP BY signal
        ORDER BY col [ASC|DESC]
        LIMIT n

        MATCH SEQUENCE
            step0: condition
            step1: condition WITHIN n CYCLES
        CAPTURE
            name = expr AT step_name
    """
    import pandas as pd

    sql = sql.strip()

    # Check for MATCH SEQUENCE
    if "MATCH SEQUENCE" in sql.upper():
        return _match_sequence_query(trace, sql)

    # Parse SQL-like query
    return _select_query(trace, sql)


def _select_query(trace: "Trace", sql: str):
    """Parse and execute a SELECT query."""
    import pandas as pd

    sql_upper = sql.upper()

    # Extract SELECT columns
    select_match = re.search(r'SELECT\s+(.+?)(?:\s+WHERE|\s+GROUP|\s+ORDER|\s+LIMIT|\s*$)',
                              sql, re.IGNORECASE | re.DOTALL)
    if not select_match:
        raise ValueError("Invalid query: missing SELECT clause")

    select_raw = select_match.group(1).strip()
    columns = [c.strip() for c in select_raw.split(",")]

    # Parse column aliases
    col_defs = []
    for col in columns:
        alias_match = re.match(r'(.+?)\s+AS\s+(\w+)', col, re.IGNORECASE)
        if alias_match:
            col_defs.append((alias_match.group(1).strip(), alias_match.group(2).strip()))
        else:
            col_defs.append((col, col))

    # Extract WHERE
    where_clause = None
    where_match = re.search(r'WHERE\s+(.+?)(?:\s+GROUP|\s+ORDER|\s+LIMIT|\s*$)',
                             sql, re.IGNORECASE | re.DOTALL)
    if where_match:
        where_clause = where_match.group(1).strip()

    # Extract GROUP BY
    group_by = None
    group_match = re.search(r'GROUP\s+BY\s+(\S+)', sql, re.IGNORECASE)
    if group_match:
        group_by = group_match.group(1).strip()

    # Extract ORDER BY
    order_by = None
    order_desc = False
    order_match = re.search(r'ORDER\s+BY\s+(\S+)(?:\s+(ASC|DESC))?', sql, re.IGNORECASE)
    if order_match:
        order_by = order_match.group(1).strip()
        if order_match.group(2) and order_match.group(2).upper() == "DESC":
            order_desc = True

    # Extract LIMIT
    limit = None
    limit_match = re.search(r'LIMIT\s+(\d+)', sql, re.IGNORECASE)
    if limit_match:
        limit = int(limit_match.group(1))

    # Build result
    if group_by:
        return _grouped_query(trace, col_defs, where_clause, group_by, order_by, order_desc, limit)

    rows = []
    for idx in range(trace.max_index + 1):
        # Evaluate WHERE
        if where_clause and not _eval_where(trace, where_clause, idx):
            continue

        row = {}
        for expr, alias in col_defs:
            row[alias] = _eval_column(trace, expr, idx)
        rows.append(row)

    df = pd.DataFrame(rows)

    if order_by and order_by in df.columns:
        df = df.sort_values(order_by, ascending=not order_desc)

    if limit:
        df = df.head(limit)

    return df


def _grouped_query(trace, col_defs, where_clause, group_by, order_by, order_desc, limit):
    """Execute a GROUP BY query with aggregations."""
    import pandas as pd

    # Collect all matching rows
    groups = {}
    for idx in range(trace.max_index + 1):
        if where_clause and not _eval_where(trace, where_clause, idx):
            continue

        key = _eval_column(trace, group_by, idx)
        if key not in groups:
            groups[key] = []
        groups[key].append(idx)

    # Compute aggregates
    rows = []
    for key, indices in groups.items():
        row = {}
        for expr, alias in col_defs:
            expr_upper = expr.upper().strip()
            if expr_upper == "COUNT(*)":
                row[alias] = len(indices)
            elif expr_upper.startswith("AVG("):
                sig = expr[4:-1].strip()
                vals = [float(trace._get_signal_value(sig, i) or 0) for i in indices]
                row[alias] = sum(vals) / len(vals) if vals else 0
            elif expr_upper.startswith("SUM("):
                sig = expr[4:-1].strip()
                vals = [float(trace._get_signal_value(sig, i) or 0) for i in indices]
                row[alias] = sum(vals)
            elif expr_upper.startswith("MIN("):
                sig = expr[4:-1].strip()
                vals = [float(trace._get_signal_value(sig, i) or 0) for i in indices]
                row[alias] = min(vals) if vals else 0
            elif expr_upper.startswith("MAX("):
                sig = expr[4:-1].strip()
                vals = [float(trace._get_signal_value(sig, i) or 0) for i in indices]
                row[alias] = max(vals) if vals else 0
            else:
                row[alias] = _eval_column(trace, expr, indices[0])
        rows.append(row)

    df = pd.DataFrame(rows)
    if order_by and order_by in df.columns:
        df = df.sort_values(order_by, ascending=not order_desc)
    if limit:
        df = df.head(limit)
    return df


def _eval_column(trace, expr: str, idx: int):
    """Evaluate a column expression at a given index."""
    expr = expr.strip()
    if expr.lower() == "index":
        return idx
    if expr.lower() == "ts" or expr.lower() == "timestamp":
        return trace._core.timestamp(idx)
    # Signal reference
    return trace._get_signal_value(expr, idx)


def _eval_where(trace, clause: str, idx: int) -> bool:
    """Evaluate a WHERE clause at a given index."""
    # Split on AND
    conditions = re.split(r'\s+AND\s+', clause, flags=re.IGNORECASE)

    for cond in conditions:
        cond = cond.strip()
        # Parse: signal op value
        m = re.match(r'(\S+)\s*(=|!=|<>|>=|<=|>|<)\s*(\S+)', cond)
        if not m:
            continue

        lhs = _eval_column(trace, m.group(1), idx)
        op = m.group(2)
        rhs_str = m.group(3)

        # Parse RHS
        try:
            rhs = int(rhs_str)
        except ValueError:
            try:
                rhs = float(rhs_str)
            except ValueError:
                rhs = rhs_str

        # Compare
        try:
            lhs_num = float(lhs) if lhs is not None else 0
            rhs_num = float(rhs)

            if op == "=" and lhs_num != rhs_num:
                return False
            if op in ("!=", "<>") and lhs_num == rhs_num:
                return False
            if op == ">" and lhs_num <= rhs_num:
                return False
            if op == "<" and lhs_num >= rhs_num:
                return False
            if op == ">=" and lhs_num < rhs_num:
                return False
            if op == "<=" and lhs_num > rhs_num:
                return False
        except (TypeError, ValueError):
            if op == "=" and lhs != rhs:
                return False

    return True


def _match_sequence_query(trace, sql: str):
    """Parse and execute a MATCH SEQUENCE temporal query."""
    import pandas as pd
    from ewal.temporal import seq, within, SequencePattern, ConditionStep, WithinStep

    # Extract steps
    steps_match = re.search(r'MATCH\s+SEQUENCE\s+(.+?)(?:CAPTURE|$)',
                            sql, re.IGNORECASE | re.DOTALL)
    if not steps_match:
        return pd.DataFrame()

    steps_text = steps_match.group(1).strip()
    step_lines = [l.strip() for l in steps_text.split("\n") if l.strip()]

    # Parse each step
    from ewal.signal import SignalRef
    parsed_steps = []
    for line in step_lines:
        # Format: step_name: signal = value [WITHIN n CYCLES]
        within_match = re.search(r'WITHIN\s+(\d+)\s+CYCLES', line, re.IGNORECASE)
        cond_part = re.sub(r'WITHIN\s+\d+\s+CYCLES', '', line, flags=re.IGNORECASE).strip()

        # Remove step label
        if ":" in cond_part:
            cond_part = cond_part.split(":", 1)[1].strip()

        # Build a callable condition
        cond_match = re.match(r'(\S+)\s*=\s*(\S+)', cond_part)
        if cond_match:
            sig_name = cond_match.group(1)
            val = int(cond_match.group(2))

            def make_cond(sn, v):
                def cond():
                    return trace._get_signal_value(sn, trace._index) == v
                return cond

            condition = make_cond(sig_name, val)

            if within_match:
                max_cycles = int(within_match.group(1))
                parsed_steps.append(WithinStep(max_cycles, condition))
            else:
                parsed_steps.append(ConditionStep(condition))

    if not parsed_steps:
        return pd.DataFrame()

    pattern = SequencePattern(parsed_steps)

    # Execute
    from ewal.temporal import match_pattern
    matches = match_pattern(trace, pattern)

    rows = [{"start": m.start, "end": m.end, **m.captures} for m in matches]
    return pd.DataFrame(rows)


def _patch_query():
    from ewal.trace import Trace

    def _query(self, sql):
        return query(self, sql)

    Trace.query = _query


_patch_query()
