"""
MCP Tools for Hierarchy-Driven Data Mart Factory (Wright Module).

Provides 4 unified MCP tools:
- wright_config: 6 actions (create, add_join_pattern, export, list, validate, version)
- wright_generate: 11 artifact types (pipeline, object, dbt_project, dbt_sources,
                   dbt_tests, dbt_metrics, dbt_ci, dbt_model, schema_yml, suggest_tests, test_queries)
- wright_analyze: 11 check types (discover_pattern, suggest_config, validate_pipeline,
                  hierarchy_quality, normalize_ids, id_source_report, filter_precedence,
                  filter_sql, ddl_compare, compare_baseline, pipeline_health)
- wright_hierarchy: 3 actions (from_hierarchy, sync, run_dbt)
"""

import logging
import re
from dataclasses import dataclass, field
from typing import Any, Dict, List, Optional, Set, Tuple
from enum import Enum

from .types import (
    MartConfig,
    JoinPattern,
    DynamicColumnMapping,
    PipelineObject,
    PipelineLayer,
    FormulaPrecedence,
)
from .config_generator import MartConfigGenerator
from .pipeline_generator import MartPipelineGenerator
from .formula_engine import FormulaPrecedenceEngine, create_standard_los_formulas
from .cortex_discovery import CortexDiscoveryAgent
from .quality_validator import HierarchyQualityValidator, validate_hierarchy_quality
from .alias_normalizer import IDSourceNormalizer, get_normalizer
from .filter_engine import GroupFilterPrecedenceEngine, analyze_group_filter_precedence
from .ddl_diff import DDLDiffComparator, compare_generated_ddl

logger = logging.getLogger(__name__)


# =============================================================================
# Test Suggestion Engine - Sophisticated dbt Test Generation
# =============================================================================


class ColumnCategory(Enum):
    """Categories of columns for test suggestion."""
    PRIMARY_KEY = "primary_key"
    SURROGATE_KEY = "surrogate_key"
    FOREIGN_KEY = "foreign_key"
    NATURAL_KEY = "natural_key"
    DATE = "date"
    DATETIME = "datetime"
    FLAG = "flag"
    CODE = "code"
    AMOUNT = "amount"
    QUANTITY = "quantity"
    PERCENT = "percent"
    NAME = "name"
    DESCRIPTION = "description"
    STATUS = "status"
    TYPE = "type"
    IDENTIFIER = "identifier"
    UNKNOWN = "unknown"


class DataClassification(Enum):
    """Data classification for meta fields."""
    PII = "pii"
    PHI = "phi"
    PCI = "pci"
    CONFIDENTIAL = "confidential"
    INTERNAL = "internal"
    PUBLIC = "public"


@dataclass
class ColumnAnalysis:
    """Analysis result for a single column."""
    name: str
    category: ColumnCategory
    data_type: Optional[str] = None
    is_nullable: bool = True
    is_unique: bool = False
    classification: DataClassification = DataClassification.INTERNAL
    referenced_table: Optional[str] = None
    referenced_column: Optional[str] = None
    accepted_values: List[str] = field(default_factory=list)
    pattern_regex: Optional[str] = None
    min_value: Optional[float] = None
    max_value: Optional[float] = None
    freshness_field: bool = False


@dataclass
class SchemaTest:
    """A schema-level test (in schema.yml)."""
    column: str
    test_type: str
    config: Dict[str, Any] = field(default_factory=dict)
    reason: str = ""


@dataclass
class SingularTest:
    """A singular test (standalone SQL file)."""
    name: str
    sql: str
    description: str
    severity: str = "warn"
    tags: List[str] = field(default_factory=list)


@dataclass
class RelationshipTest:
    """A relationship/referential integrity test."""
    from_column: str
    to_model: str
    to_column: str
    reason: str = ""


