backend/agents/sql_agent.py

from __future__ import annotations
from dataclasses import dataclass, field
from typing import Any

from backend.agents.planner import PlannerOutput
from backend.services.clickhouse_service import ClickHouseService, ColumnInfo
from backend.services.llm_service import LLMService
import re

@dataclass
class SQLAgentOutput:
    sql: str
    rationale: str
    selected_tables: list[str] = field(default_factory=list)
    selected_columns: list[dict[str, str]] = field(default_factory=list)
    result: dict[str, Any] = field(default_factory=dict)
    warnings: list[str] = field(default_factory=list)
    error: str | None = None

    def to_dict(self) -> dict[str, Any]:
        return {
            "sql": self.sql,
            "rationale": self.rationale,
            "selected_tables": self.selected_tables,
            "selected_columns": self.selected_columns,
            "result": self.result,
            "warnings": self.warnings,
            "error": self.error,
        }


class SQLAgent:
    def __init__(self, llm: LLMService, clickhouse: ClickHouseService) -> None:
        self.llm = llm
        self.clickhouse = clickhouse

    def run(self, question: str, plan: PlannerOutput) -> SQLAgentOutput:
        try:
            schema_catalog = self.clickhouse.get_schema_catalog()
        except Exception as exc:
            return SQLAgentOutput(
                sql="",
                rationale="Schema lookup failed before SQL generation.",
                warnings=["Unable to inspect ClickHouse schema."],
                error=str(exc),
            )

        query_pattern = self._infer_query_pattern(question)
        selected = self._select_relevant_columns(
            question=question,
            schema_catalog=schema_catalog,
            query_pattern=query_pattern,
        )
        selected_tables = sorted({f"{column.database}.{column.table}" for column in selected})
        fallback_sql = self._build_fallback_sql(
            question=question,
            plan=plan,
            selected=selected,
            query_pattern=query_pattern,
        )

        payload = self.llm.invoke_json(
            system_prompt=(
                "You are a senior ClickHouse SQL agent. "
                "Return strict JSON with keys sql, rationale, selected_tables, selected_columns, warnings. "
                "Use only read-only SELECT queries. Never reference columns outside the provided schema. "
                "Choose the query pattern that matches the business intent, not a generic time trend."
                "NEVER use raw aggregate functions (COUNT, AVG, SUM) inside HAVING. "
                "ALWAYS use aliases defined in SELECT inside HAVING. "
                "ALWAYS define aggregates explicitly in SELECT. "
                "Never reuse alias names without defining them. "
                "Example: AVG(star_rating) AS avg_rating must appear in SELECT before using avg_rating. "
                "Example: COUNT(*) AS total_reviews → use total_reviews in HAVING. "
                "Example: AVG(star_rating) AS avg_rating → use avg_rating in HAVING. "

            ),
            user_prompt=self._build_generation_prompt(
                question=question,
                plan=plan,
                selected=selected,
                fallback_sql=fallback_sql,
                query_pattern=query_pattern,
            ),
            fallback={
                "sql": fallback_sql,
                "rationale": "Fallback SQL generated from semantic pattern detection over the grounded schema.",
                "selected_tables": selected_tables,
                "selected_columns": [column.to_dict() for column in selected],
                "warnings": [],
            },
        )
    
        sql = (payload.get("sql") or fallback_sql).strip().rstrip(";")
        if not sql.lower().startswith("select"):
            sql = fallback_sql
        sql = self._enforce_readable_product_label(sql=sql, selected=selected, query_pattern=query_pattern)
        sql = self._fix_having_clause_aliases(sql)

        try:
            result = self.clickhouse.query(sql)
            error = None
        except Exception as exc:
            result = {"columns": [], "rows": [], "row_count": 0}
            error = str(exc)

        return SQLAgentOutput(
            sql=sql,
            rationale=payload.get("rationale") or "SQL generated for the detected business question.",
            selected_tables=payload.get("selected_tables") or selected_tables,
            selected_columns=payload.get("selected_columns") or [column.to_dict() for column in selected],
            result=result,
            warnings=payload.get("warnings") or [],
            error=error,
        )

    def _enforce_readable_product_label(
        self,
        sql: str,
        selected: list[ColumnInfo],
        query_pattern: dict[str, Any],
    ) -> str:
        if query_pattern["entity"] != "product":
            return sql

        sql_lower = sql.lower()
        if "product_title" in sql_lower:
            return sql

        product_label = next(
            (
                column.name
                for column in selected
                if column.name.lower() in {"product_title", "product_name"}
            ),
            None,
        )
        if not product_label:
            return sql

        product_key_aliases = ["product_parent", "product_id", " as product_id", " as product_key"]
        if not any(alias in sql_lower for alias in product_key_aliases):
            return sql

        if " from " not in sql_lower:
            return sql

        select_prefix, remainder = sql.split("FROM", 1) if "FROM" in sql else sql.split("from", 1)
        if "select" not in select_prefix.lower():
            return sql

        return f"{select_prefix.rstrip()}, any({product_label}) AS product_title FROM{remainder}"

    def _infer_query_pattern(self, question: str) -> dict[str, Any]:
        tokens = {token.strip(" ,.?").lower() for token in question.split()}
        return {
            "needs_time_grain": bool(
                tokens & {"trend", "over", "monthly", "month", "daily", "day", "weekly", "quarterly", "quarter", "yearly", "year"}
            ),
            "needs_popularity": bool(tokens & {"popular", "popularity", "top", "most"}),
            "needs_low_rating": bool(tokens & {"poorly", "badly", "low", "worst", "negative"}),
            "needs_high_rating": bool(tokens & {"highest", "best", "top-rated"}),
            "entity": "product" if tokens & {"product", "products"} else "category" if "category" in tokens else "generic",
        }

    def _select_relevant_columns(
        self,
        question: str,
        schema_catalog: list[ColumnInfo],
        query_pattern: dict[str, Any],
    ) -> list[ColumnInfo]:
        tokens = {token.strip(" ,.?").lower() for token in question.split()}
        scored: list[tuple[int, ColumnInfo]] = []
        for column in schema_catalog:
            score = 0
            table_name = column.table.lower()
            column_name = column.name.lower()
            column_type = column.type.lower()

            for token in tokens:
                if token and token in column_name:
                    score += 5
                if token and token in table_name:
                    score += 3

            if query_pattern["entity"] == "product" and column_name in {"product_parent", "product_title", "product_id"}:
                score += 10
            if query_pattern["entity"] == "category" and "category" in column_name:
                score += 8
            if query_pattern["needs_popularity"] and any(hint in column_name for hint in {"review", "count", "votes"}):
                score += 8
            if query_pattern["needs_low_rating"] and any(hint in column_name for hint in {"rating", "score", "star"}):
                score += 8
            if query_pattern["needs_time_grain"] and any(hint in column_name for hint in {"date", "time", "month", "year"}):
                score += 6

            if any(metric in column_name for metric in {"revenue", "sales", "amount", "price", "profit", "count", "rating", "votes", "review"}):
                score += 2
            if "int" in column_type or "float" in column_type or "decimal" in column_type:
                score += 1
            scored.append((score, column))

        scored.sort(key=lambda item: (item[0], item[1].table, item[1].name), reverse=True)
        top = [column for score, column in scored if score > 0][:10]
        if top:
            return top
        return schema_catalog[:10]

    def _build_generation_prompt(
        self,
        question: str,
        plan: PlannerOutput,
        selected: list[ColumnInfo],
        fallback_sql: str,
        query_pattern: dict[str, Any],
    ) -> str:
        schema_lines = [
            f"- table={column.database}.{column.table}, column={column.name}, type={column.type}"
            for column in selected
        ]
        return (
            f"Question: {question}\n"
            f"Intent: {plan.intent}\n"
            f"Time range: {plan.time_range}\n"
            f"Derived query pattern: {query_pattern}\n"
            "Relevant schema:\n"
            f"{chr(10).join(schema_lines)}\n\n"
            "Requirements:\n"
            "- Prefer one table unless a join is clearly necessary.\n"
            "- Use aliases that are easy to read in a UI.\n"
            "- Do not group by time unless the user clearly asked for a trend over time.\n"
            "- For popularity questions, prefer COUNT(*) or review counts over unrelated sums unless votes were explicitly requested.\n"
            "- For poorly rated questions, use AVG on the rating column and HAVING filters when the business ask implies thresholds.\n"
            "- For product questions, group by a stable product identifier and include a readable product label when available.\n"
            "- If both popularity and poor rating are requested, produce a ranking/filter query by product, not a time series.\n"
            "- Use LIMIT 20 for ranked entity lists and LIMIT 200 for trends.\n"
            f"- If unsure, use this safe fallback SQL:\n{fallback_sql}\n"
        )
    def _ensure_aggregates_exist(self, sql: str, selected: list[ColumnInfo]) -> str:
        sql_lower = sql.lower()

        if "group by" not in sql_lower or "from" not in sql_lower:
            return sql

        select_part, rest = sql.split("FROM", 1)

        # Detect aliases used
        needs_avg = "avg_" in sql_lower
        needs_count = "total_" in sql_lower or "count_" in sql_lower

        # Find candidate numeric columns dynamically
        numeric_cols = [
            col.name for col in selected
            if any(t in col.type.lower() for t in ["int", "float", "decimal"])
        ]

        # Heuristic mapping
        rating_col = next((c for c in numeric_cols if "rating" in c.lower()), None)
        count_col = next((c for c in numeric_cols if "id" in c.lower()), None)

        # Inject AVG dynamically
        if needs_avg and "avg(" not in select_part.lower() and rating_col:
            select_part += f", AVG({rating_col}) AS avg_rating"

        # Inject COUNT dynamically
        if needs_count and "count(" not in select_part.lower():
            if count_col:
                select_part += f", COUNT({count_col}) AS total_count"
            else:
                select_part += ", COUNT(*) AS total_count"

        return f"{select_part} FROM {rest}"
    
    
    def _build_fallback_sql(
        self,
        question: str,
        plan: PlannerOutput,
        selected: list[ColumnInfo],
        query_pattern: dict[str, Any],
    ) -> str:
        if not selected:
            return "SELECT 1 AS value LIMIT 1"

        tokens = {token.strip(" ,.?").lower() for token in question.split()}
        table = f"{selected[0].database}.{selected[0].table}"

        time_column = next(
            (column.name for column in selected if any(hint in column.name.lower() for hint in {"date", "time"})),
            None,
        )
        dimension_column = next(
            (column.name for column in selected if "string" in column.type.lower()),
            None,
        )
        rating_column = next(
            (column.name for column in selected if any(hint in column.name.lower() for hint in {"rating", "score", "star"})),
            None,
        )
        votes_column = next(
            (column.name for column in selected if any(hint in column.name.lower() for hint in {"votes", "helpful"})),
            None,
        )
        product_key = next(
            (column.name for column in selected if column.name.lower() in {"product_parent", "product_id"}),
            None,
        )
        product_label = next(
            (column.name for column in selected if column.name.lower() in {"product_title", "product_name"}),
            None,
        )
        category_column = next(
            (column.name for column in selected if "category" in column.name.lower()),
            None,
        )

        if query_pattern["entity"] == "product" and product_key:
            select_parts = [product_key]
            if product_label:
                select_parts.append(f"any({product_label}) AS product_title")
            select_parts.append("COUNT(*) AS total_reviews")
            if votes_column:
                select_parts.append(f"SUM({votes_column}) AS total_votes")
            if rating_column:
                select_parts.append(f"AVG({rating_column}) AS avg_rating")

            sql_parts = [
                "SELECT " + ", ".join(select_parts),
                f"FROM {table}",
                f"GROUP BY {product_key}",
            ]

            having_parts = []
            if query_pattern["needs_popularity"]:
                having_parts.append("total_reviews > 100")
            if query_pattern["needs_low_rating"] and rating_column:
                having_parts.append("avg_rating < 3")
            if query_pattern["needs_high_rating"] and rating_column:
                having_parts.append("avg_rating >= 4")
            if having_parts:
                sql_parts.append("HAVING " + " AND ".join(having_parts))

            order_parts = []
            if query_pattern["needs_popularity"]:
                order_parts.append("total_reviews DESC")
            if votes_column:
                order_parts.append("total_votes DESC")
            if query_pattern["needs_low_rating"] and rating_column and not query_pattern["needs_popularity"]:
                order_parts.append("avg_rating ASC")
            if query_pattern["needs_high_rating"] and rating_column and not query_pattern["needs_popularity"]:
                order_parts.append("avg_rating DESC")
            if order_parts:
                sql_parts.append("ORDER BY " + ", ".join(order_parts))
            sql_parts.append("LIMIT 20")
            return " ".join(sql_parts)

        if query_pattern["needs_time_grain"] and time_column:
            period_expr = (
                f"toStartOfMonth({time_column})" if {"month", "monthly"} & tokens
                else f"toStartOfYear({time_column})" if {"year", "yearly", "annual"} & tokens
                else f"toDate({time_column})"
            )
            grouping_dimension = category_column or dimension_column
            metric_expr = "COUNT(*) AS total_reviews"
            if rating_column and {"average", "avg", "mean"} & tokens:
                metric_expr = f"AVG({rating_column}) AS avg_rating"
            elif votes_column and "votes" in tokens:
                metric_expr = f"SUM({votes_column}) AS total_votes"

            dimension_sql = f", {grouping_dimension}" if grouping_dimension else ""
            group_by = f"GROUP BY period{', ' + grouping_dimension if grouping_dimension else ''}"
            order_by = f"ORDER BY period{', ' + grouping_dimension if grouping_dimension else ''}"
            return (
                f"SELECT {period_expr} AS period{dimension_sql}, {metric_expr} "
                f"FROM {table} "
                f"{group_by} "
                f"{order_by} "
            
            )

        if category_column and rating_column and query_pattern["needs_low_rating"]:
            return (
                f"SELECT {category_column} AS category, COUNT(*) AS total_reviews, AVG({rating_column}) AS avg_rating "
                f"FROM {table} "
                "GROUP BY category "
                "HAVING total_reviews > 20 AND avg_rating < 3 "
                "ORDER BY total_reviews DESC "
                "LIMIT 20"
            )

        if dimension_column:
            metric_expr = "COUNT(*) AS total_count"
            if votes_column and "votes" in tokens:
                metric_expr = f"SUM({votes_column}) AS total_votes"
            elif rating_column and {"average", "avg", "mean"} & tokens:
                metric_expr = f"AVG({rating_column}) AS avg_rating"
            return (
                f"SELECT {dimension_column} AS category, {metric_expr} "
                f"FROM {table} "
                "GROUP BY category "
                "ORDER BY 2 DESC "
                "LIMIT 20"
            )

        return f"SELECT * FROM {table} LIMIT 50"
    
    def _extract_aliases(self, sql: str) -> dict[str, str]:
        """
        Extract mapping: aggregate_expression -> alias
        Example:
        AVG(star_rating) AS avg_rating → {"avg(star_rating)": "avg_rating"}
        """
        import re

        select_match = re.search(r"SELECT(.*?)FROM", sql, re.IGNORECASE | re.DOTALL)
        if not select_match:
            return {}

        select_part = select_match.group(1)

        alias_map = {}

        # Match patterns like: AVG(col) AS alias
        matches = re.findall(
            r"(AVG|COUNT|SUM|MIN|MAX)\((.*?)\)\s+AS\s+(\w+)",
            select_part,
            re.IGNORECASE,
        )

        for func, col, alias in matches:
            key = f"{func.lower()}({col.strip()})"
            alias_map[key] = alias

        return alias_map
    

    def _fix_having_clause_aliases(self, sql: str) -> str:
        

        if "having" not in sql.lower():
            return sql

        alias_map = self._extract_aliases(sql)

        if not alias_map:
            return sql  # nothing to fix

        parts = re.split(r"\bHAVING\b", sql, flags=re.IGNORECASE)
        if len(parts) < 2:
            return sql

        before = parts[0]
        having = parts[1]

        # Replace aggregate expressions with correct aliases
        for agg_expr, alias in alias_map.items():
            pattern = re.escape(agg_expr)
            having = re.sub(pattern, alias, having, flags=re.IGNORECASE)

        return before + "HAVING " + having
    
    
    
    def _validate_sql(self, sql: str) -> str:
        

        sql_lower = sql.lower()

        # Extract SELECT aliases
        select_part = sql_lower.split("from")[0]

        aliases = set(re.findall(r"as\s+(\w+)", select_part))

        # Extract HAVING usage
        if "having" in sql_lower:
            having_part = sql_lower.split("having")[1]

            used_aliases = set(re.findall(r"\b[a-z_]+\b", having_part))

            # Find aliases used but not defined
            undefined = used_aliases - aliases

            # Ignore SQL keywords
            keywords = {"and", "or", "not", "in", "between", "like"}
            undefined = {u for u in undefined if u not in keywords}

            if undefined:
                raise ValueError(f"Undefined aliases in HAVING: {undefined}")

        return sql