import re
from typing import Union, List, Optional, Dict


class DuckDBQueryBuilder:
    """
    Unified SQL builder for DuckViz.
    Handles:
    - Standard aggregations
    - Raw SQL expressions
    - Grouping logic
    - Alias extraction
    - Optional WHERE
    """

    AGG_MAP = {
        "sum": "SUM",
        "avg": "AVG",
        "min": "MIN",
        "max": "MAX",
        "median": "MEDIAN",
        "count": "COUNT",
    }

    # --------------------------------------------------------
    # Utilities
    # --------------------------------------------------------

    def _normalize(self, value: Union[str, List[str], None]) -> List[str]:
        if value is None:
            return []
        return [value] if isinstance(value, str) else value

    def _extract_alias(self, expression: str) -> str:
        match = re.search(r"\s+AS\s+(\w+)$", expression, re.IGNORECASE)
        if match:
            return match.group(1)

        cleaned = re.sub(r"[^\w]", "_", expression)
        return cleaned[:30]

    def _where_clause(self, where: Optional[str]) -> str:
        return f"WHERE {where}" if where else ""
    
    def build(
        self,
        x: Optional[str] = None,
        y: Optional[Union[str, List[str]]] = None,
        agg: Optional[Union[str, List[str]]] = "sum",
        where: Optional[str] = None,
        group: bool = True,
    ) -> Dict:
        """
        Parameters
        ----------
        x : grouping column or SQL expression
        y : column(s) or SQL expression(s)
        agg : aggregation(s) OR None for raw SQL
        where : optional WHERE clause
        group : whether to GROUP BY x

        Returns
        -------
        dict with:
            query
            x_alias
            y_aliases
        """

        y_list = self._normalize(y)

        if agg is None:

            select_parts = []

            if x:
                select_parts.append(x)

            select_parts += y_list

            aliases = [self._extract_alias(expr) for expr in select_parts]

            query = f"""
                SELECT {", ".join(select_parts)}
                FROM data
                {self._where_clause(where)}
            """

            if group and x:
                query += "\nGROUP BY 1\nORDER BY 1"

            return {
                "query": query,
                "x_alias": aliases[0] if x else None,
                "y_aliases": aliases[1:] if x else aliases,
            }

        if isinstance(agg, str):
            agg = [agg] * len(y_list)

        select_parts = []
        y_aliases = []

        if x:
            select_parts.append(x)

        for col, a in zip(y_list, agg):
            agg_sql = self.AGG_MAP.get(a.lower(), "SUM")
            alias = f"{col}_{a.lower()}"
            select_parts.append(f"{agg_sql}({col}) AS {alias}")
            y_aliases.append(alias)

        query = f"""
            SELECT {", ".join(select_parts)}
            FROM data
            {self._where_clause(where)}
        """

        if group and x:
            query += f"\nGROUP BY {x}\nORDER BY {x}"

        return {
            "query": query,
            "x_alias": x,
            "y_aliases": y_aliases,
        }
    
    def build_kpi(
        self,
        metric: str,
        aggregation: str = "sum",
        where: Optional[str] = None,
    ):
        agg_sql = self.AGG_MAP.get(aggregation.lower(), "SUM")

        query = f"""
            SELECT {agg_sql}({metric}) AS value
            FROM data
            {self._where_clause(where)}
        """

        return {"query": query, "alias": "value"}
    
    def build_kpi_compare(
        self,
        metric: str,
        date_column: str,
        aggregation: str = "sum",
        period: str = "month",
        where: Optional[str] = None,
    ):
        agg_sql = self.AGG_MAP.get(aggregation.lower(), "SUM")

        query = f"""
            WITH base AS (
                SELECT
                    DATE_TRUNC('{period}', {date_column}) AS period,
                    {agg_sql}({metric}) AS value
                FROM data
                {self._where_clause(where)}
                GROUP BY 1
            )
            SELECT *
            FROM base
            ORDER BY period DESC
            LIMIT 2
        """

        return {"query": query}

