"""
DimensionAssistant — Interactive star schema wizard.

Guides users through a 7-step conversational flow to design, customise,
and deploy Kimball-style star schema dimensions:

  interview → scan → propose → customize → properties → preview → deploy

State is persisted per-session in data/planner/dimension_sessions/.
"""

from __future__ import annotations

import json
import logging
import os
import re
import uuid
from dataclasses import dataclass, field, asdict
from datetime import datetime, timezone
from pathlib import Path
from typing import Any, Optional

logger = logging.getLogger(__name__)

# ── Steps ────────────────────────────────────────────────────────────────

WIZARD_STEPS = [
    "interview", "scan", "propose", "customize",
    "properties", "preview", "deploy",
]

STEP_LABELS = {
    "interview": "Interview",
    "scan": "Schema Scan",
    "propose": "Propose Dimensions",
    "customize": "Customize",
    "properties": "Properties",
    "preview": "Preview",
    "deploy": "Deploy",
}

# ── State ────────────────────────────────────────────────────────────────


@dataclass
class DimensionWizardState:
    """Full wizard state, serialised to JSON between turns."""

    session_id: str = ""
    conversation_id: str = ""
    current_step: str = "interview"
    started_at: str = ""
    updated_at: str = ""
    # Interview
    company_name: str = ""
    industry: str = ""  # oil_gas | general | manufacturing | saas | custom
    departments: list[str] = field(default_factory=list)
    reporting_needs: list[str] = field(default_factory=list)
    erp_system: str = ""
    interview_complete: bool = False
    # Scan (optional)
    connection_id: str = ""
    schema_info: dict[str, Any] = field(default_factory=dict)
    classification: dict[str, str] = field(default_factory=dict)
    scan_skipped: bool = False
    scan_database: str = ""
    scan_schema: str = ""
    scan_relationships: list[dict] = field(default_factory=list)
    scan_method: str = ""  # "dimension_detector" or "table_analyzer"
    # Proposal
    matched_template: str = ""
    dimensions: list[dict[str, Any]] = field(default_factory=list)
    facts: list[dict[str, Any]] = field(default_factory=list)
    bus_matrix: dict[str, list[str]] = field(default_factory=dict)
    # Deploy
    project_id: str = ""
    hierarchy_ids: list[str] = field(default_factory=list)
    deploy_status: str = ""

    # ── serialisation helpers ──

    def to_dict(self) -> dict[str, Any]:
        return asdict(self)

    @classmethod
    def from_dict(cls, d: dict[str, Any]) -> "DimensionWizardState":
        return cls(**{k: v for k, v in d.items() if k in cls.__dataclass_fields__})


# ── System prompt (loaded from file or inline fallback) ──────────────────

_SYSTEM_PROMPT_PATH = Path(__file__).resolve().parent.parent / "skills" / "prompts" / "dimension-builder.md"

_FALLBACK_SYSTEM_PROMPT = """\
You are the DataBridge Dimension Builder assistant. You guide users through
designing Kimball-style star schema dimensions using a structured wizard.
Always be concise, use markdown formatting, and follow the step instructions."""


def _load_system_prompt() -> str:
    if _SYSTEM_PROMPT_PATH.exists():
        return _SYSTEM_PROMPT_PATH.read_text(encoding="utf-8")
    return _FALLBACK_SYSTEM_PROMPT


# ── Industry keyword maps for interview extraction ───────────────────────

INDUSTRY_KEYWORDS = {
    "oil_gas": ["oil", "gas", "petroleum", "upstream", "midstream", "well", "lease",
                "enertia", "ogsys", "wolfepak", "openinvoice", "field", "basin", "loe",
                "production", "drilling", "e&p", "energy"],
    "manufacturing": ["manufacturing", "factory", "plant", "bom", "bill of materials",
                      "production line", "inventory", "cogs", "standard cost"],
    "saas": ["saas", "subscription", "arr", "mrr", "churn", "recurring revenue",
             "software", "platform", "tenant"],
    "general": ["enterprise", "general", "company", "business"],
}

ERP_KEYWORDS = {
    "enertia": ["enertia"],
    "sap": ["sap", "s/4hana", "s4hana"],
    "oracle": ["oracle", "netsuite", "jde", "jd edwards"],
    "dynamics": ["dynamics", "d365", "navision", "great plains", "gp"],
    "quickbooks": ["quickbooks", "qb"],
    "sage": ["sage", "intacct"],
}


def _detect_industry(text: str) -> str:
    """Score text against industry keyword maps, return best match."""
    text_lower = text.lower()
    scores: dict[str, int] = {}
    for industry, kws in INDUSTRY_KEYWORDS.items():
        scores[industry] = sum(1 for kw in kws if kw in text_lower)
    best = max(scores, key=scores.get)  # type: ignore[arg-type]
    return best if scores[best] > 0 else "general"


def _detect_erp(text: str) -> str:
    text_lower = text.lower()
    for erp, kws in ERP_KEYWORDS.items():
        if any(kw in text_lower for kw in kws):
            return erp
    return ""


def _extract_departments(text: str) -> list[str]:
    """Pull department-like phrases from natural language."""
    dept_patterns = [
        "finance", "accounting", "operations", "production", "engineering",
        "hr", "human resources", "sales", "marketing", "it", "technology",
        "supply chain", "procurement", "legal", "compliance", "treasury",
        "land", "geology", "reservoir", "revenue",
    ]
    text_lower = text.lower()
    return [d.title() for d in dept_patterns if d in text_lower]


def _extract_reporting_needs(text: str) -> list[str]:
    """Pull reporting-need phrases from natural language."""
    patterns = [
        r"(loe\s*(by|per|analysis)?\s*\w*)",
        r"(revenue\s*(by|per|analysis)?\s*\w*)",
        r"(production\s*(by|per|analysis|volume)?\s*\w*)",
        r"(cost\s*(by|per|analysis|center)?\s*\w*)",
        r"(profitability\s*\w*)",
        r"(variance\s*analysis)",
        r"(budget\s*vs\s*actual)",
        r"(inventory\s*\w*)",
        r"(sales\s*(by|per|analysis)?\s*\w*)",
        r"(gl\s*reconciliation)",
        r"(trial\s*balance)",
    ]
    found = []
    text_lower = text.lower()
    for p in patterns:
        m = re.search(p, text_lower)
        if m:
            found.append(m.group(0).strip().title())
    return found


def _dedup_case_insensitive(items: list[str]) -> list[str]:
    """Deduplicate strings case-insensitively, keeping the first occurrence."""
    seen: set[str] = set()
    result = []
    for item in items:
        key = item.lower().strip()
        if key not in seen:
            seen.add(key)
            result.append(item)
    return result


# ── DimensionAssistant ───────────────────────────────────────────────────


