"""Query Builder — template-based SQL generation from NL intent.

Given a natural-language intent (e.g. "total revenue by well by month"),
this module resolves tables, measures, grain columns, and filters from
the BLCE artifact cache and skill store, then assembles a SQL template.

All logic is deterministic — no LLM calls.
"""
from __future__ import annotations

import re
from datetime import datetime
from typing import Any, Dict, List, Optional

from ..contracts import GeneratedQuery


class QueryBuilder:
    """Build SQL queries from NL intent + BLCE metadata."""

    # ------------------------------------------------------------------
    # Public API
    # ------------------------------------------------------------------

    def build_query(
        self,
        intent: str,
        *,
        artifact_cache: Optional[Dict[str, Dict[str, Any]]] = None,
        skill_store: Optional[Dict[str, Dict[str, Any]]] = None,
        vanna_client: Optional[Any] = None,
    ) -> GeneratedQuery:
        """Translate *intent* into a ``GeneratedQuery``.

        If a trained Vanna client is provided and returns a high-confidence
        result (>= 0.7), that result is used directly.  Otherwise the
        deterministic keyword-based resolution runs as before.

        Confidence scoring (deterministic path):
        - base 0.30
        - +0.20 if measures resolved
        - +0.20 if tables resolved
        - +0.15 if grain resolved
        - +0.15 if filters resolved
        """
        artifact_cache = artifact_cache or {}
        skill_store = skill_store or {}

        # Try Vanna RAG first if trained
        if vanna_client is not None:
            try:
                if vanna_client.is_trained():
                    vanna_result = vanna_client.ask(intent)
                    if vanna_result and vanna_result.get("confidence", 0) >= 0.7:
                        return GeneratedQuery(
                            intent=vanna_result.get("intent", intent),
                            sql=vanna_result.get("sql", ""),
                            tables_used=vanna_result.get("tables_used", []),
                            measures_used=vanna_result.get("measures_used", []),
                            grain_columns=vanna_result.get("grain_columns", []),
                            filters_applied=vanna_result.get("filters_applied", []),
                            confidence=vanna_result.get("confidence", 0.0),
                            warnings=vanna_result.get("warnings", []),
                        )
            except Exception:
                pass  # Fall through to deterministic path

        keywords = self._extract_keywords(intent)
        warnings: List[str] = []

        # Resolve components
        tables = self._resolve_tables(keywords, artifact_cache)
        measures = self._resolve_measures(keywords, artifact_cache, skill_store)
        grain_columns = self._resolve_grain(tables, artifact_cache)
        filters = self._resolve_filters(keywords, artifact_cache)

        # Confidence
        confidence = 0.30
        if measures:
            confidence += 0.20
        else:
            warnings.append("No measures matched — SQL may be incomplete")
        if tables:
            confidence += 0.20
        else:
            warnings.append("No tables matched — cannot generate FROM clause")
        if grain_columns:
            confidence += 0.15
        if filters:
            confidence += 0.15

        # Resolve joins via FK inference
        join_map = self._resolve_join_keys(tables, artifact_cache)

        # Assemble SQL
        sql = self._assemble_sql(
            measures=[m["expression"] for m in measures],
            tables=tables,
            grain_columns=grain_columns,
            filters=filters,
            join_map=join_map,
        )

        return GeneratedQuery(
            intent=intent,
            sql=sql,
            tables_used=tables,
            measures_used=[m.get("name") or m.get("alias", "") for m in measures],
            grain_columns=grain_columns,
            filters_applied=filters,
            confidence=round(min(confidence, 1.0), 2),
            warnings=warnings,
        )

    # ------------------------------------------------------------------
    # Resolution helpers
    # ------------------------------------------------------------------

    @staticmethod
    def _extract_keywords(text: str) -> List[str]:
        tokens = re.split(r"\W+", text.lower())
        _stop = {"a", "an", "the", "by", "for", "in", "of", "to", "and", "or", "is", "are", "total", "show", "get", "me"}
        return [t for t in tokens if t and t not in _stop]

    @staticmethod
    def _resolve_tables(
        keywords: List[str],
        artifact_cache: Dict[str, Dict[str, Any]],
    ) -> List[str]:
        """Find table names from artifact dependencies that match keywords.

        Uses both substring matching and token matching (split on ``_``).
        """
        tables: List[str] = []
        seen: set = set()
        for _aid, ad in artifact_cache.items():
            for dep in ad.get("objects", {}).get("dependencies", []):
                tname = dep.get("name", "")
                if tname and tname.lower() not in seen:
                    tname_lower = tname.lower()
                    tname_tokens = set(tname_lower.split("_"))
                    if (any(kw in tname_lower for kw in keywords)
                            or tname_tokens & set(keywords)):
                        tables.append(tname)
                        seen.add(tname_lower)
        return tables

    @staticmethod
    def _resolve_measures(
        keywords: List[str],
        artifact_cache: Dict[str, Dict[str, Any]],
        skill_store: Dict[str, Dict[str, Any]],
    ) -> List[Dict[str, Any]]:
        """Find measures whose name/alias matches keywords."""
        matches: List[Dict[str, Any]] = []
        seen: set = set()

        # From artifacts
        for _aid, ad in artifact_cache.items():
            for m in ad.get("objects", {}).get("measures", []):
                mname = (m.get("name") or m.get("alias") or "").lower()
                if mname and mname not in seen:
                    if any(kw in mname for kw in keywords):
                        matches.append(m)
                        seen.add(mname)

        # From skills
        for _sid, sd in skill_store.items():
            for mname in sd.get("measures_covered", []):
                ml = mname.lower()
                if ml not in seen and any(kw in ml for kw in keywords):
                    matches.append({"name": mname, "expression": mname, "alias": mname})
                    seen.add(ml)

        return matches

    @staticmethod
    def _resolve_grain(
        tables: List[str],
        artifact_cache: Dict[str, Dict[str, Any]],
    ) -> List[str]:
        """Extract grain columns from artifacts whose tables match."""
        grain: List[str] = []
        seen: set = set()
        table_lower = {t.lower() for t in tables}
        for _aid, ad in artifact_cache.items():
            deps = {d.get("name", "").lower() for d in ad.get("objects", {}).get("dependencies", [])}
            if deps & table_lower:
                for g in ad.get("objects", {}).get("grain_columns", []):
                    gc = g.get("column") or g.get("alias") or ""
                    if gc and gc.lower() not in seen:
                        grain.append(gc)
                        seen.add(gc.lower())
        return grain

    @staticmethod
    def _resolve_filters(
        keywords: List[str],
        artifact_cache: Dict[str, Dict[str, Any]],
    ) -> List[str]:
        """Find filter clauses that reference keyword-matched columns."""
        filters: List[str] = []
        seen: set = set()
        for _aid, ad in artifact_cache.items():
            for f in ad.get("objects", {}).get("filters", []):
                clause = f.get("source_clause") or f.get("column", "")
                if clause and clause not in seen:
                    col_lower = (f.get("column") or "").lower()
                    if any(kw in col_lower for kw in keywords):
                        filters.append(clause)
                        seen.add(clause)
        return filters

    # ------------------------------------------------------------------
    # Join resolution
    # ------------------------------------------------------------------

    @staticmethod
    def _resolve_join_keys(
        tables: List[str],
        artifact_cache: Dict[str, Dict[str, Any]],
    ) -> Dict[str, str]:
        """Resolve join ON clauses between tables using FK inference.

        Returns:
            Dict of target_table → "ON source.col = target.col"
        """
        if len(tables) < 2:
            return {}

        # Build table_columns map from artifact_cache (table profiles)
        table_columns: Dict[str, List[str]] = {}
        for t in tables:
            profile = artifact_cache.get(t, {})
            cols = profile.get("columns", [])
            if isinstance(cols, list) and cols:
                if isinstance(cols[0], dict):
                    table_columns[t] = [
                        c.get("column_name", c.get("name", "")) for c in cols
                    ]
                else:
                    table_columns[t] = list(cols)

        if not table_columns:
            return {}

        try:
            from src.data_modeling.column_fk_inferer import infer_relationships_by_name
            rels = infer_relationships_by_name(table_columns, min_score=2)
        except Exception:
            return {}

        # Build lookup: target_table → best ON clause (relative to base table)
        join_map: Dict[str, str] = {}
        base = tables[0]
        for rel in rels:
            src, tgt = rel["source_table"], rel["target_table"]
            if src == base and tgt in tables[1:] and tgt not in join_map:
                join_map[tgt] = (
                    f"ON {src}.{rel['source_column']} = {tgt}.{rel['target_column']}"
                )
            elif tgt == base and src in tables[1:] and src not in join_map:
                join_map[src] = (
                    f"ON {base}.{rel['target_column']} = {src}.{rel['source_column']}"
                )

        return join_map

    # ------------------------------------------------------------------
    # SQL assembly
    # ------------------------------------------------------------------

    @staticmethod
    def _assemble_sql(
        measures: List[str],
        tables: List[str],
        grain_columns: List[str],
        filters: List[str],
        join_map: Optional[Dict[str, str]] = None,
    ) -> str:
        """Assemble a SQL template from resolved components."""
        if not tables:
            return "-- No tables resolved; provide more context"

        select_parts: List[str] = []
        if grain_columns:
            select_parts.extend(grain_columns)
        if measures:
            select_parts.extend(measures)
        if not select_parts:
            select_parts.append("*")

        sql = f"SELECT {', '.join(select_parts)}\nFROM {tables[0]}"

        if len(tables) > 1:
            join_map = join_map or {}
            for t in tables[1:]:
                on_clause = join_map.get(
                    t, "ON 1=1  -- could not resolve join keys"
                )
                sql += f"\n  JOIN {t} {on_clause}"

        if filters:
            sql += f"\nWHERE {' AND '.join(filters)}"

        if grain_columns:
            sql += f"\nGROUP BY {', '.join(grain_columns)}"
            sql += f"\nORDER BY {', '.join(grain_columns)}"

        return sql
