"""AI-assisted data dictionary and classification."""
from __future__ import annotations

import json
import logging
import os
import re
from typing import Dict, List, Optional

logger = logging.getLogger(__name__)

from src.data_catalog.types import DataClassification


def _heuristic_classification(column: str) -> str:
    name = column.lower()
    if any(k in name for k in ["email", "ssn", "phone", "dob", "birth", "tax"]):
        return "pii"
    if name.endswith("_id") or name == "id":
        return "internal"
    if any(k in name for k in ["amount", "total", "qty", "count", "balance"]):
        return "internal"
    return "internal"


def _heuristic_description(column: str) -> str:
    name = column.replace("_", " ").title()
    return f"{name} field."


class SnowflakeQueryClient:
    """Minimal Snowflake query client for Cortex calls.

    Accepts an optional pre-existing connection to avoid SSO prompts.
    """

    def __init__(self, database: str, schema: str,
                 connection_name: Optional[str] = None,
                 sf_connection=None):
        self._database = database
        self._schema = schema
        self._connection_name = connection_name
        self._sf_connection = sf_connection

        if sf_connection is None:
            from .snowflake_utils import ensure_snowflake_env_from_cli
            if not os.getenv("SNOWFLAKE_ACCOUNT"):
                ensure_snowflake_env_from_cli(connection_name=connection_name)

    def execute_query(self, sql: str, limit: int = 1):
        if self._sf_connection is not None:
            conn = self._sf_connection
            cur = conn.cursor()
            cur.execute(sql)
            cols = [c[0] for c in cur.description] if cur.description else []
            rows = cur.fetchmany(limit)
            cur.close()
            return [dict(zip(cols, r)) for r in rows]

        from .snowflake_session_manager import get_snowflake_session_manager
        manager = get_snowflake_session_manager()
        return manager.execute_query(
            sql,
            connection_name=self._connection_name,
            database=self._database,
            schema=self._schema,
            limit=limit,
            use_env=True,
        )


def _cortex_classify(
    columns: List[Dict[str, str]],
    sample_values: Dict[str, List],
    database: str,
    schema: str,
    connection_name: Optional[str] = None,
    sf_connection=None,
) -> Dict[str, str]:
    from src.datashield.ai_classifier import get_ai_classifier

    sf_client = SnowflakeQueryClient(database, schema, connection_name=connection_name,
                                     sf_connection=sf_connection)
    classifier = get_ai_classifier("snowflake_cortex", sf_client=sf_client)
    existing = {c["name"]: "safe" for c in columns}
    return classifier.classify(columns, existing, sample_values=sample_values)


def _anthropic_classify(
    columns: List[Dict[str, str]],
    sample_values: Dict[str, List],
) -> Dict[str, str]:
    from src.datashield.ai_classifier import get_ai_classifier

    api_key = os.getenv("ANTHROPIC_API_KEY", "")
    classifier = get_ai_classifier("anthropic", api_key=api_key)
    existing = {c["name"]: "safe" for c in columns}
    return classifier.classify(columns, existing, sample_values=sample_values)


def classify_columns(
    columns: List[Dict[str, str]],
    sample_values: Dict[str, List],
    database: Optional[str] = None,
    schema: Optional[str] = None,
    connection_name: Optional[str] = None,
    sf_connection=None,
) -> Dict[str, str]:
    # Try Cortex first
    if database and schema:
        try:
            result = _cortex_classify(columns, sample_values, database, schema,
                                      connection_name=connection_name,
                                      sf_connection=sf_connection)
            logger.debug("Cortex classification succeeded (%d columns)", len(result))
            return result
        except Exception as exc:
            logger.warning("Cortex classify failed, trying Anthropic: %s", exc)
    # Then Anthropic
    if os.getenv("ANTHROPIC_API_KEY"):
        try:
            result = _anthropic_classify(columns, sample_values)
            logger.debug("Anthropic classification succeeded")
            return result
        except Exception as exc:
            logger.warning("Anthropic classify failed, using heuristics: %s", exc)
    # Fallback heuristics ("codex" backup)
    logger.debug("Using heuristic classification (%d columns)", len(columns))
    return {c["name"]: _heuristic_classification(c["name"]) for c in columns}


def generate_dictionary(
    table: str,
    columns: List[Dict[str, str]],
    sample_values: Dict[str, List],
    classifications: Dict[str, str],
    database: Optional[str] = None,
    schema: Optional[str] = None,
    connection_name: Optional[str] = None,
    sf_connection=None,
) -> Dict[str, Dict[str, str]]:
    dictionary: Dict[str, Dict[str, str]] = {}

    for col in columns:
        name = col["name"]
        classification = classifications.get(name, "internal")
        desc = _heuristic_description(name)
        dictionary[name] = {
            "classification": classification,
            "description": desc,
        }

    # Optional Cortex enrichment (single prompt)
    if database and schema:
        try:
            from src.cortex_agent.cortex_client import CortexClient

            def query_func(connection_id: str, query: str):
                client = SnowflakeQueryClient(database, schema, connection_name=connection_name,
                                              sf_connection=sf_connection)
                return client.execute_query(query)

            cortex = CortexClient(connection_id="local", query_func=query_func)
            prompt = {
                "table": table,
                "columns": [
                    {
                        "name": c["name"],
                        "sample_values": sample_values.get(c["name"], [])[:3],
                        "current_description": dictionary[c["name"]]["description"],
                    }
                    for c in columns
                ],
                "task": (
                    "Improve descriptions for a data dictionary. "
                    "Respond with ONLY a valid JSON object, no markdown, no explanation. "
                    'Format: {"COLUMN_NAME": "improved description", ...}'
                ),
            }
            result = cortex.complete(prompt=json.dumps(prompt, default=str), max_tokens=1200)
            if result.success and isinstance(result.result, str):
                raw = result.result.strip()
                # Strip markdown code fences if present
                md_match = re.search(r"```(?:json)?\s*\n?(.*?)```", raw, re.DOTALL)
                if md_match:
                    raw = md_match.group(1).strip()
                try:
                    improved = json.loads(raw)
                    for col_name, desc in improved.items():
                        if col_name in dictionary:
                            dictionary[col_name]["description"] = str(desc)
                except json.JSONDecodeError:
                    logger.debug("Cortex returned non-JSON for table %s enrichment (non-critical)", table)
        except Exception as exc:
            logger.warning("Cortex enrichment failed for %s: %s", table, exc)

    return dictionary


def to_catalog_classification(value: str) -> DataClassification:
    normalized = (value or "").lower()
    if normalized in ("sensitive_pii", "pii"):
        return DataClassification.PII
    if normalized in ("phi",):
        return DataClassification.PHI
    if normalized in ("pci",):
        return DataClassification.PCI
    try:
        return DataClassification(normalized)
    except Exception:
        return DataClassification.INTERNAL