class DimensionAssistant:
    """Multi-turn conversational dimension wizard."""

    def __init__(self, anthropic_client=None):
        self._client = anthropic_client
        self._templates: dict[str, dict] = {}
        self._system_prompt = _load_system_prompt()

        # Lazy-init Anthropic client
        if self._client is None:
            try:
                from anthropic import Anthropic
                api_key = os.getenv("ANTHROPIC_API_KEY")
                if api_key:
                    self._client = Anthropic(api_key=api_key)
            except ImportError:
                pass

    # ── Template loading ─────────────────────────────────────────────────

    def _load_templates(self):
        if self._templates:
            return
        template_dir = Path(__file__).resolve().parent.parent / "skills" / "templates" / "dimensions"
        if not template_dir.exists():
            logger.warning("Dimension template directory not found: %s", template_dir)
            return
        for f in template_dir.glob("*.json"):
            try:
                tpl = json.loads(f.read_text(encoding="utf-8"))
                self._templates[tpl["id"]] = tpl
            except Exception as exc:
                logger.warning("Failed to load template %s: %s", f.name, exc)

    def _match_template(self, industry: str, departments: list[str],
                        reporting_needs: list[str]) -> str:
        """Score templates against interview answers, return best match ID."""
        self._load_templates()
        if not self._templates:
            return ""

        best_id = ""
        best_score = -1.0
        dept_lower = [d.lower() for d in departments]
        need_lower = [n.lower() for n in reporting_needs]

        for tid, tpl in self._templates.items():
            score = 0.0
            # Industry match
            if tpl.get("industry") == industry:
                score += 3.0
            elif tpl.get("industry") == "general":
                score += 0.5

            # Dimension name overlap with departments
            dim_names = [d["name"].lower() for d in tpl.get("dimensions", [])]
            for dn in dim_names:
                if any(dp in dn or dn in dp for dp in dept_lower):
                    score += 0.5

            # Fact/measure overlap with reporting needs
            for fact in tpl.get("facts", []):
                fact_words = fact["name"].lower().split()
                for need in need_lower:
                    if any(w in need for w in fact_words):
                        score += 0.5
                for measure in fact.get("measures", []):
                    m_words = measure.replace("_", " ").lower()
                    for need in need_lower:
                        if any(w in need for w in m_words.split()):
                            score += 0.3

            if score > best_score:
                best_score = score
                best_id = tid

        return best_id if best_score > 0 else ""

    # ── Main entry point ─────────────────────────────────────────────────

    def process_turn(self, message: str, state_dict: Optional[dict] = None,
                     nav_action: Optional[str] = None,
                     database: str = "", schema: str = "") -> dict[str, Any]:
        """Process one wizard turn.

        Returns: {session_id, step, response, state, preview, actions, complete}
        """
        # Initialise or restore state
        if state_dict:
            state = DimensionWizardState.from_dict(state_dict)
        else:
            state = DimensionWizardState(
                session_id=str(uuid.uuid4())[:8],
                conversation_id=str(uuid.uuid4())[:8],
                started_at=datetime.now(timezone.utc).isoformat(),
                current_step="interview",
            )
        state.updated_at = datetime.now(timezone.utc).isoformat()

        # Handle navigation actions
        if nav_action == "back":
            idx = WIZARD_STEPS.index(state.current_step)
            if idx > 0:
                state.current_step = WIZARD_STEPS[idx - 1]
        elif nav_action == "skip" and state.current_step == "scan":
            state.scan_skipped = True
            state.current_step = "propose"
        elif nav_action == "next":
            idx = WIZARD_STEPS.index(state.current_step)
            if idx < len(WIZARD_STEPS) - 1:
                state.current_step = WIZARD_STEPS[idx + 1]
        elif nav_action == "deploy":
            state.current_step = "deploy"
        elif nav_action == "scan":
            state.current_step = "scan"

        # Inject database/schema from request if provided
        if database:
            state.scan_database = database
        if schema:
            state.scan_schema = schema

        # Dispatch to step handler
        handler = {
            "interview": self._step_interview,
            "scan": self._step_scan,
            "propose": self._step_propose,
            "customize": self._step_customize,
            "properties": self._step_properties,
            "preview": self._step_preview,
            "deploy": self._step_deploy,
        }.get(state.current_step, self._step_interview)

        result = handler(message, state)

        # Build response
        step_idx = WIZARD_STEPS.index(state.current_step)
        progress = [
            {"name": STEP_LABELS[s], "status": "done" if i < step_idx else ("active" if i == step_idx else "pending")}
            for i, s in enumerate(WIZARD_STEPS)
        ]

        response = {
            "session_id": state.session_id,
            "conversation_id": state.conversation_id,
            "step": state.current_step,
            "step_label": STEP_LABELS.get(state.current_step, state.current_step),
            "response": result.get("response", ""),
            "state": state.to_dict(),
            "preview": self._build_preview(state),
            "actions": result.get("actions", []),
            "progress": progress,
            "complete": state.current_step == "deploy" and state.deploy_status == "complete",
        }
        # Pass through extra fields from step handlers (connections, scan_summary)
        for extra_key in ("connections", "scan_summary"):
            if extra_key in result:
                response[extra_key] = result[extra_key]
        return response

    # ── Step handlers ────────────────────────────────────────────────────

    def _step_interview(self, message: str, state: DimensionWizardState) -> dict:
        """Gather company, industry, departments, reporting needs."""
        if not message.strip():
            return {
                "response": (
                    "Welcome to the **Dimension Builder** wizard!\n\n"
                    "I'll help you design a star schema tailored to your business. "
                    "Let's start with some questions:\n\n"
                    "1. What **industry** is your company in? (e.g. Oil & Gas, Manufacturing, SaaS, General)\n"
                    "2. What **ERP system** do you use? (e.g. Enertia, SAP, Oracle, Dynamics)\n"
                    "3. Which **departments** need reporting? (e.g. Finance, Operations, Production)\n"
                    "4. What **reports** do you need? (e.g. LOE by well, Revenue by product, Cost variance)\n\n"
                    "You can answer all at once or one at a time."
                ),
                "actions": [],
            }

        # Extract structured data from message
        if not state.industry:
            state.industry = _detect_industry(message)
        if not state.erp_system:
            state.erp_system = _detect_erp(message)
        new_depts = _extract_departments(message)
        if new_depts:
            state.departments = _dedup_case_insensitive(state.departments + new_depts)
        new_needs = _extract_reporting_needs(message)
        if new_needs:
            state.reporting_needs = _dedup_case_insensitive(state.reporting_needs + new_needs)

        # Use LLM to extract any remaining info
        if self._client:
            extracted = self._llm_extract_interview(message, state)
            if extracted:
                if extracted.get("company_name") and not state.company_name:
                    state.company_name = extracted["company_name"]
                if extracted.get("industry") and not state.industry:
                    state.industry = extracted["industry"]
                if extracted.get("erp_system") and not state.erp_system:
                    state.erp_system = extracted["erp_system"]
                if extracted.get("departments"):
                    state.departments = _dedup_case_insensitive(state.departments + extracted["departments"])
                if extracted.get("reporting_needs"):
                    state.reporting_needs = _dedup_case_insensitive(state.reporting_needs + extracted["reporting_needs"])

        # Check completeness
        has_enough = bool(state.industry and (state.departments or state.reporting_needs))

        if has_enough:
            state.interview_complete = True
            summary = self._format_interview_summary(state)
            # Advance to scan
            state.current_step = "scan"
            return {
                "response": (
                    f"Great, here's what I gathered:\n\n{summary}\n\n"
                    "**Next step:** Would you like me to **scan your database** to discover "
                    "existing tables and classify them as dimensions/facts? "
                    "Or you can **skip** to go straight to dimension proposals based on "
                    "industry templates."
                ),
                "actions": [
                    {"id": "scan", "label": "Scan Database", "icon": "search"},
                    {"id": "skip", "label": "Skip to Proposals", "icon": "forward"},
                ],
            }

        # Need more info
        missing = []
        if not state.industry or state.industry == "general":
            missing.append("your **industry**")
        if not state.departments:
            missing.append("which **departments** need reporting")
        if not state.reporting_needs:
            missing.append("what **reports/metrics** you need")

        return {
            "response": (
                f"Thanks! I still need a bit more info — specifically: {', '.join(missing)}.\n\n"
                "Just describe your reporting needs naturally and I'll extract the details."
            ),
            "actions": [],
        }

    def _step_scan(self, message: str, state: DimensionWizardState) -> dict:
        """Connect to database and discover/classify tables."""
        if state.scan_skipped:
            state.current_step = "propose"
            return self._step_propose("", state)

        # Parse connection + database.schema from user message
        connection_id = state.connection_id
        database = state.scan_database
        schema = state.scan_schema
        msg = message.strip()

        if msg:
            # Parse patterns: "default", "default ENERTIA_RAW.ENERTIA_DBO", "ENERTIA_RAW.ENERTIA_DBO"
            parts = msg.split(None, 1)
            if len(parts) >= 1:
                # First token could be connection name or DATABASE.SCHEMA
                if "." in parts[0] and not connection_id:
                    # Looks like DATABASE.SCHEMA without connection, assume "default"
                    connection_id = "default"
                    db_schema = parts[0]
                else:
                    connection_id = parts[0]
                    db_schema = parts[1] if len(parts) > 1 else ""
                # Parse DATABASE.SCHEMA
                if db_schema and "." in db_schema:
                    db_parts = db_schema.split(".", 1)
                    database = db_parts[0].strip()
                    schema = db_parts[1].strip()

        # If no connection yet, show available connections
        if not connection_id:
            connections = self._list_available_connections()
            if connections:
                return {
                    "response": (
                        "Select a **Snowflake connection** to scan, or type a connection name "
                        "with database and schema:\n\n"
                        "Example: `default ENERTIA_RAW.ENERTIA_DBO`"
                    ),
                    "connections": connections,
                    "actions": [
                        {"id": "skip", "label": "Skip Scan", "icon": "forward"},
                    ],
                }
            return {
                "response": (
                    "No Snowflake connections found in `~/.snowflake/config.toml` or environment.\n\n"
                    "Type a connection name with database and schema, e.g.:\n"
                    "`default ENERTIA_RAW.ENERTIA_DBO`\n\n"
                    "Or click **Skip** to proceed with template-based proposals."
                ),
                "actions": [
                    {"id": "skip", "label": "Skip Scan", "icon": "forward"},
                ],
            }

        # If connection provided but no database/schema, try to pull from config
        if not database or not schema:
            connections = self._list_available_connections()
            for c in connections:
                if c["id"] == connection_id:
                    database = database or c.get("database", "")
                    schema = schema or c.get("schema", "")
                    break

        state.connection_id = connection_id
        state.scan_database = database
        state.scan_schema = schema

        # Run the scan
        scan_result = self._run_schema_scan(connection_id, database, schema)

        if scan_result.get("error"):
            # Reset connection so user can try again
            state.connection_id = ""
            return {
                "response": (
                    f"**Scan issue:** {scan_result['error']}\n\n"
                    "You can try a different connection or skip to template proposals."
                ),
                "actions": [
                    {"id": "scan", "label": "Try Again", "icon": "search"},
                    {"id": "skip", "label": "Skip to Proposals", "icon": "forward"},
                ],
            }

        # Store scan results in state
        state.schema_info = scan_result.get("schema_info", {})
        state.classification = scan_result.get("classification", {})
        state.scan_relationships = scan_result.get("relationships", [])
        state.scan_method = scan_result.get("method", "")

        # Build rich results summary
        table_count = state.schema_info.get("table_count", len(state.schema_info.get("tables", [])))
        dim_count = sum(1 for v in state.classification.values() if v == "dimension")
        fact_count = sum(1 for v in state.classification.values() if v == "fact")
        bridge_count = sum(1 for v in state.classification.values() if v == "bridge")
        unknown_count = table_count - dim_count - fact_count - bridge_count
        rel_count = len(state.scan_relationships)

        # Top dimensions by inbound references
        top_dims_text = ""
        if state.scan_relationships:
            dim_refs: dict[str, int] = {}
            for rel in state.scan_relationships:
                tgt = rel.get("target_table", "")
                if state.classification.get(tgt) == "dimension":
                    dim_refs[tgt] = dim_refs.get(tgt, 0) + 1
            if dim_refs:
                sorted_dims = sorted(dim_refs.items(), key=lambda x: x[1], reverse=True)[:5]
                top_dims_text = "\n**Top dimensions** (most referenced):\n"
                for tbl, refs in sorted_dims:
                    top_dims_text += f"- {tbl} ({refs} inbound refs)\n"

        method_label = {
            "dimension_detector": "DimensionDetector (FK-based)",
            "table_analyzer": "TableAnalyzer (heuristic)",
        }.get(state.scan_method, state.scan_method)

        state.current_step = "propose"

        return {
            "response": (
                f"**Scan complete!** Found **{table_count} tables** in {database}.{schema}:\n"
                f"- {dim_count} classified as **dimensions**\n"
                f"- {fact_count} classified as **facts**\n"
                f"- {bridge_count} classified as **bridge** tables\n"
                f"- {unknown_count} other (unknown/config)\n"
                f"- {rel_count} relationships inferred\n"
                f"{top_dims_text}\n"
                f"Classification method: {method_label}\n\n"
                "I'll now merge scan results with industry templates to propose a star schema."
            ),
            "actions": [],
            "scan_summary": {
                "table_count": table_count,
                "dim_count": dim_count,
                "fact_count": fact_count,
                "bridge_count": bridge_count,
                "unknown_count": unknown_count,
                "relationship_count": rel_count,
                "method": method_label,
            },
        }

    def _step_propose(self, message: str, state: DimensionWizardState) -> dict:
        """Load templates, score against interview, generate dimension proposals."""
        self._load_templates()

        # Match template
        if not state.matched_template:
            state.matched_template = self._match_template(
                state.industry, state.departments, state.reporting_needs
            )

        template = self._templates.get(state.matched_template)

        if template:
            # Start from template
            state.dimensions = [dict(d) for d in template.get("dimensions", [])]
            state.facts = [dict(f) for f in template.get("facts", [])]
            state.bus_matrix = dict(template.get("bus_matrix", {}))

            # If we have scan data, merge/enhance
            if state.schema_info:
                self._merge_scan_with_template(state, template)

            tpl_name = template.get("name", state.matched_template)
            dim_list = "\n".join(
                f"  - **{d['name']}** ({d.get('type', 'conformed')}) — "
                f"{len(d.get('levels', []))} levels"
                for d in state.dimensions
            )
            fact_list = "\n".join(
                f"  - **{f['name']}** — measures: {', '.join(f.get('measures', []))}"
                for f in state.facts
            )

            state.current_step = "customize"

            return {
                "response": (
                    f"Based on your input, I'm using the **{tpl_name}** template.\n\n"
                    f"### Proposed Dimensions\n{dim_list}\n\n"
                    f"### Proposed Fact Tables\n{fact_list}\n\n"
                    "The dimension tree and bus matrix are shown in the preview panel.\n\n"
                    "**Customize:** Tell me what to add, remove, or change. For example:\n"
                    "- *\"Add a Geography dimension with State > County > Township\"*\n"
                    "- *\"Remove the Partner dimension\"*\n"
                    "- *\"Add a capex measure to LOE facts\"*\n\n"
                    "Or click **Next** to accept and configure properties."
                ),
                "actions": [
                    {"id": "next", "label": "Accept & Continue", "icon": "check"},
                    {"id": "back", "label": "Back to Interview", "icon": "back"},
                ],
            }

        # No template match — use LLM or create minimal defaults
        state.dimensions, state.facts, state.bus_matrix = self._generate_default_schema(state)
        state.current_step = "customize"

        dim_list = "\n".join(
            f"  - **{d['name']}** — {len(d.get('levels', []))} levels"
            for d in state.dimensions
        )
        fact_list = "\n".join(
            f"  - **{f['name']}** — measures: {', '.join(f.get('measures', []))}"
            for f in state.facts
        )

        return {
            "response": (
                "I've generated a custom schema proposal:\n\n"
                f"### Dimensions\n{dim_list}\n\n"
                f"### Fact Tables\n{fact_list}\n\n"
                "Feel free to customize, or click **Next** to proceed."
            ),
            "actions": [
                {"id": "next", "label": "Accept & Continue", "icon": "check"},
            ],
        }

    def _step_customize(self, message: str, state: DimensionWizardState) -> dict:
        """Handle natural language edits to dimensions/facts."""
        if not message.strip():
            state.current_step = "properties"
            return self._step_properties("", state)

        # Use LLM to interpret edits
        edits = self._llm_interpret_edits(message, state)

        if edits.get("error"):
            return {
                "response": (
                    f"I wasn't sure how to apply that change: {edits['error']}\n\n"
                    "Try being more specific, e.g.:\n"
                    "- *\"Add a Product dimension with Category > Subcategory > Product\"*\n"
                    "- *\"Remove dim_partner\"*\n"
                    "- *\"Add water_volume measure to Production facts\"*"
                ),
                "actions": [
                    {"id": "next", "label": "Done Customizing", "icon": "check"},
                ],
            }

        # Apply edits
        changes_made = []
        for edit in edits.get("edits", []):
            change = self._apply_edit(edit, state)
            if change:
                changes_made.append(change)

        if changes_made:
            changes_text = "\n".join(f"- {c}" for c in changes_made)
            return {
                "response": (
                    f"Applied changes:\n{changes_text}\n\n"
                    "The preview panel has been updated. Make more changes or click **Next**."
                ),
                "actions": [
                    {"id": "next", "label": "Done Customizing", "icon": "check"},
                    {"id": "back", "label": "Back", "icon": "back"},
                ],
            }

        return {
            "response": "No changes detected. Describe what you'd like to modify, or click **Next** to proceed.",
            "actions": [
                {"id": "next", "label": "Done Customizing", "icon": "check"},
            ],
        }

    def _step_properties(self, message: str, state: DimensionWizardState) -> dict:
        """Configure dimension/fact properties (aggregation, drill, format)."""
        if not message.strip() and state.current_step == "properties":
            # Show current properties with suggested defaults
            props = self._suggest_properties(state)
            prop_text = self._format_properties(props)

            state.current_step = "properties"

            return {
                "response": (
                    "### Property Configuration\n\n"
                    f"{prop_text}\n\n"
                    "These are suggested defaults based on Kimball best practices. "
                    "You can adjust any property — just tell me what to change, e.g.:\n"
                    "- *\"Set LOE aggregation to AVG\"*\n"
                    "- *\"Disable drill on Date dimension\"*\n\n"
                    "Or click **Next** to accept defaults and preview."
                ),
                "actions": [
                    {"id": "next", "label": "Accept & Preview", "icon": "check"},
                    {"id": "back", "label": "Back to Customize", "icon": "back"},
                ],
            }

        if message.strip():
            # Apply property changes
            changes = self._apply_property_changes(message, state)
            if changes:
                return {
                    "response": f"Updated properties:\n{changes}\n\nMake more changes or click **Next**.",
                    "actions": [
                        {"id": "next", "label": "Accept & Preview", "icon": "check"},
                    ],
                }

        # Advance to preview
        state.current_step = "preview"
        return self._step_preview("", state)

    def _step_preview(self, message: str, state: DimensionWizardState) -> dict:
        """Full preview: tree + bus matrix + DDL + report suggestions."""
        dim_count = len(state.dimensions)
        fact_count = len(state.facts)
        level_count = sum(len(d.get("levels", [])) for d in state.dimensions)
        measure_count = sum(len(f.get("measures", [])) for f in state.facts)

        ddl_preview = self._generate_ddl_preview(state)

        state.current_step = "preview"

        return {
            "response": (
                "### Star Schema Preview\n\n"
                f"- **{dim_count}** dimensions with **{level_count}** total levels\n"
                f"- **{fact_count}** fact tables with **{measure_count}** measures\n"
                f"- **{len(state.bus_matrix)}** bus matrix relationships\n\n"
                "Check the preview panel for the full dimension tree and bus matrix.\n\n"
                f"### DDL Preview\n```sql\n{ddl_preview}\n```\n\n"
                "When ready, click **Deploy** to create hierarchy projects and "
                "optionally deploy to Snowflake."
            ),
            "actions": [
                {"id": "deploy", "label": "Deploy", "icon": "rocket"},
                {"id": "back", "label": "Back to Customize", "icon": "back"},
            ],
        }

    def _step_deploy(self, message: str, state: DimensionWizardState) -> dict:
        """Create hierarchy projects and deploy to database."""
        if state.deploy_status == "complete":
            return {
                "response": (
                    "Deployment is already complete!\n\n"
                    f"- Project ID: `{state.project_id}`\n"
                    f"- Hierarchies created: {len(state.hierarchy_ids)}\n\n"
                    "You can view your hierarchies in the Hierarchy Builder page."
                ),
                "actions": [],
            }

        # Create hierarchy projects
        deploy_result = self._run_deploy(state)

        if deploy_result.get("error"):
            return {
                "response": (
                    f"Deployment encountered an issue: {deploy_result['error']}\n\n"
                    "You can try again or go back to adjust the schema."
                ),
                "actions": [
                    {"id": "deploy", "label": "Retry Deploy", "icon": "rocket"},
                    {"id": "back", "label": "Back", "icon": "back"},
                ],
            }

        state.project_id = deploy_result.get("project_id", "")
        state.hierarchy_ids = deploy_result.get("hierarchy_ids", [])
        state.deploy_status = "complete"

        return {
            "response": (
                "### Deployment Complete!\n\n"
                f"- **Project:** `{state.project_id}`\n"
                f"- **Hierarchies created:** {len(state.hierarchy_ids)}\n"
                f"  - {chr(10).join(state.hierarchy_ids)}\n\n"
                "Your star schema dimensions have been created as hierarchy projects. "
                "You can now:\n"
                "- View them in the **Hierarchy Builder** page\n"
                "- Run a **Wright pipeline** to generate dbt models\n"
                "- Use **Data Quality** tools to validate the schema"
            ),
            "actions": [],
        }

    # ── LLM helpers ──────────────────────────────────────────────────────

    def _llm_extract_interview(self, message: str, state: DimensionWizardState) -> dict:
        """Use Claude to extract structured interview data."""
        if not self._client:
            return {}

        try:
            response = self._client.messages.create(
                model="claude-sonnet-4-20250514",
                max_tokens=1024,
                temperature=0.2,
                system=(
                    "Extract structured information from the user's message about their company. "
                    "Return JSON with keys: company_name, industry (oil_gas|general|manufacturing|saas|custom), "
                    "erp_system, departments (list), reporting_needs (list). "
                    "Only include keys where you found information. Return {} if nothing found."
                ),
                messages=[{"role": "user", "content": message}],
            )
            text = response.content[0].text.strip()
            # Extract JSON from response
            json_match = re.search(r'\{[^{}]*\}', text, re.DOTALL)
            if json_match:
                return json.loads(json_match.group())
        except Exception as exc:
            logger.warning("LLM interview extraction failed: %s", exc)

        return {}

    def _llm_interpret_edits(self, message: str, state: DimensionWizardState) -> dict:
        """Use Claude to interpret natural language edits."""
        if not self._client:
            return self._heuristic_interpret_edits(message, state)

        current_dims = json.dumps([{"id": d["id"], "name": d["name"]} for d in state.dimensions])
        current_facts = json.dumps([{"id": f["id"], "name": f["name"]} for f in state.facts])

        try:
            response = self._client.messages.create(
                model="claude-sonnet-4-20250514",
                max_tokens=1024,
                temperature=0.2,
                system=(
                    "You interpret edits to a star schema. Current dimensions: " + current_dims +
                    ". Current facts: " + current_facts + ". "
                    "Return JSON: {edits: [{action: 'add'|'remove'|'modify', "
                    "target: 'dimension'|'fact'|'measure'|'level', "
                    "id: str, name: str, details: {}}]}. "
                    "For add dimension, include levels array [{name, level}]. "
                    "For add measure, include fact_id and measure_name."
                ),
                messages=[{"role": "user", "content": message}],
            )
            text = response.content[0].text.strip()
            json_match = re.search(r'\{.*\}', text, re.DOTALL)
            if json_match:
                return json.loads(json_match.group())
        except Exception as exc:
            logger.warning("LLM edit interpretation failed: %s", exc)

        return self._heuristic_interpret_edits(message, state)

    # ── Heuristic helpers ────────────────────────────────────────────────

    def _heuristic_interpret_edits(self, message: str, state: DimensionWizardState) -> dict:
        """Fallback edit interpretation without LLM."""
        edits = []
        msg_lower = message.lower()

        # Add dimension pattern
        add_dim = re.search(r'add\s+(?:a\s+)?(\w[\w\s]*?)\s+dimension', msg_lower)
        if add_dim:
            name = add_dim.group(1).strip().title()
            # Look for levels pattern: "with X > Y > Z"
            levels = []
            level_match = re.search(r'with\s+([\w\s>]+)', msg_lower)
            if level_match:
                parts = [p.strip().title() for p in level_match.group(1).split(">")]
                levels = [{"name": p, "level": i + 1} for i, p in enumerate(parts)]
            dim_id = f"dim_{name.lower().replace(' ', '_')}"
            edits.append({
                "action": "add", "target": "dimension",
                "id": dim_id, "name": name,
                "details": {"levels": levels, "type": "conformed"},
            })

        # Remove dimension pattern
        remove_dim = re.search(r'remove\s+(?:the\s+)?(\w[\w\s]*?)\s*(?:dimension)?$', msg_lower)
        if remove_dim:
            name = remove_dim.group(1).strip()
            # Find matching dim
            for d in state.dimensions:
                if name.lower() in d["name"].lower() or name.lower() in d["id"].lower():
                    edits.append({"action": "remove", "target": "dimension", "id": d["id"], "name": d["name"]})
                    break

        # Add measure pattern
        add_measure = re.search(r'add\s+(?:a\s+)?(\w[\w\s]*?)\s+(?:measure|metric)\s+to\s+(\w[\w\s]*)', msg_lower)
        if add_measure:
            measure = add_measure.group(1).strip().lower().replace(" ", "_")
            fact_name = add_measure.group(2).strip()
            for f in state.facts:
                if fact_name.lower() in f["name"].lower() or fact_name.lower() in f["id"].lower():
                    edits.append({
                        "action": "add", "target": "measure",
                        "id": f["id"], "name": measure,
                        "details": {"fact_id": f["id"], "measure_name": measure},
                    })
                    break

        if not edits:
            return {"error": "Could not parse edit request"}

        return {"edits": edits}

    def _apply_edit(self, edit: dict, state: DimensionWizardState) -> str:
        """Apply a single edit to state. Returns change description."""
        action = edit.get("action")
        target = edit.get("target")

        if action == "add" and target == "dimension":
            new_dim = {
                "id": edit.get("id", f"dim_{edit['name'].lower().replace(' ', '_')}"),
                "name": edit["name"],
                "type": edit.get("details", {}).get("type", "conformed"),
                "levels": edit.get("details", {}).get("levels", []),
                "properties": {"drill_enabled": True, "aggregation": "SUM"},
            }
            state.dimensions.append(new_dim)
            level_names = " > ".join(l["name"] for l in new_dim["levels"]) if new_dim["levels"] else "no levels"
            return f"Added **{new_dim['name']}** dimension ({level_names})"

        if action == "remove" and target == "dimension":
            before = len(state.dimensions)
            state.dimensions = [d for d in state.dimensions if d["id"] != edit["id"]]
            # Also remove from bus matrix
            for fact_id in list(state.bus_matrix.keys()):
                if edit["id"] in state.bus_matrix[fact_id]:
                    state.bus_matrix[fact_id].remove(edit["id"])
            if len(state.dimensions) < before:
                return f"Removed **{edit['name']}** dimension"

        if action == "add" and target == "measure":
            fact_id = edit.get("details", {}).get("fact_id")
            measure_name = edit.get("details", {}).get("measure_name", edit.get("name"))
            for f in state.facts:
                if f["id"] == fact_id:
                    f.setdefault("measures", []).append(measure_name)
                    return f"Added measure `{measure_name}` to **{f['name']}**"

        if action == "add" and target == "fact":
            new_fact = {
                "id": edit.get("id", f"fct_{edit['name'].lower().replace(' ', '_')}"),
                "name": edit["name"],
                "measures": edit.get("details", {}).get("measures", []),
            }
            state.facts.append(new_fact)
            return f"Added **{new_fact['name']}** fact table"

        if action == "remove" and target == "fact":
            before = len(state.facts)
            state.facts = [f for f in state.facts if f["id"] != edit["id"]]
            state.bus_matrix.pop(edit["id"], None)
            if len(state.facts) < before:
                return f"Removed **{edit['name']}** fact table"

        return ""

    def _format_interview_summary(self, state: DimensionWizardState) -> str:
        lines = []
        if state.company_name:
            lines.append(f"- **Company:** {state.company_name}")
        lines.append(f"- **Industry:** {state.industry.replace('_', ' ').title()}")
        if state.erp_system:
            lines.append(f"- **ERP:** {state.erp_system.title()}")
        if state.departments:
            lines.append(f"- **Departments:** {', '.join(state.departments)}")
        if state.reporting_needs:
            lines.append(f"- **Reporting needs:** {', '.join(state.reporting_needs)}")
        return "\n".join(lines)

    def _suggest_properties(self, state: DimensionWizardState) -> dict:
        """Suggest default properties for dimensions and facts."""
        props: dict[str, Any] = {"dimensions": {}, "facts": {}}

        for dim in state.dimensions:
            dim_props = {
                "drill_enabled": True,
                "aggregation": "SUM",
                "sort_order": "name",
            }
            if dim.get("type") == "role-playing":
                dim_props["role_playing"] = True
            props["dimensions"][dim["id"]] = dim_props
            dim.setdefault("properties", {}).update(dim_props)

        for fact in state.facts:
            fact_props = {}
            for measure in fact.get("measures", []):
                m_lower = measure.lower()
                if any(w in m_lower for w in ["count", "qty", "quantity", "volume"]):
                    fact_props[measure] = {"type": "additive", "aggregation": "SUM", "format": "#,##0"}
                elif any(w in m_lower for w in ["rate", "per_", "ratio", "percent"]):
                    fact_props[measure] = {"type": "non-additive", "aggregation": "AVG", "format": "#,##0.00"}
                elif any(w in m_lower for w in ["cost", "revenue", "amount", "price", "tax"]):
                    fact_props[measure] = {"type": "additive", "aggregation": "SUM", "format": "$#,##0.00"}
                else:
                    fact_props[measure] = {"type": "additive", "aggregation": "SUM", "format": "#,##0.00"}
            props["facts"][fact["id"]] = fact_props
            fact.setdefault("properties", {}).update(fact_props)

        return props

    def _format_properties(self, props: dict) -> str:
        lines = []
        if props.get("dimensions"):
            lines.append("**Dimension Properties:**")
            for dim_id, dp in props["dimensions"].items():
                lines.append(f"- `{dim_id}`: drill={dp.get('drill_enabled')}, agg={dp.get('aggregation')}")

        if props.get("facts"):
            lines.append("\n**Fact Measure Properties:**")
            for fact_id, measures in props["facts"].items():
                lines.append(f"- `{fact_id}`:")
                for m_name, mp in measures.items():
                    lines.append(f"  - `{m_name}`: type={mp.get('type')}, agg={mp.get('aggregation')}, format={mp.get('format')}")

        return "\n".join(lines)

    def _apply_property_changes(self, message: str, state: DimensionWizardState) -> str:
        """Parse and apply property changes from natural language."""
        changes = []
        msg_lower = message.lower()

        # Aggregation change
        agg_match = re.search(r'set\s+(\w+)\s+aggregation\s+to\s+(\w+)', msg_lower)
        if agg_match:
            target = agg_match.group(1)
            agg = agg_match.group(2).upper()
            for dim in state.dimensions:
                if target in dim["id"].lower() or target in dim["name"].lower():
                    dim.setdefault("properties", {})["aggregation"] = agg
                    changes.append(f"- Set **{dim['name']}** aggregation to `{agg}`")
                    break
            for fact in state.facts:
                if target in fact["id"].lower() or target in fact["name"].lower():
                    for m in fact.get("measures", []):
                        fact.setdefault("properties", {})[m] = fact.get("properties", {}).get(m, {})
                        fact["properties"][m]["aggregation"] = agg
                    changes.append(f"- Set **{fact['name']}** measure aggregation to `{agg}`")
                    break

        # Drill toggle
        drill_match = re.search(r'(enable|disable)\s+drill\s+(?:on\s+)?(\w+)', msg_lower)
        if drill_match:
            enabled = drill_match.group(1) == "enable"
            target = drill_match.group(2)
            for dim in state.dimensions:
                if target in dim["id"].lower() or target in dim["name"].lower():
                    dim.setdefault("properties", {})["drill_enabled"] = enabled
                    changes.append(f"- {'Enabled' if enabled else 'Disabled'} drill on **{dim['name']}**")
                    break

        return "\n".join(changes) if changes else ""

    def _generate_ddl_preview(self, state: DimensionWizardState) -> str:
        """Generate simplified DDL preview."""
        lines = []
        for dim in state.dimensions:
            cols = ["id INT PRIMARY KEY"]
            for level in dim.get("levels", []):
                col_name = level["name"].lower().replace(" ", "_")
                cols.append(f"{col_name} VARCHAR(255)")
            cols.append("created_at TIMESTAMP DEFAULT CURRENT_TIMESTAMP")
            col_str = ",\n  ".join(cols)
            lines.append(f"CREATE TABLE {dim['id']} (\n  {col_str}\n);")

        for fact in state.facts:
            cols = ["id INT PRIMARY KEY"]
            # FK columns for each dimension in bus matrix
            dim_ids = state.bus_matrix.get(fact["id"], [])
            for dim_id in dim_ids:
                cols.append(f"{dim_id}_key INT REFERENCES {dim_id}(id)")
            for measure in fact.get("measures", []):
                cols.append(f"{measure} DECIMAL(18,4)")
            cols.append("created_at TIMESTAMP DEFAULT CURRENT_TIMESTAMP")
            col_str = ",\n  ".join(cols)
            lines.append(f"CREATE TABLE {fact['id']} (\n  {col_str}\n);")

        return "\n\n".join(lines) if lines else "-- No tables defined yet"

    def _generate_default_schema(self, state: DimensionWizardState) -> tuple:
        """Generate a minimal default schema based on interview data."""
        dims = [
            {"id": "dim_date", "name": "Date", "type": "role-playing",
             "levels": [{"name": "Year", "level": 1}, {"name": "Quarter", "level": 2},
                        {"name": "Month", "level": 3}, {"name": "Day", "level": 4}],
             "properties": {"drill_enabled": True, "aggregation": "SUM"}},
            {"id": "dim_account", "name": "Account", "type": "conformed",
             "levels": [{"name": "Account Group", "level": 1}, {"name": "Account", "level": 2}],
             "properties": {"drill_enabled": True, "aggregation": "SUM"}},
        ]

        if any("operations" in d.lower() for d in state.departments):
            dims.append({
                "id": "dim_entity", "name": "Entity", "type": "conformed",
                "levels": [{"name": "Company", "level": 1}, {"name": "Department", "level": 2}],
                "properties": {"drill_enabled": True, "aggregation": "SUM"},
            })

        facts = [
            {"id": "fct_gl", "name": "General Ledger",
             "measures": ["debit_amount", "credit_amount", "net_amount"]},
        ]

        bus_matrix = {"fct_gl": [d["id"] for d in dims]}

        return dims, facts, bus_matrix

    def _merge_scan_with_template(self, state: DimensionWizardState, template: dict):
        """Enrich template dimensions with scan-discovered data.

        1. Match scan-discovered dimensions to template dims by name similarity
        2. Add scan-discovered facts not in template
        3. Populate source_tables on matched dimensions
        4. Merge scan bus_matrix with template bus_matrix
        5. Add columns_by_table to dimension metadata
        """
        tables = state.schema_info.get("tables", [])
        classification = state.classification
        columns_by_table = state.schema_info.get("columns_by_table", {})

        # 1. Match scan dimensions to template dimensions using source_hints + name similarity
        for dim in state.dimensions:
            hints = dim.get("source_hints", [])
            dim_name_lower = dim["name"].lower().replace(" ", "_")
            matched_tables = []

            for table in tables:
                tname = table if isinstance(table, str) else table.get("name", "")
                if classification.get(tname) != "dimension":
                    continue

                # Check source_hints first
                hint_match = False
                for hint in hints:
                    pattern = hint.replace("*", ".*")
                    if re.match(pattern, tname, re.IGNORECASE):
                        hint_match = True
                        break

                # Name similarity: dim "Well" matches table "WELL", "WELL_HEADER", etc.
                tname_lower = tname.lower()
                name_match = (
                    dim_name_lower in tname_lower
                    or tname_lower in dim_name_lower
                    or tname_lower.startswith(dim_name_lower + "_")
                    or tname_lower.endswith("_" + dim_name_lower)
                )

                if hint_match or name_match:
                    matched_tables.append(tname)

            if matched_tables:
                dim["source_tables"] = matched_tables
                # Add column info from first matched source table
                first_table = matched_tables[0]
                if first_table in columns_by_table:
                    dim["source_columns"] = columns_by_table[first_table]

        # 2. Add scan-discovered facts not already in template
        template_fact_ids = {f["id"] for f in state.facts}
        for table in tables:
            tname = table if isinstance(table, str) else table.get("name", "")
            if classification.get(tname) != "fact":
                continue
            fact_id = f"fct_{tname.lower()}"
            if fact_id not in template_fact_ids:
                # Build measures from numeric columns
                measures = []
                if tname in columns_by_table:
                    for col in columns_by_table[tname]:
                        col_lower = col.lower()
                        if any(kw in col_lower for kw in ["amount", "qty", "quantity", "cost", "price",
                                                           "rate", "volume", "count", "total", "revenue"]):
                            measures.append(col.lower())
                state.facts.append({
                    "id": fact_id, "name": tname.replace("_", " ").title(),
                    "measures": measures[:8],  # cap at 8
                    "source_table": tname,
                })

        # 3. Merge scan bus_matrix with template bus_matrix
        bus_result = getattr(state, "_scan_bus_matrix_result", None)
        if not bus_result:
            bus_result = state.schema_info.get("bus_matrix_result", {})
        if bus_result and bus_result.get("matrix"):
            scan_matrix = bus_result["matrix"]
            dim_id_map = {d["name"].upper(): d["id"] for d in state.dimensions}
            for fact_table, dim_links in scan_matrix.items():
                fact_id = f"fct_{fact_table.lower()}"
                for ft in state.facts:
                    if ft.get("source_table", "").upper() == fact_table.upper() or ft["id"] == fact_id:
                        for dim_table, linked in dim_links.items():
                            if linked:
                                dim_id = dim_id_map.get(dim_table.upper())
                                if dim_id:
                                    state.bus_matrix.setdefault(ft["id"], [])
                                    if dim_id not in state.bus_matrix[ft["id"]]:
                                        state.bus_matrix[ft["id"]].append(dim_id)

    def _build_preview(self, state: DimensionWizardState) -> dict:
        """Build preview data for the UI."""
        tree = []

        # Dimensions
        for dim in state.dimensions:
            node = {
                "id": dim["id"],
                "name": dim["name"],
                "type": "dimension",
                "dim_type": dim.get("type", "conformed"),
                "children": [
                    {"id": f"{dim['id']}_L{l.get('level', i+1)}", "name": l["name"], "type": "level"}
                    for i, l in enumerate(dim.get("levels", []))
                ],
            }
            tree.append(node)

        # Facts
        for fact in state.facts:
            node = {
                "id": fact["id"],
                "name": fact["name"],
                "type": "fact",
                "children": [
                    {"id": f"{fact['id']}_{m}", "name": m.replace("_", " ").title(), "type": "measure"}
                    for m in fact.get("measures", [])
                ],
            }
            tree.append(node)

        return {
            "tree": tree,
            "bus_matrix": state.bus_matrix,
            "dimensions": [{"id": d["id"], "name": d["name"]} for d in state.dimensions],
            "facts": [{"id": f["id"], "name": f["name"]} for f in state.facts],
        }

    # ── External integrations (scan, deploy) ─────────────────────────────

    def _list_available_connections(self) -> list[dict]:
        """List Snowflake connections from config.toml + env vars."""
        connections = []
        try:
            config_path = Path(os.environ.get("USERPROFILE", os.environ.get("HOME", str(Path.home())))) / ".snowflake" / "config.toml"
            if config_path.exists():
                # Load TOML inline to avoid triggering heavy data_modeling __init__
                try:
                    import tomllib  # Python 3.11+
                except ImportError:
                    import tomli as tomllib  # type: ignore[no-redef]
                with open(config_path, "rb") as f:
                    config = tomllib.load(f)
                for name, params in config.get("connections", {}).items():
                    if isinstance(params, dict):
                        connections.append({
                            "id": name, "name": name,
                            "database": params.get("database", ""),
                            "schema": params.get("schema", ""),
                        })
        except Exception as exc:
            logger.warning("Failed to list connections from config.toml: %s", exc)
        # Environment variable fallback
        if os.getenv("SNOWFLAKE_ACCOUNT"):
            connections.append({
                "id": "__env__", "name": "Environment Variables",
                "database": os.getenv("SNOWFLAKE_DATABASE", ""),
                "schema": os.getenv("SNOWFLAKE_SCHEMA", ""),
            })
        return connections

    def _run_schema_scan(self, connection_id: str, database: str = "", schema: str = "") -> dict:
        """Run full schema discovery pipeline.

        Stages:
          1. Connect via sf_pool
          2. Discover tables via INFORMATION_SCHEMA.TABLES
          3. Fetch column metadata via INFORMATION_SCHEMA.COLUMNS
          4. Load sample data via load_snowflake_samples()
          5. Infer relationships via infer_relationships()
          6. Classify tables via DimensionDetector (FK-based) or TableAnalyzer (fallback)
          7. Generate bus matrix via BusMatrixGenerator
        """
        try:
            from src.data_modeling.snowflake_pool import sf_pool
        except ImportError:
            return {"error": "Data modeling modules not available (sf_pool)"}

        # Stage 1 — Connect
        try:
            conn = sf_pool.get(connection_id, database=database, schema=schema)
            if not conn:
                return {"error": f"Connection '{connection_id}' not found or unavailable"}
        except Exception as exc:
            logger.warning("Connection failed for '%s': %s", connection_id, exc)
            return {"error": f"Connection failed: {exc}"}

        # Resolve database/schema from connection if not provided
        if not database:
            try:
                database = conn.database or ""
            except Exception:
                pass
        if not schema:
            try:
                schema = conn.schema or ""
            except Exception:
                pass
        if not database or not schema:
            return {"error": f"Database and schema are required. Got database='{database}', schema='{schema}'. "
                    "Please specify as CONNECTION DATABASE.SCHEMA (e.g. 'default ENERTIA_RAW.ENERTIA_DBO')"}

        # Stage 2 — Discover tables
        try:
            cur = conn.cursor()
            cur.execute(
                f"SELECT TABLE_NAME FROM {database}.INFORMATION_SCHEMA.TABLES "
                f"WHERE TABLE_SCHEMA = '{schema}' AND TABLE_TYPE = 'BASE TABLE' "
                f"ORDER BY TABLE_NAME"
            )
            table_list = [row[0] for row in cur.fetchall()]
            cur.close()
        except Exception as exc:
            logger.warning("Table discovery failed: %s", exc)
            return {"error": f"Table discovery failed: {exc}"}

        if not table_list:
            return {"error": f"No tables found in {database}.{schema}"}

        # Stage 3 — Column metadata
        column_metadata = []
        columns_by_table: dict[str, list[str]] = {}
        try:
            cur = conn.cursor()
            table_in = ",".join(f"'{t}'" for t in table_list)
            cur.execute(
                f"SELECT TABLE_NAME, COLUMN_NAME, DATA_TYPE, ORDINAL_POSITION "
                f"FROM {database}.INFORMATION_SCHEMA.COLUMNS "
                f"WHERE TABLE_SCHEMA = '{schema}' AND TABLE_NAME IN ({table_in}) "
                f"ORDER BY TABLE_NAME, ORDINAL_POSITION"
            )
            for row in cur.fetchall():
                column_metadata.append({
                    "TABLE_NAME": row[0], "COLUMN_NAME": row[1],
                    "DATA_TYPE": row[2], "ORDINAL_POSITION": row[3],
                })
                columns_by_table.setdefault(row[0], []).append(row[1])
            cur.close()
        except Exception as exc:
            logger.warning("Column metadata query failed: %s", exc)
            # Continue — we can still classify without full column metadata

        # Stage 4 — Load sample data
        sample_tables = {}
        try:
            from src.data_modeling.sample_loader import load_snowflake_samples
            sample_tables = load_snowflake_samples(
                database, schema, table_list, limit=100, sf_connection=conn
            )
        except ImportError:
            logger.warning("sample_loader not available, skipping sample loading")
        except Exception as exc:
            logger.warning("Sample loading failed (continuing): %s", exc)

        # Stage 5 — Infer relationships
        relationships = []
        if sample_tables:
            try:
                from src.data_modeling.relationship_inferer import infer_relationships
                relationships = infer_relationships(sample_tables, min_overlap=0.3)
            except ImportError:
                logger.warning("relationship_inferer not available")
            except Exception as exc:
                logger.warning("Relationship inference failed (continuing): %s", exc)

        # Stage 6 — Classify tables
        classification: dict[str, str] = {}
        method = "none"
        try:
            if relationships:
                from src.data_modeling.dimension_detector import DimensionDetector
                classification = DimensionDetector().classify_tables(
                    relationships, threshold=2, column_metadata=columns_by_table
                )
                method = "dimension_detector"
            else:
                # Fallback: heuristic-based classification without FK data
                from src.data_modeling.table_analyzer import TableAnalyzer
                analyzer = TableAnalyzer(column_metadata=column_metadata)
                enhanced = analyzer.analyze_all(table_list)
                classification = enhanced.classification
                method = "table_analyzer"
        except ImportError as exc:
            logger.warning("Classification module not available: %s", exc)
        except Exception as exc:
            logger.warning("Classification failed (continuing): %s", exc)

        # Stage 7 — Bus matrix
        bus_matrix_result = {}
        if relationships and classification:
            try:
                from src.data_modeling.bus_matrix import BusMatrixGenerator
                bus_matrix_result = BusMatrixGenerator().generate(relationships, classification)
            except ImportError:
                logger.warning("BusMatrixGenerator not available")
            except Exception as exc:
                logger.warning("Bus matrix generation failed (continuing): %s", exc)

        return {
            "schema_info": {
                "tables": table_list,
                "column_metadata": column_metadata,
                "columns_by_table": columns_by_table,
                "table_count": len(table_list),
                "sample_row_count": sum(len(v.get("rows", [])) for v in sample_tables.values()),
            },
            "classification": classification,
            "relationships": relationships,
            "bus_matrix_result": bus_matrix_result,
            "method": method,
        }

    def _run_deploy(self, state: DimensionWizardState) -> dict:
        """Create hierarchy projects from the star schema."""
        try:
            from src.hierarchy.dispatch import (
                dispatch_hierarchy_manage,
                dispatch_hierarchy_property,
            )

            # Create a project for the star schema
            project_name = f"{state.company_name or 'Star Schema'} - {state.industry.replace('_', ' ').title()}"
            project_result = dispatch_hierarchy_manage(
                action="create_project",
                name=project_name,
                description=f"Star schema generated by Dimension Builder wizard (session {state.session_id})",
            )

            project_id = project_result.get("project_id", project_result.get("id", state.session_id))
            hierarchy_ids = []

            # Create a hierarchy for each dimension
            for dim in state.dimensions:
                h_result = dispatch_hierarchy_manage(
                    action="create_hierarchy",
                    project_id=project_id,
                    name=dim["name"],
                    hierarchy_type=dim.get("type", "conformed"),
                )
                h_id = h_result.get("hierarchy_id", h_result.get("id", dim["id"]))
                hierarchy_ids.append(h_id)

                # Add levels as nodes
                for level in dim.get("levels", []):
                    dispatch_hierarchy_manage(
                        action="add_node",
                        hierarchy_id=h_id,
                        name=level["name"],
                        level=level.get("level", 1),
                    )

                # Set properties
                props = dim.get("properties", {})
                if props:
                    dispatch_hierarchy_property(
                        action="set_properties",
                        hierarchy_id=h_id,
                        properties=props,
                    )

            return {"project_id": project_id, "hierarchy_ids": hierarchy_ids}

        except ImportError:
            return {"error": "Hierarchy modules not available. Deploy manually from the preview DDL."}
        except Exception as exc:
            logger.warning("Deploy failed: %s", exc)
            return {"error": str(exc)}
