"""Vanna RAG Text-to-SQL Client for DataBridge AI.

Thin wrapper around Vanna that manages training data (DDL, docs, Q&A),
SQL generation from natural language, and result caching.

Falls back gracefully when Vanna is not installed.
"""
from __future__ import annotations

import json
import logging
import os
from datetime import datetime, timezone
from pathlib import Path
from typing import Any, Dict, List, Optional

logger = logging.getLogger(__name__)

# ------------------------------------------------------------------
# Conditional Vanna import
# ------------------------------------------------------------------

try:
    # Vanna v2 moved to vanna.legacy.*
    try:
        from vanna.legacy.chromadb.chromadb_vector import ChromaDB_VectorStore
        from vanna.legacy.anthropic.anthropic_chat import Anthropic_Chat
    except ImportError:
        # Vanna v0.x / v1.x
        from vanna.chromadb import ChromaDB_VectorStore  # type: ignore[no-redef]
        from vanna.anthropic import Anthropic_Chat  # type: ignore[no-redef]

    VANNA_AVAILABLE = True
except ImportError:
    VANNA_AVAILABLE = False
    ChromaDB_VectorStore = object  # type: ignore[misc,assignment]
    Anthropic_Chat = object  # type: ignore[misc,assignment]

# ------------------------------------------------------------------
# Data directory
# ------------------------------------------------------------------

_DATA_DIR = Path(os.getenv("DATABRIDGE_DATA_DIR", "data")) / "vanna"
_STATE_FILE = _DATA_DIR / "training_state.json"


# ------------------------------------------------------------------
# Vanna client class
# ------------------------------------------------------------------

if VANNA_AVAILABLE:
    class _VannaBackend(ChromaDB_VectorStore, Anthropic_Chat):  # type: ignore[misc]
        """Vanna backend using ChromaDB for vectors + Anthropic Claude for LLM."""

        def __init__(self, config: Optional[Dict[str, Any]] = None):
            config = config or {}
            # ChromaDB persistence
            chroma_path = str(config.get("chroma_path", _DATA_DIR / "chromadb"))
            Path(chroma_path).mkdir(parents=True, exist_ok=True)
            ChromaDB_VectorStore.__init__(self, config={"path": chroma_path})
            # Anthropic LLM
            api_key = config.get("api_key") or os.getenv("ANTHROPIC_API_KEY", "")
            model = config.get("model", "claude-sonnet-4-20250514")
            Anthropic_Chat.__init__(self, config={
                "api_key": api_key,
                "model": model,
            })
else:
    _VannaBackend = None  # type: ignore[misc,assignment]


