"""DataShield AI Classifier — LLM-enhanced column classification.

Supports two providers:
  1. Snowflake Cortex — runs inference inside Snowflake (data never leaves)
  2. Anthropic Claude  — external API call via anthropic SDK

Both return a dict of column overrides: {column_name: new_classification}
for columns where the AI disagrees with the current (rule-based) classification.
"""

import json
import logging
from typing import Any, Dict, List, Optional

from .types import ColumnClassification

logger = logging.getLogger(__name__)

# =========================================================================
# Shared prompt
# =========================================================================

_SYSTEM_PROMPT = """You are a data classification expert. Your job is to classify database columns
into exactly one of these categories:

- measure: numeric quantities (amounts, totals, counts, percentages)
- fact_dimension: foreign keys linking to dimension tables
- descriptive: human-readable labels (names, descriptions, titles)
- geographic: locations (cities, states, countries, addresses, zip codes)
- temporal: dates and timestamps
- identifier: unique row identifiers (IDs, primary keys, GUIDs)
- code: short coded values (status codes, type codes, flags)
- sensitive_pii: personally identifiable information (SSN, email, phone, DOB, salary)
- safe: non-sensitive, low-risk columns (booleans, system flags)

Rules:
- Only return overrides for columns where your classification DIFFERS from the current one.
- If all current classifications are correct, return an empty JSON object {}.
- Return ONLY valid JSON — no markdown, no explanation, no extra text.
- Be conservative: only override when you are confident the current classification is wrong.
- Pay special attention to columns that look like PII but may be classified as something else."""


def _build_user_prompt(
    columns: List[Dict[str, Any]],
    existing_rules: Dict[str, str],
    sample_values: Optional[Dict[str, List]] = None,
) -> str:
    """Build the user prompt with column info and samples."""
    lines = ["Analyze these database columns and their current classifications.\n"]
    lines.append("| Column | Data Type | Current Classification | Sample Values |")
    lines.append("|--------|-----------|----------------------|---------------|")

    for col in columns:
        name = col.get("name", col.get("column_name", ""))
        dtype = col.get("data_type", "VARCHAR")
        current = existing_rules.get(name, "safe")
        samples = ""
        if sample_values and name in sample_values:
            vals = [str(v) for v in sample_values[name][:5] if v is not None]
            samples = ", ".join(vals[:5])
        lines.append(f"| {name} | {dtype} | {current} | {samples} |")

    lines.append("\nReturn a JSON object with ONLY the columns that should be reclassified.")
    lines.append('Format: {"COLUMN_NAME": "new_classification", ...}')
    lines.append("If all classifications are correct, return: {}")
    return "\n".join(lines)


def _parse_response(raw_text: str) -> Dict[str, str]:
    """Extract a JSON dict from the LLM response, validating values."""
    # Try direct parse first
    text = raw_text.strip()
    # Strip markdown code fences if present
    if text.startswith("```"):
        lines = text.split("\n")
        lines = [l for l in lines if not l.strip().startswith("```")]
        text = "\n".join(lines).strip()

    try:
        result = json.loads(text)
    except json.JSONDecodeError:
        # Try to find JSON object in the text
        start = text.find("{")
        end = text.rfind("}") + 1
        if start >= 0 and end > start:
            try:
                result = json.loads(text[start:end])
            except json.JSONDecodeError:
                logger.warning("Failed to parse AI classifier response: %s", text[:200])
                return {}
        else:
            return {}

    if not isinstance(result, dict):
        return {}

    # Validate classifications
    valid_values = {c.value for c in ColumnClassification}
    validated = {}
    for col, cls in result.items():
        cls_lower = cls.lower().strip()
        if cls_lower in valid_values:
            validated[col] = cls_lower
        else:
            logger.warning("AI returned invalid classification '%s' for column '%s'", cls, col)

    return validated


# =========================================================================
# Snowflake Cortex classifier
# =========================================================================

