"""Data quality + validation specialist agent."""

from __future__ import annotations

import logging
from typing import Any, Dict

from .base_agent import BaseDataBridgeAgent
from .state_schema import WorkflowState
from .databridge_tools import tool_run_quality_check

logger = logging.getLogger(__name__)


class QAAgent(BaseDataBridgeAgent):
    """Specialist agent for data quality validation and observability."""

    name = "qa_agent"
    description = "Runs data quality validations and sets up observability baselines"
    phase_name = "quality"

    def __init__(self, **kwargs):
        super().__init__(
            tools=[tool_run_quality_check],
            system_prompt=(
                "You are the Quality Assurance agent for DataBridge AI. "
                "You generate expectation suites from column classifications, "
                "run validations, and set up observability baselines."
            ),
            **kwargs,
        )

    async def run(self, state: WorkflowState) -> WorkflowState:
        """Run quality checks based on classification data."""
        phase = state.get("current_phase", self.phase_name)
        config = state.get("config", {})
        ctx = state.get("context", {})

        if phase == "observability":
            return await self._run_observability(state, config)

        return await self._run_quality(state, config)

    async def _run_quality(self, state: WorkflowState, config: Dict) -> WorkflowState:
        """Phase: generate and run quality expectations."""
        ctx = state.get("context", {})
        dim_result = ctx.get("detect_dimensions", {})
        classification = dim_result.get("classification", {})

        if not classification:
            return self._update_context(state, "quality", {
                "skipped": True, "reason": "no classification data from detect_dimensions"
            })

        try:
            from src.data_quality.classification_bridge import ClassificationExpectationBridge

            bridge = ClassificationExpectationBridge()
            suites_info = []
            for table, cls_type in classification.items():
                col_map = {}
                if cls_type == "dimension":
                    col_map[f"{table}_KEY"] = "identifier"
                elif cls_type == "fact":
                    col_map[f"{table}_AMOUNT"] = "measure"
                if col_map:
                    suite = bridge.generate_suite(
                        suite_name=f"auto_{table}",
                        classifications=col_map,
                        table_name=table,
                        database=config.get("database", ""),
                        schema_name=config.get("schema", ""),
                    )
                    suites_info.append({
                        "table": table,
                        "expectations": len(suite.expectations),
                    })

            return self._update_context(state, "quality", {
                "suites_generated": len(suites_info),
                "suites": suites_info,
            })
        except Exception as e:
            return self._update_context(state, "quality", {"error": str(e)})

    async def _run_observability(self, state: WorkflowState, config: Dict) -> WorkflowState:
        """Phase: set up observability baselines."""
        return self._update_context(state, "observability", {
            "status": "placeholder",
            "note": "Use obs_record_metric via MCP for observability setup",
        })