class TestSuggestionEngine:
    """
    Sophisticated test suggestion engine for dbt models.

    Analyzes:
    - Column naming patterns (FK_, SK_, _ID, _DATE, _FLAG, _CODE, _AMOUNT)
    - Wright formula groups for calculation tests
    - Referential integrity based on join patterns
    - SQL structure for data quality tests
    """

    # Column pattern definitions
    COLUMN_PATTERNS: Dict[str, Tuple[ColumnCategory, List[str]]] = {
        # Primary/Surrogate Keys
        r'^SK_': (ColumnCategory.SURROGATE_KEY, ['unique', 'not_null']),
        r'^SURROGATE_KEY$': (ColumnCategory.SURROGATE_KEY, ['unique', 'not_null']),
        r'^PK_': (ColumnCategory.PRIMARY_KEY, ['unique', 'not_null']),

        # Foreign Keys
        r'^FK_': (ColumnCategory.FOREIGN_KEY, ['not_null']),
        r'_KEY$': (ColumnCategory.FOREIGN_KEY, ['not_null']),

        # Identifiers
        r'_ID$': (ColumnCategory.IDENTIFIER, ['not_null']),
        r'^ID$': (ColumnCategory.IDENTIFIER, ['not_null']),
        r'_UID$': (ColumnCategory.IDENTIFIER, []),

        # Dates and Times
        r'_DATE$': (ColumnCategory.DATE, []),
        r'_DT$': (ColumnCategory.DATETIME, []),
        r'_DATETIME$': (ColumnCategory.DATETIME, []),
        r'_TIMESTAMP$': (ColumnCategory.DATETIME, []),
        r'^CREATED_AT$': (ColumnCategory.DATETIME, ['not_null']),
        r'^UPDATED_AT$': (ColumnCategory.DATETIME, []),
        r'^LOADED_AT$': (ColumnCategory.DATETIME, ['not_null']),
        r'^VALID_FROM$': (ColumnCategory.DATETIME, ['not_null']),
        r'^VALID_TO$': (ColumnCategory.DATETIME, []),

        # Flags and Booleans
        r'_FLAG$': (ColumnCategory.FLAG, []),
        r'^IS_': (ColumnCategory.FLAG, []),
        r'^HAS_': (ColumnCategory.FLAG, []),
        r'_IND$': (ColumnCategory.FLAG, []),

        # Codes and Types
        r'_CODE$': (ColumnCategory.CODE, ['not_null']),
        r'_TYPE$': (ColumnCategory.TYPE, []),
        r'_STATUS$': (ColumnCategory.STATUS, []),
        r'_CATEGORY$': (ColumnCategory.TYPE, []),

        # Amounts and Quantities
        r'_AMOUNT$': (ColumnCategory.AMOUNT, []),
        r'_AMT$': (ColumnCategory.AMOUNT, []),
        r'_QUANTITY$': (ColumnCategory.QUANTITY, []),
        r'_QTY$': (ColumnCategory.QUANTITY, []),
        r'_COUNT$': (ColumnCategory.QUANTITY, []),
        r'_PERCENT$': (ColumnCategory.PERCENT, []),
        r'_PCT$': (ColumnCategory.PERCENT, []),
        r'_RATE$': (ColumnCategory.PERCENT, []),

        # Names and Descriptions
        r'_NAME$': (ColumnCategory.NAME, []),
        r'_DESC$': (ColumnCategory.DESCRIPTION, []),
        r'_DESCRIPTION$': (ColumnCategory.DESCRIPTION, []),
    }

    # PII column patterns for classification
    PII_PATTERNS = [
        r'EMAIL', r'PHONE', r'SSN', r'SOCIAL_SECURITY', r'DOB', r'BIRTH_DATE',
        r'ADDRESS', r'ZIP', r'POSTAL', r'FIRST_NAME', r'LAST_NAME', r'FULL_NAME',
        r'DRIVER_LICENSE', r'PASSPORT', r'TAX_ID', r'NATIONAL_ID'
    ]

    # PHI patterns
    PHI_PATTERNS = [
        r'DIAGNOSIS', r'TREATMENT', r'MEDICATION', r'MEDICAL', r'HEALTH',
        r'PATIENT', r'PRESCRIPTION', r'INSURANCE_ID', r'MEMBER_ID'
    ]

    # PCI patterns
    PCI_PATTERNS = [
        r'CREDIT_CARD', r'CARD_NUMBER', r'CVV', r'CVC', r'EXPIRY', r'PAN',
        r'ACCOUNT_NUMBER', r'ROUTING_NUMBER', r'BANK_ACCOUNT'
    ]

    # Status value patterns
    STATUS_VALUES = {
        'STATUS': ['ACTIVE', 'INACTIVE', 'PENDING', 'DELETED', 'ARCHIVED'],
        'TYPE': ['PRIMARY', 'SECONDARY', 'TERTIARY'],
        'FLAG': ['Y', 'N', 'TRUE', 'FALSE', '0', '1'],
    }

    def __init__(self, config: Optional[MartConfig] = None):
        """Initialize with optional Wright config for context."""
        self.config = config
        self._column_cache: Dict[str, ColumnAnalysis] = {}

    def analyze_sql(self, sql: str) -> Dict[str, Any]:
        """
        Parse SQL thoroughly to detect patterns.

        Returns:
            Dict with columns, joins, CTEs, aggregations detected
        """
        sql_upper = sql.upper()

        analysis = {
            "columns": [],
            "tables": [],
            "joins": [],
            "ctes": [],
            "aggregations": [],
            "where_clauses": [],
            "group_by_columns": [],
            "has_surrogate_key": False,
            "has_foreign_keys": False,
            "is_incremental": False,
            "has_window_functions": False,
        }

        # Detect CTEs
        cte_pattern = r'WITH\s+(\w+)\s+AS\s*\('
        analysis["ctes"] = re.findall(cte_pattern, sql_upper)

        # Detect SELECT columns - improved parsing
        # First, try to find the main SELECT...FROM block
        # Handle CTEs by finding the last SELECT before FROM
        select_matches = list(re.finditer(
            r'SELECT\s+([\s\S]*?)\s+FROM\s',
            sql_upper,
            re.IGNORECASE
        ))

        columns = []
        if select_matches:
            # Use the last SELECT (main query, not CTEs)
            select_clause = select_matches[-1].group(1)

            # Method 1: Extract explicit AS aliases (most reliable)
            as_cols = re.findall(r'\sAS\s+([A-Z_][A-Z0-9_]*)', select_clause)
            columns.extend(as_cols)

            # Method 2: Extract simple column references (TABLE.COLUMN or COLUMN)
            # Split by comma first, then extract the last identifier
            parts = re.split(r',(?![^(]*\))', select_clause)  # Split by comma, not inside parens
            for part in parts:
                part = part.strip()
                if not part:
                    continue

                # If has AS, we already got it
                if ' AS ' in part:
                    continue

                # Extract last identifier (handles TABLE.COLUMN -> COLUMN)
                # Also handles expressions like SUM(AMOUNT) -> skip function
                if not re.search(r'[A-Z_]+\s*\(', part):  # Not a function
                    match = re.search(r'\.?([A-Z_][A-Z0-9_]*)$', part)
                    if match:
                        columns.append(match.group(1))

            # Remove duplicates and exclude keywords
            keywords = {'DISTINCT', 'ALL', 'TOP', 'CASE', 'WHEN', 'THEN', 'ELSE', 'END', 'AND', 'OR', 'AS'}
            analysis["columns"] = list(set(c for c in columns if c not in keywords))

        # Detect table references
        from_pattern = r'FROM\s+([A-Z_][A-Z0-9_\.]+)'
        analysis["tables"] = re.findall(from_pattern, sql_upper)

        # Detect JOINs
        join_pattern = r'(LEFT|RIGHT|INNER|OUTER|CROSS)?\s*JOIN\s+([A-Z_][A-Z0-9_\.]+)\s+(?:AS\s+)?(\w+)?\s+ON\s+([^WHERE|LEFT|RIGHT|INNER|GROUP|ORDER]+)'
        joins = re.findall(join_pattern, sql_upper, re.DOTALL)
        for join in joins:
            analysis["joins"].append({
                "type": join[0] or "INNER",
                "table": join[1],
                "alias": join[2],
                "condition": join[3].strip(),
            })

        # Detect aggregations
        agg_pattern = r'(SUM|COUNT|AVG|MIN|MAX|LISTAGG)\s*\('
        analysis["aggregations"] = list(set(re.findall(agg_pattern, sql_upper)))

        # Detect GROUP BY
        group_pattern = r'GROUP\s+BY\s+([^ORDER|HAVING|LIMIT|;]+)'
        group_match = re.search(group_pattern, sql_upper)
        if group_match:
            analysis["group_by_columns"] = [
                c.strip() for c in group_match.group(1).split(',')
            ]

        # Detect window functions
        if re.search(r'OVER\s*\(', sql_upper):
            analysis["has_window_functions"] = True

        # Detect incremental pattern
        if 'IS_INCREMENTAL' in sql_upper or 'INCREMENTAL' in sql_upper:
            analysis["is_incremental"] = True

        # Check for surrogate/foreign keys
        analysis["has_surrogate_key"] = bool(
            re.search(r'(SURROGATE_KEY|SK_)', sql_upper)
        )
        analysis["has_foreign_keys"] = bool(
            re.search(r'FK_\w+', sql_upper)
        )

        return analysis

    def analyze_column(self, col_name: str) -> ColumnAnalysis:
        """Analyze a column by its name pattern."""
        if col_name in self._column_cache:
            return self._column_cache[col_name]

        col_upper = col_name.upper()
        category = ColumnCategory.UNKNOWN
        tests = []
        classification = DataClassification.INTERNAL
        accepted_values: List[str] = []
        pattern_regex: Optional[str] = None
        referenced_table: Optional[str] = None
        referenced_column: Optional[str] = None

        # Check column patterns
        for pattern, (cat, default_tests) in self.COLUMN_PATTERNS.items():
            if re.search(pattern, col_upper):
                category = cat
                tests = default_tests
                break

        # Detect data classification
        for pii_pattern in self.PII_PATTERNS:
            if re.search(pii_pattern, col_upper):
                classification = DataClassification.PII
                break
        for phi_pattern in self.PHI_PATTERNS:
            if re.search(phi_pattern, col_upper):
                classification = DataClassification.PHI
                break
        for pci_pattern in self.PCI_PATTERNS:
            if re.search(pci_pattern, col_upper):
                classification = DataClassification.PCI
                break

        # Determine accepted values for known patterns
        if category == ColumnCategory.FLAG:
            accepted_values = ['true', 'false']
        elif category == ColumnCategory.STATUS:
            accepted_values = ['ACTIVE', 'INACTIVE', 'PENDING', 'DELETED']

        # Detect foreign key references
        if category == ColumnCategory.FOREIGN_KEY:
            # FK_ACCOUNT_KEY -> ref('dim_account')
            fk_match = re.match(r'FK_(\w+)_KEY', col_upper)
            if fk_match:
                referenced_table = f"dim_{fk_match.group(1).lower()}"
                referenced_column = f"{fk_match.group(1).lower()}_key"

        # Detect freshness fields
        freshness_field = col_upper in [
            'LOADED_AT', 'UPDATED_AT', 'CREATED_AT', 'ETL_TIMESTAMP',
            'LAST_MODIFIED', 'MODIFIED_DATE', 'LOAD_DATE'
        ]

        analysis = ColumnAnalysis(
            name=col_name,
            category=category,
            classification=classification,
            accepted_values=accepted_values,
            pattern_regex=pattern_regex,
            referenced_table=referenced_table,
            referenced_column=referenced_column,
            freshness_field=freshness_field,
        )

        self._column_cache[col_name] = analysis
        return analysis

    def suggest_schema_tests(
        self,
        columns: List[str],
        sql_analysis: Dict[str, Any],
    ) -> List[SchemaTest]:
        """
        Suggest schema-level tests for columns.

        Returns list of SchemaTest objects for schema.yml generation.
        """
        tests: List[SchemaTest] = []

        for col_name in columns:
            analysis = self.analyze_column(col_name)

            # Basic tests based on category
            if analysis.category == ColumnCategory.SURROGATE_KEY:
                tests.append(SchemaTest(
                    column=col_name,
                    test_type="unique",
                    reason="Surrogate key must be unique"
                ))
                tests.append(SchemaTest(
                    column=col_name,
                    test_type="not_null",
                    reason="Surrogate key must not be null"
                ))

            elif analysis.category == ColumnCategory.PRIMARY_KEY:
                tests.append(SchemaTest(
                    column=col_name,
                    test_type="unique",
                    reason="Primary key must be unique"
                ))
                tests.append(SchemaTest(
                    column=col_name,
                    test_type="not_null",
                    reason="Primary key must not be null"
                ))

            elif analysis.category == ColumnCategory.FOREIGN_KEY:
                tests.append(SchemaTest(
                    column=col_name,
                    test_type="not_null",
                    reason="Foreign key should not be null for referential integrity"
                ))
                # Add relationship test if we can infer the reference
                if analysis.referenced_table:
                    tests.append(SchemaTest(
                        column=col_name,
                        test_type="relationships",
                        config={
                            "to": f"ref('{analysis.referenced_table}')",
                            "field": analysis.referenced_column or col_name.lower(),
                        },
                        reason=f"Referential integrity to {analysis.referenced_table}"
                    ))

            elif analysis.category == ColumnCategory.FLAG:
                tests.append(SchemaTest(
                    column=col_name,
                    test_type="accepted_values",
                    config={"values": analysis.accepted_values or ['true', 'false']},
                    reason="Boolean flag should only contain true/false"
                ))

            elif analysis.category == ColumnCategory.STATUS:
                tests.append(SchemaTest(
                    column=col_name,
                    test_type="accepted_values",
                    config={
                        "values": analysis.accepted_values or
                                  ['ACTIVE', 'INACTIVE', 'PENDING', 'DELETED']
                    },
                    reason="Status column should contain known values"
                ))

            elif analysis.category == ColumnCategory.CODE:
                tests.append(SchemaTest(
                    column=col_name,
                    test_type="not_null",
                    reason="Code columns typically should not be null"
                ))

            elif analysis.category == ColumnCategory.DATE:
                # Date recency test for certain columns
                if any(kw in col_name.upper() for kw in ['LOADED', 'UPDATED', 'CREATED']):
                    tests.append(SchemaTest(
                        column=col_name,
                        test_type="dbt_utils.recency",
                        config={
                            "datepart": "day",
                            "field": col_name.lower(),
                            "interval": 1,
                        },
                        reason="Freshness check - data should be recent"
                    ))

            elif analysis.category == ColumnCategory.PERCENT:
                tests.append(SchemaTest(
                    column=col_name,
                    test_type="dbt_utils.accepted_range",
                    config={"min_value": 0, "max_value": 100},
                    reason="Percentage should be between 0 and 100"
                ))

        return tests

    def suggest_relationship_tests(
        self,
        sql_analysis: Dict[str, Any],
    ) -> List[RelationshipTest]:
        """
        Suggest relationship tests based on JOIN patterns.
        """
        tests: List[RelationshipTest] = []

        for join in sql_analysis.get("joins", []):
            condition = join.get("condition", "")
            table = join.get("table", "")

            # Parse ON condition for column mappings
            # e.g., "a.FK_ACCOUNT_KEY = dim.ACCOUNT_KEY"
            on_match = re.search(
                r'(\w+)\.(\w+)\s*=\s*(\w+)\.(\w+)',
                condition,
                re.IGNORECASE
            )
            if on_match:
                from_col = on_match.group(2)
                to_col = on_match.group(4)

                # Infer model name from table
                model_name = table.split('.')[-1].lower()
                if model_name.startswith('tbl_'):
                    model_name = model_name[4:]

                tests.append(RelationshipTest(
                    from_column=from_col,
                    to_model=model_name,
                    to_column=to_col,
                    reason=f"Referential integrity from JOIN to {table}"
                ))

        return tests

    def suggest_singular_tests(
        self,
        model_name: str,
        sql_analysis: Dict[str, Any],
        config: Optional[MartConfig] = None,
    ) -> List[SingularTest]:
        """
        Suggest singular tests (SQL files) for complex validations.
        """
        tests: List[SingularTest] = []

        # 1. Duplicate check on surrogate key
        if sql_analysis.get("has_surrogate_key"):
            tests.append(SingularTest(
                name=f"assert_{model_name}_no_duplicate_keys",
                sql=f"""-- Assert no duplicate surrogate keys
SELECT
    SURROGATE_KEY,
    COUNT(*) as occurrence_count
FROM {{{{ ref('{model_name}') }}}}
GROUP BY SURROGATE_KEY
HAVING COUNT(*) > 1
""",
                description="Ensure surrogate keys are unique across all records",
                severity="error",
                tags=["data-quality", "keys"]
            ))

        # 2. Foreign key orphan check
        for col in sql_analysis.get("columns", []):
            if col.upper().startswith("FK_"):
                fk_match = re.match(r'FK_(\w+)_KEY', col.upper())
                if fk_match:
                    dim_name = fk_match.group(1).lower()
                    tests.append(SingularTest(
                        name=f"assert_{model_name}_{col.lower()}_valid",
                        sql=f"""-- Assert all foreign keys have matching dimension records
SELECT
    src.{col},
    COUNT(*) as orphan_count
FROM {{{{ ref('{model_name}') }}}} src
LEFT JOIN {{{{ ref('dim_{dim_name}') }}}} dim
    ON src.{col} = dim.{dim_name}_key
WHERE dim.{dim_name}_key IS NULL
  AND src.{col} IS NOT NULL
GROUP BY src.{col}
""",
                        description=f"Ensure all {col} values exist in dim_{dim_name}",
                        severity="warn",
                        tags=["referential-integrity", "data-quality"]
                    ))

        # 3. Date sanity checks
        date_cols = [
            c for c in sql_analysis.get("columns", [])
            if any(d in c.upper() for d in ['_DATE', '_DT', '_TIMESTAMP'])
        ]
        if date_cols:
            col = date_cols[0]
            tests.append(SingularTest(
                name=f"assert_{model_name}_date_sanity",
                sql=f"""-- Assert dates are within reasonable range
SELECT
    '{col}' as column_name,
    MIN({col}) as min_date,
    MAX({col}) as max_date,
    COUNT(CASE WHEN {col} > CURRENT_DATE() THEN 1 END) as future_dates,
    COUNT(CASE WHEN {col} < '1900-01-01' THEN 1 END) as ancient_dates
FROM {{{{ ref('{model_name}') }}}}
HAVING future_dates > 0 OR ancient_dates > 0
""",
                description="Ensure dates are within reasonable historical range",
                severity="warn",
                tags=["data-quality", "dates"]
            ))

        # 4. Amount/quantity non-negative check
        amount_cols = [
            c for c in sql_analysis.get("columns", [])
            if any(a in c.upper() for a in ['_AMOUNT', '_AMT', '_QTY', '_QUANTITY', '_COUNT'])
        ]
        if amount_cols:
            col = amount_cols[0]
            tests.append(SingularTest(
                name=f"assert_{model_name}_amounts_valid",
                sql=f"""-- Check for unexpected negative amounts
SELECT
    '{col}' as column_name,
    COUNT(*) as negative_count,
    SUM({col}) as negative_total
FROM {{{{ ref('{model_name}') }}}}
WHERE {col} < 0
HAVING negative_count > 0
""",
                description="Identify unexpected negative values in amount columns",
                severity="warn",
                tags=["data-quality", "amounts"]
            ))

        # 5. Wright formula validation (if config provided)
        if config:
            measure_prefix = config.effective_measure_prefix
            tests.append(SingularTest(
                name=f"assert_{model_name}_formula_integrity",
                sql=f"""-- Validate Wright formula calculations
WITH base_checks AS (
    SELECT
        COUNT(*) as total_rows,
        COUNT(SURROGATE_KEY) as non_null_keys,
        COUNT(DISTINCT SURROGATE_KEY) as unique_keys,
        SUM(CASE WHEN {measure_prefix}AMOUNT IS NULL THEN 1 ELSE 0 END) as null_amounts
    FROM {{{{ ref('{model_name}') }}}}
)
SELECT
    *,
    CASE
        WHEN total_rows != non_null_keys THEN 'MISSING_KEYS'
        WHEN total_rows != unique_keys THEN 'DUPLICATE_KEYS'
        WHEN null_amounts > total_rows * 0.5 THEN 'EXCESSIVE_NULL_AMOUNTS'
        ELSE 'OK'
    END as validation_result
FROM base_checks
WHERE total_rows != non_null_keys
   OR total_rows != unique_keys
   OR null_amounts > total_rows * 0.5
""",
                description="Validate Wright pipeline formula calculations and key integrity",
                severity="error",
                tags=["wright", "formulas", "data-quality"]
            ))

            # 6. Formula group balance check
            tests.append(SingularTest(
                name=f"assert_{model_name}_formula_group_balance",
                sql=f"""-- Validate formula group calculations balance
WITH group_totals AS (
    SELECT
        FORMULA_GROUP,
        SUM({measure_prefix}AMOUNT) as group_total
    FROM {{{{ ref('{model_name}') }}}}
    WHERE FORMULA_GROUP IS NOT NULL
    GROUP BY FORMULA_GROUP
),
calculated_rows AS (
    SELECT
        HIERARCHY_NAME,
        {measure_prefix}AMOUNT as calculated_amount,
        FORMULA_GROUP
    FROM {{{{ ref('{model_name}') }}}}
    WHERE PRECEDENCE_LEVEL IS NOT NULL
)
SELECT
    c.HIERARCHY_NAME,
    c.calculated_amount,
    c.FORMULA_GROUP,
    g.group_total,
    ABS(c.calculated_amount - g.group_total) as variance
FROM calculated_rows c
JOIN group_totals g ON c.FORMULA_GROUP = g.FORMULA_GROUP
WHERE ABS(c.calculated_amount - g.group_total) > 0.01
""",
                description="Ensure formula group calculated rows match component sums",
                severity="warn",
                tags=["wright", "formulas", "calculations"]
            ))

        # 7. Row count trend check (for incremental models)
        if sql_analysis.get("is_incremental"):
            tests.append(SingularTest(
                name=f"assert_{model_name}_row_count_trend",
                sql=f"""-- Check for unexpected row count changes
{{%- set old_count = 1000 -%}}  -- Set baseline
SELECT
    (SELECT COUNT(*) FROM {{{{ ref('{model_name}') }}}}) as current_count,
    {{{{ old_count }}}} as baseline_count,
    ABS((SELECT COUNT(*) FROM {{{{ ref('{model_name}') }}}}) - {{{{ old_count }}}}) / NULLIF({{{{ old_count }}}}, 0) * 100 as pct_change
WHERE ABS((SELECT COUNT(*) FROM {{{{ ref('{model_name}') }}}}) - {{{{ old_count }}}}) / NULLIF({{{{ old_count }}}}, 0) > 0.5
""",
                description="Alert on significant row count changes (>50%)",
                severity="warn",
                tags=["trend", "monitoring"]
            ))

        return tests

    def suggest_formula_tests(
        self,
        model_name: str,
        formulas: List[FormulaPrecedence],
    ) -> List[SingularTest]:
        """
        Generate tests based on Wright formula definitions.
        """
        tests: List[SingularTest] = []

        for formula in formulas:
            # Build the expected calculation SQL
            if formula.logic.value == "SUM":
                expected = f"SUM({formula.param_ref})"
            elif formula.logic.value == "SUBTRACT":
                expected = f"({formula.param_ref}) - ({formula.param2_ref})"
            else:
                continue

            tests.append(SingularTest(
                name=f"assert_{model_name}_formula_{formula.formula_group.lower().replace(' ', '_')}",
                sql=f"""-- Validate {formula.formula_group} calculation
WITH expected AS (
    SELECT
        {expected} as expected_value
    FROM {{{{ ref('{model_name}') }}}}
),
actual AS (
    SELECT
        SUM(AMOUNT) as actual_value
    FROM {{{{ ref('{model_name}') }}}}
    WHERE FORMULA_GROUP = '{formula.formula_group}'
)
SELECT
    e.expected_value,
    a.actual_value,
    ABS(e.expected_value - a.actual_value) as variance
FROM expected e, actual a
WHERE ABS(e.expected_value - a.actual_value) > 0.01
""",
                description=f"Validate {formula.formula_group} formula (P{formula.precedence_level})",
                severity="error",
                tags=["wright", "formula", f"p{formula.precedence_level}"]
            ))

        return tests

    def generate_all_tests(
        self,
        model_name: str,
        sql: str,
        config: Optional[MartConfig] = None,
        formulas: Optional[List[FormulaPrecedence]] = None,
    ) -> Dict[str, Any]:
        """
        Generate all test suggestions for a model.

        Returns:
            Dict with schema_tests, singular_tests, relationship_tests
        """
        sql_analysis = self.analyze_sql(sql)
        columns = sql_analysis.get("columns", [])

        schema_tests = self.suggest_schema_tests(columns, sql_analysis)
        relationship_tests = self.suggest_relationship_tests(sql_analysis)
        singular_tests = self.suggest_singular_tests(
            model_name, sql_analysis, config
        )

        # Add formula tests if formulas provided
        if formulas:
            singular_tests.extend(
                self.suggest_formula_tests(model_name, formulas)
            )

        return {
            "schema_tests": schema_tests,
            "singular_tests": singular_tests,
            "relationship_tests": relationship_tests,
            "sql_analysis": sql_analysis,
        }