class SnowflakeCortexClassifier:
    """Classify columns using Snowflake Cortex COMPLETE function.

    Data never leaves Snowflake — the LLM call runs inside the platform.
    """

    def __init__(self, sf_client, model: str = "mistral-large"):
        self._client = sf_client
        self._model = model

    def classify(
        self,
        columns: List[Dict[str, Any]],
        existing_rules: Dict[str, str],
        sample_values: Optional[Dict[str, List]] = None,
    ) -> Dict[str, str]:
        """Run AI classification via Cortex COMPLETE.

        Args:
            columns: Column metadata list [{name, data_type}, ...]
            existing_rules: Current {column_name: classification} mapping
            sample_values: Optional {column_name: [values]} for context

        Returns:
            Dict of overrides {column_name: new_classification}
        """
        user_prompt = _build_user_prompt(columns, existing_rules, sample_values)

        # Escape single quotes for SQL
        system_escaped = _SYSTEM_PROMPT.replace("'", "''")
        user_escaped = user_prompt.replace("'", "''")

        # Use the 2-argument COMPLETE(model, prompt) overload.
        # The 3-argument overload (COMPLETE$V6) requires the model arg to
        # be a compile-time string literal which fails with dynamic SQL.
        combined_prompt = (
            f"[SYSTEM] {_SYSTEM_PROMPT}\n\n[USER] {user_prompt}"
        ).replace("'", "''")

        sql = f"""
        SELECT SNOWFLAKE.CORTEX.COMPLETE(
            '{self._model}',
            '{combined_prompt}'
        ) AS response;
        """

        try:
            result = self._client.execute_query(sql, limit=1)
            if isinstance(result, list):
                rows = result
            else:
                rows = result.get("rows", [])
            if not rows or not rows[0]:
                return {}

            first_row = rows[0]
            if isinstance(first_row, dict):
                raw_response = first_row.get("RESPONSE") or next(iter(first_row.values()), "")
            else:
                raw_response = first_row[0] if first_row else ""
            # Cortex returns JSON with a 'choices' array
            try:
                resp_obj = json.loads(raw_response)
                content = resp_obj.get("choices", [{}])[0].get("messages", "")
                if not content:
                    content = resp_obj.get("choices", [{}])[0].get("message", {}).get("content", "")
                if not content:
                    content = raw_response
            except (json.JSONDecodeError, IndexError, KeyError):
                content = raw_response

            return _parse_response(str(content))
        except Exception as e:
            logger.error("Cortex classification failed: %s", e)
            raise


# =========================================================================
# Anthropic classifier
# =========================================================================

class AnthropicClassifier:
    """Classify columns using the Anthropic Claude API."""

    def __init__(self, api_key: str, model: str = "claude-sonnet-4-5-20250929"):
        try:
            from anthropic import Anthropic
        except ImportError:
            raise ImportError("anthropic package required: pip install anthropic")
        self._client = Anthropic(api_key=api_key)
        self._model = model
        self._last_usage = None

    def classify(
        self,
        columns: List[Dict[str, Any]],
        existing_rules: Dict[str, str],
        sample_values: Optional[Dict[str, List]] = None,
    ) -> Dict[str, str]:
        """Run AI classification via Anthropic API.

        Args:
            columns: Column metadata list [{name, data_type}, ...]
            existing_rules: Current {column_name: classification} mapping
            sample_values: Optional {column_name: [values]} for context

        Returns:
            Dict of overrides {column_name: new_classification}
        """
        user_prompt = _build_user_prompt(columns, existing_rules, sample_values)

        try:
            response = self._client.messages.create(
                model=self._model,
                max_tokens=2000,
                temperature=0.1,
                system=_SYSTEM_PROMPT,
                messages=[{"role": "user", "content": user_prompt}],
            )
            self._last_usage = {
                "input_tokens": response.usage.input_tokens,
                "output_tokens": response.usage.output_tokens,
                "model": self._model,
            }
            content = response.content[0].text
            return _parse_response(content)
        except Exception as e:
            logger.error("Anthropic classification failed: %s", e)
            raise


# =========================================================================
# Factory
# =========================================================================

def get_ai_classifier(
    provider: str,
    sf_client=None,
    api_key: Optional[str] = None,
    model: Optional[str] = None,
):
    """Create an AI classifier instance.

    Args:
        provider: "snowflake_cortex" or "anthropic"
        sf_client: SnowflakeClient (required for Cortex)
        api_key: Anthropic API key (required for Anthropic)
        model: Optional model override

    Returns:
        Classifier instance with a .classify() method
    """
    if provider == "snowflake_cortex":
        if sf_client is None:
            raise ValueError("sf_client required for Snowflake Cortex classifier")
        return SnowflakeCortexClassifier(sf_client, model=model or "mistral-large")

    elif provider == "anthropic":
        if not api_key:
            raise ValueError("api_key required for Anthropic classifier")
        return AnthropicClassifier(api_key, model=model or "claude-sonnet-4-5-20250929")

    else:
        raise ValueError(f"Unknown AI classifier provider: {provider}")