class DataBridgeVanna:
    """DataBridge wrapper around Vanna for RAG-powered Text-to-SQL.

    Usage::

        vanna = DataBridgeVanna()
        vanna.train_from_ddl("CREATE TABLE orders (id INT, amount DECIMAL)")
        result = vanna.ask("total orders by month")
        print(result)  # GeneratedQuery-compatible dict
    """

    def __init__(self, config: Optional[Dict[str, Any]] = None):
        self._config = config or {}
        self._vn: Any = None
        self._training_stats = {"ddl": 0, "docs": 0, "qa": 0}
        self._state_file = Path(self._config.get("state_file", _STATE_FILE))
        self._state_loaded = False
        self._state_persistence_enabled = True

        if VANNA_AVAILABLE and _VannaBackend is not None:
            try:
                self._vn = _VannaBackend(config=self._config)
                logger.info("Vanna RAG client initialized (ChromaDB + Anthropic)")
            except Exception as exc:
                logger.warning("Vanna initialization failed: %s", exc)
        else:
            logger.info("Vanna not installed — RAG Text-to-SQL unavailable")

        self._load_training_state()

    # ------------------------------------------------------------------
    # Training methods
    # ------------------------------------------------------------------

    def train_from_ddl(self, ddl: str) -> Dict[str, Any]:
        """Train on a CREATE TABLE (or other DDL) statement.

        Args:
            ddl: DDL string (e.g. ``CREATE TABLE orders (...)``).

        Returns:
            Status dict with training_id.
        """
        if not self._vn:
            return {"status": "skipped", "reason": "Vanna not available"}

        try:
            training_id = self._vn.train(ddl=ddl)
            self._training_stats["ddl"] += 1
            self._save_training_state()
            return {"status": "ok", "training_id": str(training_id), "type": "ddl"}
        except Exception as exc:
            logger.error("Vanna DDL training failed: %s", exc)
            return {"status": "error", "error": str(exc)}

    def train_from_documentation(self, doc: str) -> Dict[str, Any]:
        """Train on business context documentation.

        Args:
            doc: Free-text documentation about the data model.
        """
        if not self._vn:
            return {"status": "skipped", "reason": "Vanna not available"}

        try:
            training_id = self._vn.train(documentation=doc)
            self._training_stats["docs"] += 1
            self._save_training_state()
            return {"status": "ok", "training_id": str(training_id), "type": "documentation"}
        except Exception as exc:
            logger.error("Vanna doc training failed: %s", exc)
            return {"status": "error", "error": str(exc)}

    def train_from_query_history(self, sql: str, question: str) -> Dict[str, Any]:
        """Train on a past question + SQL pair.

        Args:
            sql: The SQL that answered *question*.
            question: The original natural-language question.
        """
        if not self._vn:
            return {"status": "skipped", "reason": "Vanna not available"}

        try:
            training_id = self._vn.train(sql=sql, question=question)
            self._training_stats["qa"] += 1
            self._save_training_state()
            return {"status": "ok", "training_id": str(training_id), "type": "qa"}
        except Exception as exc:
            logger.error("Vanna Q&A training failed: %s", exc)
            return {"status": "error", "error": str(exc)}

    def train_from_snowflake(
        self,
        account: str,
        database: str,
        schema: str,
        *,
        user: Optional[str] = None,
        password: Optional[str] = None,
        warehouse: Optional[str] = None,
        role: Optional[str] = None,
    ) -> Dict[str, Any]:
        """Auto-extract DDL from a Snowflake schema and train.

        Uses ``INFORMATION_SCHEMA`` to pull all table DDL for the given
        database/schema pair.

        Returns:
            Status dict with count of DDL statements trained.
        """
        if not self._vn:
            return {"status": "skipped", "reason": "Vanna not available"}

        try:
            import snowflake.connector

            conn = snowflake.connector.connect(
                account=account,
                user=user or os.getenv("SNOWFLAKE_USER", ""),
                password=password or os.getenv("SNOWFLAKE_PASSWORD", ""),
                warehouse=warehouse,
                role=role,
                database=database,
                schema=schema,
                client_store_temporary_credential=True,
                session_parameters={"CLIENT_SESSION_KEEP_ALIVE": True},
            )

            cur = conn.cursor()
            cur.execute(
                "SELECT TABLE_NAME FROM INFORMATION_SCHEMA.TABLES "
                "WHERE TABLE_SCHEMA = %s ORDER BY TABLE_NAME",
                (schema.upper(),),
            )
            tables = [row[0] for row in cur.fetchall()]

            trained = 0
            for tbl in tables:
                try:
                    cur.execute(f'SELECT GET_DDL(\'TABLE\', \'"{database}"."{schema}"."{tbl}"\')')
                    ddl = cur.fetchone()[0]
                    result = self.train_from_ddl(ddl)
                    if result.get("status") == "ok":
                        trained += 1
                except Exception as tbl_exc:
                    logger.warning("Skipping DDL for %s: %s", tbl, tbl_exc)

            cur.close()
            conn.close()

            return {
                "status": "ok",
                "tables_found": len(tables),
                "tables_trained": trained,
                "schema": f"{database}.{schema}",
            }

        except ImportError:
            return {"status": "error", "error": "snowflake-connector-python not installed"}
        except Exception as exc:
            return {"status": "error", "error": str(exc)}

    # ------------------------------------------------------------------
    # SQL generation
    # ------------------------------------------------------------------

    def ask(self, question: str) -> Optional[Dict[str, Any]]:
        """Generate SQL from a natural-language question.

        Returns a dict compatible with BLCE's ``GeneratedQuery`` model,
        or ``None`` if generation fails.

        Args:
            question: Natural-language question (e.g. "total orders by month").
        """
        if not self._vn:
            return None

        try:
            sql = self._generate_sql_with_fallbacks(question)
            if not sql or not sql.strip():
                return None

            # Extract table names from SQL (basic parser)
            tables_used = self._extract_tables_from_sql(sql)

            # Estimate confidence based on training data richness
            confidence = self._estimate_confidence()

            return {
                "intent": question,
                "sql": sql.strip(),
                "tables_used": tables_used,
                "measures_used": [],
                "grain_columns": [],
                "filters_applied": [],
                "confidence": confidence,
                "warnings": ["Generated by Vanna RAG — review before execution"],
                "generated_at": datetime.now(timezone.utc).isoformat(),
                "source": "vanna_rag",
            }
        except Exception as exc:
            logger.error("Vanna SQL generation failed: %s", exc)
            return None

    def _generate_sql_with_fallbacks(self, question: str) -> str:
        """Try multiple backend methods to maximize compatibility across Vanna versions."""
        if not self._vn:
            return ""

        # Preferred path for most versions.
        sql = self._call_generate_sql(question)
        if sql:
            return sql

        # Fallback to ask() if exposed by backend wrappers.
        sql = self._call_ask(question)
        if sql:
            return sql

        # Fallback to retrieval cache for exact/similar trained questions.
        sql = self._call_similar_question_sql(question)
        if sql:
            return sql
        return ""

    def _call_generate_sql(self, question: str) -> str:
        fn = getattr(self._vn, "generate_sql", None)
        if not callable(fn):
            return ""
        try:
            out = fn(question=question)
        except TypeError:
            try:
                out = fn(question)
            except Exception:
                return ""
        except Exception:
            return ""
        return self._extract_sql_text(out)

    def _call_ask(self, question: str) -> str:
        fn = getattr(self._vn, "ask", None)
        if not callable(fn):
            return ""
        try:
            out = fn(question=question)
        except TypeError:
            try:
                out = fn(question)
            except Exception:
                return ""
        except Exception:
            return ""
        return self._extract_sql_text(out)

    def _call_similar_question_sql(self, question: str) -> str:
        fn = getattr(self._vn, "get_similar_question_sql", None)
        if not callable(fn):
            return ""
        try:
            out = fn(question=question)
        except TypeError:
            try:
                out = fn(question)
            except Exception:
                return ""
        except Exception:
            return ""

        # Common shapes:
        # - "SELECT ..."
        # - {"sql": "..."}
        # - [{"sql":"..."}, ...]
        # - [("question", "sql"), ...]
        if isinstance(out, list) and out:
            first = out[0]
            if isinstance(first, tuple) and len(first) >= 2 and isinstance(first[1], str):
                return first[1].strip()
            if isinstance(first, dict):
                sql = first.get("sql")
                if isinstance(sql, str):
                    return sql.strip()
            if isinstance(first, str):
                return first.strip()
        return self._extract_sql_text(out)

    @staticmethod
    def _extract_sql_text(value: Any) -> str:
        if isinstance(value, str):
            return value.strip()
        if isinstance(value, dict):
            for key in ("sql", "query", "text"):
                v = value.get(key)
                if isinstance(v, str) and v.strip():
                    return v.strip()
        return ""

    # ------------------------------------------------------------------
    # Status / readiness
    # ------------------------------------------------------------------

    def get_training_status(self) -> Dict[str, Any]:
        """Return training data statistics."""
        total = sum(self._training_stats.values())
        state_file = getattr(self, "_state_file", _STATE_FILE)
        state_enabled = bool(getattr(self, "_state_persistence_enabled", False))
        return {
            "available": self._vn is not None,
            "ddl_count": self._training_stats["ddl"],
            "doc_count": self._training_stats["docs"],
            "qa_count": self._training_stats["qa"],
            "total_training_items": total,
            "is_trained": total >= 1,
            "data_dir": str(_DATA_DIR),
            "state_file": str(state_file),
            "state_persisted": state_enabled and state_file.exists(),
            "state_loaded": bool(getattr(self, "_state_loaded", False)),
        }

    def is_trained(self) -> bool:
        """Whether enough training data exists for useful generation."""
        if not self._vn:
            return False
        return sum(self._training_stats.values()) >= 1

    def _load_training_state(self) -> None:
        """Load training counters from disk if persistence is enabled."""
        if not getattr(self, "_state_persistence_enabled", False):
            return

        state_file = getattr(self, "_state_file", _STATE_FILE)
        self._state_loaded = False

        if not state_file.exists():
            return

        try:
            data = json.loads(state_file.read_text(encoding="utf-8"))
            self._training_stats = {
                "ddl": int(data.get("ddl", 0)),
                "docs": int(data.get("docs", 0)),
                "qa": int(data.get("qa", 0)),
            }
            self._state_loaded = True
        except Exception as exc:
            logger.warning("Vanna training state load failed (%s); resetting counters", exc)
            self._training_stats = {"ddl": 0, "docs": 0, "qa": 0}

    def _save_training_state(self) -> bool:
        """Persist training counters to disk if persistence is enabled."""
        if not getattr(self, "_state_persistence_enabled", False):
            return False

        state_file = getattr(self, "_state_file", _STATE_FILE)
        try:
            state_file.parent.mkdir(parents=True, exist_ok=True)
            payload = {
                "ddl": int(self._training_stats.get("ddl", 0)),
                "docs": int(self._training_stats.get("docs", 0)),
                "qa": int(self._training_stats.get("qa", 0)),
                "updated_at": datetime.now(timezone.utc).isoformat(),
            }
            state_file.write_text(json.dumps(payload, indent=2), encoding="utf-8")
            return True
        except Exception as exc:
            logger.warning("Vanna training state save failed: %s", exc)
            return False

    # ------------------------------------------------------------------
    # Internal helpers
    # ------------------------------------------------------------------

    @staticmethod
    def _extract_tables_from_sql(sql: str) -> List[str]:
        """Basic extraction of table names from SQL."""
        import re

        tables: List[str] = []
        seen: set = set()
        # Match FROM/JOIN followed by table name
        for match in re.finditer(
            r"(?:FROM|JOIN)\s+\"?(\w+)\"?\.?\"?(\w+)?\"?\.?\"?(\w+)?\"?",
            sql,
            re.IGNORECASE,
        ):
            parts = [g for g in match.groups() if g]
            table_name = parts[-1] if parts else ""
            if table_name and table_name.upper() not in seen:
                tables.append(table_name)
                seen.add(table_name.upper())
        return tables

    def _estimate_confidence(self) -> float:
        """Estimate confidence based on training data volume."""
        stats = self._training_stats
        score = 0.4  # base
        if stats["ddl"] >= 5:
            score += 0.2
        elif stats["ddl"] >= 1:
            score += 0.1
        if stats["docs"] >= 1:
            score += 0.1
        if stats["qa"] >= 5:
            score += 0.2
        elif stats["qa"] >= 1:
            score += 0.1
        return min(round(score, 2), 1.0)


# ------------------------------------------------------------------
# Module-level singleton
# ------------------------------------------------------------------

_vanna_client: Optional[DataBridgeVanna] = None


def get_vanna_client() -> DataBridgeVanna:
    """Get or create the singleton Vanna client."""
    global _vanna_client
    if _vanna_client is None:
        _vanna_client = DataBridgeVanna()
    return _vanna_client
