"""Bridge from DataShield column classifications to data quality expectations.

Maps ColumnClassification values to appropriate ExpectationType rules
so that classified columns automatically get relevant quality checks.
"""

from __future__ import annotations

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

from src.datashield.types import ColumnClassification
from .types import (
    Expectation,
    ExpectationSuite,
    ExpectationType,
    SeverityLevel,
)
from .suite_generator import ExpectationSuiteGenerator

logger = logging.getLogger(__name__)

# -----------------------------------------------------------------------
# Classification → Expectation mapping
# -----------------------------------------------------------------------

CLASSIFICATION_EXPECTATIONS: Dict[str, List[Dict[str, Any]]] = {
    ColumnClassification.IDENTIFIER: [
        {"type": ExpectationType.UNIQUE, "severity": SeverityLevel.HIGH},
        {"type": ExpectationType.NOT_NULL, "severity": SeverityLevel.CRITICAL},
    ],
    ColumnClassification.SENSITIVE_PII: [
        {"type": ExpectationType.NOT_NULL, "severity": SeverityLevel.HIGH},
        {"type": ExpectationType.MATCH_REGEX, "severity": SeverityLevel.MEDIUM,
         "kwargs": {"regex": ".+"}, "description": "PII columns must have a value format"},
    ],
    ColumnClassification.TEMPORAL: [
        {"type": ExpectationType.NOT_NULL, "severity": SeverityLevel.HIGH},
        {"type": ExpectationType.BETWEEN, "severity": SeverityLevel.MEDIUM,
         "kwargs": {"min_value": "1900-01-01", "max_value": "2099-12-31"},
         "description": "Date within reasonable range"},
    ],
    ColumnClassification.MEASURE: [
        {"type": ExpectationType.NOT_NULL, "severity": SeverityLevel.MEDIUM},
        {"type": ExpectationType.IN_TYPE_LIST, "severity": SeverityLevel.MEDIUM,
         "kwargs": {"type_list": ["int", "float", "int64", "float64", "Decimal"]},
         "description": "Measure must be numeric"},
    ],
    ColumnClassification.DESCRIPTIVE: [
        {"type": ExpectationType.NOT_NULL, "severity": SeverityLevel.LOW},
    ],
    ColumnClassification.CODE: [
        {"type": ExpectationType.NOT_NULL, "severity": SeverityLevel.HIGH},
        {"type": ExpectationType.IN_SET, "severity": SeverityLevel.MEDIUM,
         "kwargs": {"value_set": []}, "description": "Code column — populate allowed values from data"},
    ],
    ColumnClassification.FACT_DIMENSION: [
        {"type": ExpectationType.NOT_NULL, "severity": SeverityLevel.MEDIUM},
    ],
    ColumnClassification.GEOGRAPHIC: [
        {"type": ExpectationType.NOT_NULL, "severity": SeverityLevel.LOW},
    ],
    ColumnClassification.SAFE: [],
}


class ClassificationExpectationBridge:
    """Generate expectations from DataShield column classifications."""

    def __init__(self, generator: Optional[ExpectationSuiteGenerator] = None):
        self.generator = generator or ExpectationSuiteGenerator()

    def generate_suite(
        self,
        suite_name: str,
        classifications: Dict[str, str],
        table_name: Optional[str] = None,
        database: Optional[str] = None,
        schema_name: Optional[str] = None,
    ) -> ExpectationSuite:
        """Create an expectation suite from column classifications.

        Args:
            suite_name: Name for the suite.
            classifications: Dict mapping column_name -> classification string
                (e.g. {"WELL_ID": "identifier", "AMOUNT": "measure"}).
            table_name: Target table.
            database: Target database.
            schema_name: Target schema.

        Returns:
            Populated ExpectationSuite.
        """
        suite = self.generator.get_suite(suite_name)
        if not suite:
            suite = self.generator.create_suite(
                name=suite_name,
                description=f"Auto-generated from DataShield classifications for {table_name or 'unknown'}",
                database=database,
                schema_name=schema_name,
                table_name=table_name,
            )

        for col_name, classification_str in classifications.items():
            try:
                cls_enum = ColumnClassification(classification_str)
            except ValueError:
                logger.warning("Unknown classification '%s' for column %s — skipping", classification_str, col_name)
                continue

            templates = CLASSIFICATION_EXPECTATIONS.get(cls_enum, [])
            for tmpl in templates:
                exp = Expectation(
                    expectation_type=tmpl["type"],
                    column=col_name,
                    severity=tmpl.get("severity", SeverityLevel.MEDIUM),
                    description=tmpl.get("description", f"{cls_enum.value} → {tmpl['type'].value}"),
                    kwargs=tmpl.get("kwargs", {}),
                )
                suite.add_expectation(exp)

        return suite
