"""AI Advisor for Data Modeling — optional LLM enrichment layer.

Adds AI-powered insights on top of deterministic analysis when available.
Falls back silently to no-op when no LLM is configured.

Uses Snowflake Cortex (via CortexClient) as primary LLM, with optional
Anthropic API fallback.
"""

from __future__ import annotations

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

logger = logging.getLogger(__name__)


class AiAdvisor:
    """Optional AI layer for data modeling operations.

    Provides enrichment methods for 4 core agents:
    - DimensionDetector (classification)
    - TableAnalyzer (domain/role analysis)
    - DMSpecGenerator (spec generation)
    - column_fk_inferer (relationship inference)

    All methods return None when no LLM is available or on failure,
    allowing callers to silently fall back to deterministic results.
    """

    def __init__(
        self,
        cortex_client=None,
        anthropic_key: Optional[str] = None,
        model: str = "mistral-large",
    ):
        self._cortex = cortex_client
        self._anthropic_key = anthropic_key
        self._model = model
        self._available = cortex_client is not None or anthropic_key is not None

    @property
    def available(self) -> bool:
        return self._available

    # ------------------------------------------------------------------
    # Internal LLM call
    # ------------------------------------------------------------------

    def _ask(self, prompt: str, max_tokens: int = 1000) -> Optional[str]:
        """Single LLM call with Cortex-first, Anthropic-fallback, None on failure."""
        # Try Cortex
        if self._cortex:
            try:
                result = self._cortex.complete(prompt, model=self._model)
                if result and result.success and result.result:
                    text = result.result if isinstance(result.result, str) else str(result.result)
                    return text
            except Exception as e:
                logger.debug("Cortex advisor call failed: %s", e)

        # Try Anthropic
        if self._anthropic_key:
            try:
                import anthropic
                client = anthropic.Anthropic(api_key=self._anthropic_key)
                response = client.messages.create(
                    model="claude-sonnet-4-20250514",
                    max_tokens=max_tokens,
                    messages=[{"role": "user", "content": prompt}],
                )
                return response.content[0].text
            except Exception as e:
                logger.debug("Anthropic advisor call failed: %s", e)

        return None

    def _ask_json(self, prompt: str, max_tokens: int = 1000) -> Optional[Dict]:
        """Ask LLM and parse JSON response. Returns None on failure."""
        text = self._ask(prompt, max_tokens=max_tokens)
        if not text:
            return None
        try:
            # Strip markdown code fences if present
            cleaned = text.strip()
            if cleaned.startswith("```"):
                first_nl = cleaned.index("\n")
                last_fence = cleaned.rfind("```")
                if last_fence > first_nl:
                    cleaned = cleaned[first_nl + 1:last_fence].strip()
            return json.loads(cleaned)
        except (json.JSONDecodeError, ValueError) as e:
            logger.debug("AI advisor JSON parse failed: %s", e)
            return None

    # ------------------------------------------------------------------
    # 1. DimensionDetector enrichment
    # ------------------------------------------------------------------

    def enhance_classification(
        self,
        classification: Dict[str, str],
        relationships: List[Dict],
        column_metadata: Optional[Dict[str, List[str]]] = None,
        sample_data: Optional[Any] = None,
    ) -> Optional[Dict[str, Any]]:
        """Enrich DimensionDetector output with AI reasoning.

        Returns dict with:
        - suggestions: list of {table, current, suggested, reason}
        - scd_types: dict of {table: "SCD1"|"SCD2"|...}
        - notes: free-text AI observations
        Returns None if unavailable or on failure.
        """
        if not self._available:
            return None

        # Build compact relationship summary (top 20)
        rel_summary = relationships[:20] if relationships else []

        # Build column samples (first 5 columns per table)
        col_samples = {}
        if column_metadata:
            for table, cols in column_metadata.items():
                col_samples[table] = cols[:5] if isinstance(cols, list) else []

        prompt = f"""You are a Kimball dimensional modeling expert. Given these table classifications \
(dimension/fact/bridge/unknown) derived from FK heuristics, review them and suggest \
corrections. Also recommend SCD types for dimension tables.

Tables and current classifications:
{json.dumps(classification, indent=2)}

Relationships (top 20):
{json.dumps(rel_summary, indent=2)}

Column samples (first 5 columns per table):
{json.dumps(col_samples, indent=2)}

Respond ONLY as JSON: {{"suggestions": [{{"table": "...", "current": "...", "suggested": "...", "reason": "..."}}], "scd_types": {{"TABLE": "SCD1|SCD2"}}, "notes": "..."}}"""

        return self._ask_json(prompt, max_tokens=1500)

    # ------------------------------------------------------------------
    # 2. TableAnalyzer enrichment
    # ------------------------------------------------------------------

    def enhance_table_analysis(
        self,
        tables: Any,
        column_metadata: Optional[Dict[str, List[str]]] = None,
        sample_data: Optional[Any] = None,
    ) -> Optional[Dict[str, Any]]:
        """Enrich TableAnalyzer output with AI domain inference.

        Returns dict with:
        - enrichments: list of {table, domain, purpose, related}
        Returns None if unavailable or on failure.
        """
        if not self._available:
            return None

        # Build table→columns summary
        table_col_json = {}
        if column_metadata:
            for table, cols in column_metadata.items():
                table_col_json[table] = cols[:10] if isinstance(cols, list) else []

        prompt = f"""You are analyzing database tables from an ERP system. For each table, suggest:
1. Business domain (Finance, Operations, HR, etc.)
2. One-sentence purpose description
3. Likely related tables

Tables with columns:
{json.dumps(table_col_json, indent=2)}

Respond ONLY as JSON: {{"enrichments": [{{"table": "...", "domain": "...", "purpose": "...", "related": ["..."]}}]}}"""

        return self._ask_json(prompt, max_tokens=2000)

    # ------------------------------------------------------------------
    # 3. DMSpecGenerator enrichment
    # ------------------------------------------------------------------

    def enhance_spec_generation(
        self,
        specs: List[Dict[str, Any]],
        classification: Dict[str, str],
        column_metadata: Optional[Dict[str, List[Dict]]] = None,
    ) -> Optional[Dict[str, Any]]:
        """Enrich DMSpecGenerator output with SCD type and grain recommendations.

        Returns dict with:
        - grain_analysis: {table: "grain description"}
        - suggested_measures: {table: ["col1", "col2"]}
        - naming_suggestions: {table: {"OLD_NAME": "better_name"}}
        Returns None if unavailable or on failure.
        """
        if not self._available:
            return None

        # Compact spec summary (target_table, source_tables, is_transaction, column names)
        spec_summary = []
        for s in specs[:15]:
            spec_summary.append({
                "target": s.get("target_table", ""),
                "sources": s.get("source_tables", []),
                "is_fact": s.get("is_transaction", False),
                "columns": [c.get("alias", "") for c in s.get("columns", [])][:15],
            })

        col_meta_summary = {}
        if column_metadata:
            for table, cols in list(column_metadata.items())[:15]:
                col_meta_summary[table] = [
                    c.get("name", c.get("COLUMN_NAME", "")) for c in cols[:15]
                ]

        prompt = f"""You are reviewing auto-generated dimensional model specifications. For each table:
1. Confirm or suggest the correct grain
2. Identify any missed measure columns
3. Suggest better column aliases where names are unclear

Table specs:
{json.dumps(spec_summary, indent=2)}

Source column metadata:
{json.dumps(col_meta_summary, indent=2)}

Respond ONLY as JSON: {{"grain_analysis": {{"TABLE": "grain description"}}, "suggested_measures": {{"TABLE": ["col"]}}, "naming_suggestions": {{"TABLE": {{"OLD": "new"}}}}}}"""

        return self._ask_json(prompt, max_tokens=2000)

    # ------------------------------------------------------------------
    # 4. Relationship inference enrichment
    # ------------------------------------------------------------------

    def enhance_relationship_inference(
        self,
        inferred_rels: List[Dict],
        table_columns: Dict[str, List[str]],
        sample_data: Optional[Any] = None,
    ) -> Optional[Dict[str, Any]]:
        """Find semantic relationships missed by name-matching heuristics.

        Returns dict with:
        - false_positives: list of {source_table, source_column, target_table, reason}
        - semantic_relationships: list of relationship dicts
        - confidence_scores: {"{src}.{col}->{tgt}": float}
        Returns None if unavailable or on failure.
        """
        if not self._available:
            return None

        # Compact relationship summary (top 30)
        rel_summary = []
        for r in inferred_rels[:30]:
            rel_summary.append({
                "src": r.get("source_table", ""),
                "src_col": r.get("source_column", ""),
                "tgt": r.get("target_table", ""),
                "tgt_col": r.get("target_column", ""),
                "method": r.get("method", ""),
            })

        # Compact table columns (first 10 cols each, max 20 tables)
        tc_summary = {}
        for table, cols in list(table_columns.items())[:20]:
            tc_summary[table] = cols[:10]

        prompt = f"""You are analyzing inferred FK relationships between database tables. These were found \
by column name matching (e.g., ACCOUNT_ID in table A -> ACCOUNT_ID in table B).

Review these relationships and:
1. Flag likely false positives (coincidental name matches)
2. Suggest semantic relationships that name-matching missed
3. Rate confidence for each relationship (0.0-1.0)

Inferred relationships:
{json.dumps(rel_summary, indent=2)}

Table columns:
{json.dumps(tc_summary, indent=2)}

Respond ONLY as JSON: {{"false_positives": [{{"source_table": "...", "source_column": "...", "target_table": "...", "reason": "..."}}], "semantic_relationships": [{{"source_table": "...", "source_column": "...", "target_table": "...", "target_column": "...", "reason": "..."}}], "confidence_scores": {{}}}}"""

        return self._ask_json(prompt, max_tokens=2000)
