"""Schema discovery + relationship detection agent.

Handles the load_metadata, relationship_discovery, group_tables,
classify, detect_dimensions, and artifact_bundle phases.
"""

from __future__ import annotations

import json
import logging
from typing import Any, Dict

from .base_agent import BaseDataBridgeAgent
from .state_schema import WorkflowState
from .databridge_tools import (
    tool_load_metadata,
    tool_discover_relationships,
    tool_classify_columns,
    tool_detect_dimensions,
    tool_build_artifact,
)

logger = logging.getLogger(__name__)


class DataModelingAgent(BaseDataBridgeAgent):
    """Specialist agent for schema discovery and data modeling."""

    name = "data_modeling_agent"
    description = "Discovers schemas, classifies columns, detects dimensions, and builds artifacts"
    phase_name = "load_metadata"

    def __init__(self, **kwargs):
        super().__init__(
            tools=[
                tool_load_metadata,
                tool_discover_relationships,
                tool_classify_columns,
                tool_detect_dimensions,
                tool_build_artifact,
            ],
            system_prompt=(
                "You are the Data Modeling agent for DataBridge AI. "
                "You handle schema discovery, relationship detection, "
                "column classification, dimension detection, and artifact generation."
            ),
            **kwargs,
        )

    async def run(self, state: WorkflowState) -> WorkflowState:
        """Execute the data modeling phase appropriate to current_phase."""
        phase = state.get("current_phase", self.phase_name)
        config = state.get("config", {})

        if phase == "load_metadata":
            return await self._run_load_metadata(state, config)
        elif phase == "relationship_discovery":
            return await self._run_relationship_discovery(state, config)
        elif phase == "ai_classify":
            return await self._run_classify(state, config)
        elif phase == "detect_dimensions":
            return await self._run_detect_dimensions(state, config)
        elif phase == "artifact_bundle":
            return await self._run_artifact_bundle(state, config)
        else:
            return await self._run_load_metadata(state, config)

    async def _run_load_metadata(self, state: WorkflowState, config: Dict) -> WorkflowState:
        """Phase: load table metadata and samples."""
        try:
            from databridge_data_modeling.mcp_adapter import _discover_data_model

            result = _discover_data_model(
                source_type=config.get("source_type", "csv"),
                csv_folder=config.get("csv_folder"),
                database=config.get("database", ""),
                schema=config.get("schema", ""),
                tables=config.get("tables", "*"),
                sample_limit=config.get("sample_limit", 200),
                min_overlap=config.get("min_overlap", 0.5),
                persist=True,
                snowflake_connection=config.get("snowflake_connection"),
                return_samples=True,
            )
            return self._update_context(state, "load_metadata", result)
        except Exception as e:
            logger.error("load_metadata failed: %s", e)
            return self._update_context(state, "load_metadata", {"error": str(e)})

    async def _run_relationship_discovery(self, state: WorkflowState, config: Dict) -> WorkflowState:
        """Phase: discover FK relationships by column name inference."""
        ctx = state.get("context", {})
        meta = ctx.get("load_metadata", {})
        sample_tables = meta.get("sample_tables", {})

        table_columns = {
            t: data.get("columns", [])
            for t, data in sample_tables.items()
        } if sample_tables else {}

        if not table_columns:
            col_meta = meta.get("column_metadata", [])
            for cm in col_meta:
                tbl = cm.get("TABLE_NAME", "")
                col = cm.get("COLUMN_NAME", "")
                if tbl and col:
                    table_columns.setdefault(tbl, []).append(col)

        if not table_columns:
            return self._update_context(state, "relationship_discovery", {
                "skipped": True, "reason": "no table columns available"
            })

        try:
            from src.data_modeling.column_fk_inferer import infer_relationships_by_name
            from src.data_modeling.erp_configs.registry import get_erp_config

            cfg = None
            erp_entry = get_erp_config(config.get("erp", "enertia"))
            if erp_entry:
                cfg = erp_entry.inference

            rels = infer_relationships_by_name(table_columns, config=cfg)
            return self._update_context(state, "relationship_discovery", {
                "relationships": rels, "count": len(rels)
            })
        except Exception as e:
            return self._update_context(state, "relationship_discovery", {"error": str(e)})

    async def _run_classify(self, state: WorkflowState, config: Dict) -> WorkflowState:
        """Phase: classify columns using DataShield auto_classify_columns."""
        ctx = state.get("context", {})
        meta = ctx.get("load_metadata", {})
        sample_tables = meta.get("sample_tables", {})

        if not sample_tables:
            return self._update_context(state, "ai_classify", {
                "skipped": True, "reason": "no sample_tables"
            })

        try:
            from collections import Counter
            from src.datashield.classifier import auto_classify_columns

            col_meta = meta.get("column_metadata", [])
            col_meta_by_table = {}
            for cm in col_meta:
                tbl = cm["TABLE_NAME"]
                col_meta_by_table.setdefault(tbl, []).append({
                    "name": cm["COLUMN_NAME"],
                    "data_type": cm.get("DATA_TYPE", "VARCHAR"),
                })

            classifications = {}
            counts = Counter()
            for table, data in sample_tables.items():
                cols = col_meta_by_table.get(table, [
                    {"name": c, "data_type": "VARCHAR"} for c in data.get("columns", [])
                ])
                rows = data.get("rows", [])
                sample_vals = {}
                for row in rows:
                    for col, val in row.items():
                        sample_vals.setdefault(col, []).append(val)

                rules = auto_classify_columns(cols, sample_data=sample_vals, row_count=len(rows))
                classifications[table] = [r.model_dump() for r in rules]
                for rule in rules:
                    counts[rule.classification.value] += 1

            result = {
                "classifications": classifications,
                "classification_counts": dict(counts),
                "tables_classified": len(classifications),
            }
            return self._update_context(state, "ai_classify", result)
        except Exception as e:
            return self._update_context(state, "ai_classify", {"error": str(e)})

    async def _run_detect_dimensions(self, state: WorkflowState, config: Dict) -> WorkflowState:
        """Phase: detect dimension vs fact tables."""
        ctx = state.get("context", {})
        meta = ctx.get("load_metadata", {})
        rels = meta.get("relationships", [])

        # Also check relationship_discovery phase
        rel_disc = ctx.get("relationship_discovery", {})
        if not rels and rel_disc.get("relationships"):
            rels = rel_disc["relationships"]

        threshold = config.get("dimension_threshold", 3)

        if rels:
            try:
                from src.data_modeling.dimension_detector import DimensionDetector

                detector = DimensionDetector()
                classification = detector.classify_tables(rels, threshold=threshold)
                return self._update_context(state, "detect_dimensions", {
                    "classification": classification,
                    "relationship_count": len(rels),
                    "method": "dimension_detector",
                })
            except Exception as e:
                return self._update_context(state, "detect_dimensions", {"error": str(e)})

        return self._update_context(state, "detect_dimensions", {
            "skipped": True, "reason": "no relationships available"
        })

    async def _run_artifact_bundle(self, state: WorkflowState, config: Dict) -> WorkflowState:
        """Phase: generate artifact bundle."""
        try:
            from pathlib import Path
            from src.artifacts.bundle import ArtifactBundle

            run_id = state.get("run_id", "unknown")
            output_dir = config.get("output_dir", "data/workflow_runs")
            bundle = ArtifactBundle(
                run_id=run_id,
                output_dir=Path(output_dir) / "artifacts",
            )
            # Add run summary
            summary = {
                "run_id": run_id,
                "status": state.get("status", "unknown"),
                "phases": state.get("completed_phases", []),
                "config": config,
            }
            bundle.add("run_summary", summary, fmt="json")
            manifest = bundle.generate_manifest()
            return self._update_context(state, "artifact_bundle", {
                "manifest": str(manifest),
                "artifacts": len(bundle.artifacts),
            })
        except Exception as e:
            return self._update_context(state, "artifact_bundle", {"error": str(e)})