class SchemaYmlGenerator:
    """
    Enhanced schema.yml generator with sophisticated column analysis.
    """

    def __init__(self, test_engine: TestSuggestionEngine):
        self.test_engine = test_engine

    def generate_column_description(self, col_name: str) -> str:
        """Generate rich description based on column patterns."""
        analysis = self.test_engine.analyze_column(col_name)
        col_upper = col_name.upper()

        # Build description based on category
        descriptions = {
            ColumnCategory.SURROGATE_KEY: "Unique surrogate key generated for this record. Used as the primary identifier for joins and lookups.",
            ColumnCategory.PRIMARY_KEY: f"Primary key identifier for the {col_upper.replace('PK_', '').replace('_', ' ').lower()} entity.",
            ColumnCategory.FOREIGN_KEY: f"Foreign key linking to the {col_upper.replace('FK_', '').replace('_KEY', '').lower()} dimension table.",
            ColumnCategory.IDENTIFIER: f"Business identifier for {col_upper.replace('_ID', '').replace('_', ' ').lower()}.",
            ColumnCategory.DATE: f"Date value representing {col_upper.replace('_DATE', '').replace('_', ' ').lower()}.",
            ColumnCategory.DATETIME: f"Timestamp recording when {col_upper.replace('_DT', '').replace('_DATETIME', '').replace('_TIMESTAMP', '').replace('_', ' ').lower()} occurred.",
            ColumnCategory.FLAG: f"Boolean indicator for {col_upper.replace('_FLAG', '').replace('IS_', '').replace('HAS_', '').replace('_', ' ').lower()}. Values: true/false.",
            ColumnCategory.CODE: f"Code value representing {col_upper.replace('_CODE', '').replace('_', ' ').lower()}.",
            ColumnCategory.STATUS: f"Current status of the record. Expected values: ACTIVE, INACTIVE, PENDING, etc.",
            ColumnCategory.TYPE: f"Type classification for {col_upper.replace('_TYPE', '').replace('_CATEGORY', '').replace('_', ' ').lower()}.",
            ColumnCategory.AMOUNT: f"Monetary amount for {col_upper.replace('_AMOUNT', '').replace('_AMT', '').replace('_', ' ').lower()}. Currency: USD.",
            ColumnCategory.QUANTITY: f"Quantity/count of {col_upper.replace('_QUANTITY', '').replace('_QTY', '').replace('_COUNT', '').replace('_', ' ').lower()}.",
            ColumnCategory.PERCENT: f"Percentage value for {col_upper.replace('_PERCENT', '').replace('_PCT', '').replace('_RATE', '').replace('_', ' ').lower()}. Range: 0-100.",
            ColumnCategory.NAME: f"Display name for {col_upper.replace('_NAME', '').replace('_', ' ').lower()}.",
            ColumnCategory.DESCRIPTION: f"Descriptive text for {col_upper.replace('_DESC', '').replace('_DESCRIPTION', '').replace('_', ' ').lower()}.",
        }

        base_desc = descriptions.get(
            analysis.category,
            col_upper.replace('_', ' ').title()
        )

        # Add classification note if sensitive
        if analysis.classification == DataClassification.PII:
            base_desc += " **Contains PII - handle with care.**"
        elif analysis.classification == DataClassification.PHI:
            base_desc += " **Contains PHI - HIPAA compliance required.**"
        elif analysis.classification == DataClassification.PCI:
            base_desc += " **Contains PCI data - PCI-DSS compliance required.**"

        return base_desc

    def generate_meta_fields(self, col_name: str) -> Dict[str, Any]:
        """Generate meta fields for data classification and ownership."""
        analysis = self.test_engine.analyze_column(col_name)

        meta = {
            "contains_pii": analysis.classification == DataClassification.PII,
        }

        # Add classification
        if analysis.classification != DataClassification.INTERNAL:
            meta["data_classification"] = analysis.classification.value

        # Add ownership hints based on patterns
        col_upper = col_name.upper()
        if any(f in col_upper for f in ['FINANCE', 'REVENUE', 'COST', 'AMOUNT']):
            meta["domain"] = "finance"
        elif any(f in col_upper for f in ['CUSTOMER', 'CLIENT', 'USER']):
            meta["domain"] = "customer"
        elif any(f in col_upper for f in ['PRODUCT', 'SKU', 'ITEM']):
            meta["domain"] = "product"
        elif any(f in col_upper for f in ['ORDER', 'TRANSACTION', 'INVOICE']):
            meta["domain"] = "orders"

        return meta

    def generate_freshness_config(
        self,
        columns: List[str],
        default_hours: int = 24,
    ) -> Optional[Dict[str, Any]]:
        """
        Generate freshness configuration if date columns detected.
        """
        # Look for load/update timestamp columns
        freshness_cols = [
            c for c in columns
            if any(f in c.upper() for f in [
                'LOADED_AT', 'UPDATED_AT', 'CREATED_AT', 'ETL_TIMESTAMP',
                'LAST_MODIFIED', 'MODIFIED_DATE', 'LOAD_DATE'
            ])
        ]

        if freshness_cols:
            return {
                "loaded_at_field": freshness_cols[0].lower(),
                "warn_after": {"count": default_hours, "period": "hour"},
                "error_after": {"count": default_hours * 2, "period": "hour"},
            }
        return None

    def generate_schema_yml(
        self,
        model_name: str,
        columns: List[str],
        sql_analysis: Dict[str, Any],
        schema_tests: List[SchemaTest],
        config: Optional[MartConfig] = None,
        include_freshness: bool = True,
    ) -> str:
        """
        Generate complete schema.yml content.
        """
        # Build column definitions
        column_yaml_parts = []
        for col in columns[:50]:  # Limit to 50 columns
            desc = self.generate_column_description(col)
            meta = self.generate_meta_fields(col)

            # Find tests for this column
            col_tests = [t for t in schema_tests if t.column == col]

            col_entry = f"      - name: {col}\n"
            col_entry += f'        description: "{desc}"\n'

            # Add meta if present
            if meta:
                col_entry += "        meta:\n"
                for key, value in meta.items():
                    if isinstance(value, bool):
                        col_entry += f"          {key}: {str(value).lower()}\n"
                    else:
                        col_entry += f'          {key}: "{value}"\n'

            # Add tests
            if col_tests:
                col_entry += "        tests:\n"
                for test in col_tests:
                    if test.config:
                        col_entry += f"          - {test.test_type}:\n"
                        for key, value in test.config.items():
                            if isinstance(value, list):
                                col_entry += f"              {key}: {value}\n"
                            else:
                                col_entry += f"              {key}: {value}\n"
                    else:
                        col_entry += f"          - {test.test_type}\n"

            column_yaml_parts.append(col_entry)

        # Build model description
        config_desc = ""
        if config:
            config_desc = f" for {config.report_type} {config.project_name} reporting"

        # Build freshness config
        freshness_yml = ""
        if include_freshness:
            freshness = self.generate_freshness_config(columns)
            if freshness:
                freshness_yml = f"""
    freshness:
      loaded_at_field: {freshness['loaded_at_field']}
      warn_after:
        count: {freshness['warn_after']['count']}
        period: {freshness['warn_after']['period']}
      error_after:
        count: {freshness['error_after']['count']}
        period: {freshness['error_after']['period']}
"""

        # Assemble full YAML
        schema_yml = f'''version: 2

models:
  - name: {model_name}
    description: "Wright pipeline model{config_desc}. Auto-generated schema documentation."
    config:
      tags: ['wright', 'data-mart']
      materialized: table{freshness_yml}
    meta:
      owner: "data-engineering"
      sla: "daily"
      data_quality_tier: "gold"
    columns:
{chr(10).join(column_yaml_parts)}
'''

        return schema_yml



def register_mart_factory_tools(mcp, settings=None) -> Dict[str, Any]:
    """
    Register Mart Factory MCP tools (unified only).

    Registers 4 unified tools.

    Args:
        mcp: The FastMCP instance
        settings: Optional settings

    Returns:
        Dict with registration info
    """
    from .unified import (
        dispatch_wright_config,
        dispatch_wright_generate,
        dispatch_wright_analyze,
        dispatch_wright_hierarchy,
        register_unified_wright_tools,
    )

    # Register 4 unified tools
    register_unified_wright_tools(mcp, settings)

    logger.info("Registered 4 Wright tools (4 unified)")

    return {
        "tools_registered": 4,
        "unified_tools": ["wright_config", "wright_generate", "wright_analyze", "wright_hierarchy"],
    }
