import os
import re
import logging
import yaml
from typing import List, Dict, Any, Optional
from app.core.config import settings

logger = logging.getLogger(__name__)


def infer_data_type(col_name: str) -> str:
    """Infer standard Lakehouse data type for dbt contract enforcement."""
    name = col_name.lower().strip()
    if name.endswith("_date") or name == "date":
        return "date"
    if name.endswith("_at") or "timestamp" in name or name.endswith("_time"):
        return "timestamp with time zone"
    if any(k in name for k in ["score", "rate", "amount", "percent", "price", "val", "total", "count"]):
        return "numeric"
    if name.endswith("_id") or name.endswith("_number") or name == "id" or "code" in name or "status" in name or "type" in name:
        return "character varying(50)"
    return "character varying(255)"


def generate_model_schema_dict(
    model_name: str,
    test_configs: Optional[Dict[str, Any]] = None,
    sql_query: Optional[str] = None
) -> Dict[str, Any]:
    """Generate dbt schema dictionary for a model, configuring tests and contract enforcement."""
    clean_name = model_name.replace(".sql", "")
    cfg = test_configs or {}
    
    unique_cols = cfg.get("unique_columns") or cfg.get("unique_cols") or []
    if isinstance(unique_cols, str):
        unique_cols = [c.strip() for c in unique_cols.split(",") if c.strip()]
        
    not_null_cols = cfg.get("not_null_columns") or cfg.get("not_null_cols") or []
    if isinstance(not_null_cols, str):
        not_null_cols = [c.strip() for c in not_null_cols.split(",") if c.strip()]
        
    accepted_values = cfg.get("accepted_values") or {}
    relationships = cfg.get("relationships") or {}
    enforce_contract = cfg.get("enforce_contract", False)
    column_types = cfg.get("column_types") or {}

    # Collect all configured columns
    all_cols = set(unique_cols) | set(not_null_cols) | set(accepted_values.keys()) | set(relationships.keys()) | set(column_types.keys())
    
    # If contract is enforced, extract all output columns from SQL query to satisfy dbt contract rules
    effective_sql = sql_query or cfg.get("sql_query") or cfg.get("query") or ""
    if enforce_contract and effective_sql:
        m = re.search(r'SELECT\s+(.*?)\s+FROM', effective_sql, re.IGNORECASE | re.DOTALL)
        if m:
            for raw_col in m.group(1).split(","):
                raw_col = raw_col.strip()
                if not raw_col:
                    continue
                if " as " in raw_col.lower():
                    col_id = re.split(r'\s+as\s+', raw_col, flags=re.IGNORECASE)[-1].strip().strip('"').strip("'")
                else:
                    col_id = raw_col.split(".")[-1].strip().strip('"').strip("'")
                if col_id and re.match(r'^[a-zA-Z_][a-zA-Z0-9_]*$', col_id):
                    all_cols.add(col_id)

    columns_list = []
    for col in sorted(all_cols):
        col_tests = []
        if col in unique_cols:
            col_tests.append("unique")
        if col in not_null_cols:
            col_tests.append("not_null")
        if col in accepted_values and accepted_values[col]:
            vals = accepted_values[col]
            if isinstance(vals, str):
                vals = [v.strip() for v in vals.split(",") if v.strip()]
            col_tests.append({"accepted_values": {"values": vals}})
        if col in relationships and relationships[col]:
            rel_info = relationships[col]
            if isinstance(rel_info, dict):
                col_tests.append({"relationships": {
                    "to": rel_info.get("to", "ref('upstream')"),
                    "field": rel_info.get("field", "id")
                }})

        col_dict: Dict[str, Any] = {"name": col}
        if col_tests:
            col_dict["tests"] = col_tests
        if col in column_types:
            col_dict["data_type"] = column_types[col]
        elif enforce_contract:
            col_dict["data_type"] = infer_data_type(col)
        columns_list.append(col_dict)

    model_entry: Dict[str, Any] = {
        "name": clean_name,
        "description": f"dbt Model {clean_name} auto-registered with Shift-Left quality assertions"
    }
    
    if enforce_contract:
        model_entry["config"] = {"contract": {"enforced": True}}
    if columns_list:
        model_entry["columns"] = columns_list

    return model_entry


def compile_dbt_model(
    model_name: str,
    layer: str = "silver",
    sql_query: str = "",
    test_configs: Optional[Dict[str, Any]] = None
) -> Dict[str, Any]:
    """
    Perform pre-flight compilation check for dbt models:
    - Verifies Jinja delimiters pairing.
    - Extracts {{ ref('...') }} and {{ source('...', '...') }}.
    - Compiles Jinja refs into simulated physical lakehouse identifiers.
    - Validates SQL structure (SELECT/WITH statements, parentheses matching).
    - Generates corresponding dbt schema.yml definition with quality tests.
    """
    clean_name = (model_name or "model").replace(".sql", "")
    errors: List[str] = []
    warnings: List[str] = []
    
    if not sql_query or not sql_query.strip():
        errors.append("SQL query cannot be empty.")

    # 1. Jinja delimiter syntax checking
    open_brackets = len(re.findall(r'\{\{', sql_query))
    close_brackets = len(re.findall(r'\}\}', sql_query))
    if open_brackets != close_brackets:
        errors.append(f"Mismatched Jinja expression delimiters: {open_brackets} '{{{{' vs {close_brackets} '}}}}'.")

    open_blocks = len(re.findall(r'\{%', sql_query))
    close_blocks = len(re.findall(r'%\}', sql_query))
    if open_blocks != close_blocks:
        errors.append(f"Mismatched Jinja block delimiters: {open_blocks} '{{%' vs {close_blocks} '%}}'.")

    # 2. Extract references & sources
    extracted_refs = re.findall(r"ref\(\s*['\"]([^'\"]+)['\"]\s*\)", sql_query)
    extracted_sources = re.findall(r"source\(\s*['\"]([^'\"]+)['\"]\s*,\s*['\"]([^'\"]+)['\"]\s*\)", sql_query)

    # 3. Compile simulated SQL (substituting Jinja with lakehouse catalogue paths)
    catalog = os.getenv("ICEBERG_CATALOG_NAME", "dlh_lakehouse")
    compiled_sql = sql_query

    for ref in extracted_refs:
        escaped_ref = re.escape(ref)
        replacement = f"{catalog}.{layer}.{ref}" if not ref.startswith("bronze") else f"{catalog}.bronze.{ref}"
        pattern = re.compile(r"\{\{\s*ref\(\s*['\"]" + escaped_ref + r"['\"]\s*\)\s*\}\}")
        compiled_sql = pattern.sub(replacement, compiled_sql)

    for src_schema, src_tbl in extracted_sources:
        escaped_schema = re.escape(src_schema)
        escaped_tbl = re.escape(src_tbl)
        replacement = f"{catalog}.bronze.{src_tbl}"
        pattern = re.compile(r"\{\{\s*source\(\s*['\"]" + escaped_schema + r"['\"]\s*,\s*['\"]" + escaped_tbl + r"['\"]\s*\)\s*\}\}")
        compiled_sql = pattern.sub(replacement, compiled_sql)

    # Clean up dbt config() macro in compiled preview
    compiled_sql = re.sub(r"\{\{\s*config\([^)]*\)\s*\}\}", "-- [dbt config applied]", compiled_sql)

    # 4. Basic SQL structural validation
    clean_sql = re.sub(r'--.*', '', compiled_sql).strip()
    if clean_sql:
        if not re.search(r'\b(select|with)\b', clean_sql, re.IGNORECASE):
            warnings.append("SQL query does not start with SELECT or WITH clause; ensure it produces a valid dataset.")
        if clean_sql.count('(') != clean_sql.count(')'):
            errors.append("Unbalanced parentheses '(' and ')' detected in SQL query.")

    # 5. Generate Schema YAML testing representation
    schema_dict = generate_model_schema_dict(clean_name, test_configs, sql_query)
    wrapper = {"version": 2, "models": [schema_dict]}
    schema_yml_str = yaml.dump(wrapper, sort_keys=False, default_flow_style=False)

    is_valid = len(errors) == 0
    return {
        "status": "success" if is_valid else "error",
        "valid": is_valid,
        "model_name": clean_name,
        "layer": layer,
        "compiled_sql": compiled_sql.strip(),
        "schema_yml": schema_yml_str,
        "extracted_refs": extracted_refs,
        "extracted_sources": [{"schema": s[0], "table": s[1]} for s in extracted_sources],
        "warnings": warnings,
        "errors": errors
    }


def sync_dbt_models_from_nodes(nodes: List[Dict[str, Any]]) -> List[str]:
    """Sync dbt SQL model files, schema.yml (with tests), and sources.yml physically in both dbt target directories."""
    synced_files = []
    
    target_dbt_dirs = [settings.DBT_DIR]
    if hasattr(settings, "DEFAULT_DBT_DIR") and settings.DEFAULT_DBT_DIR and settings.DEFAULT_DBT_DIR != settings.DBT_DIR:
        target_dbt_dirs.append(settings.DEFAULT_DBT_DIR)

    for base_dbt_dir in target_dbt_dirs:
        try:
            models_dir = os.path.join(base_dbt_dir, "models")
            os.makedirs(models_dir, exist_ok=True)

            models_by_layer: Dict[str, List[Dict[str, Any]]] = {}

            # 1. Sync dbt SQL models
            for node in nodes:
                node_type = node.get("data", {}).get("type", "")
                cfg = node.get("data", {}).get("config", {})
                if node_type == "dbt_transform":
                    layer = (cfg.get("dbt_layer") or cfg.get("sel") or "silver").lower()
                    model_name = cfg.get("dbt_model") or cfg.get("sel") or node.get("data", {}).get("label", "").lower().replace(" ", "_")
                    if not model_name:
                        continue
                    if not model_name.endswith(".sql"):
                        model_name += ".sql"
                    
                    target_dir = os.path.join(models_dir, layer)
                    os.makedirs(target_dir, exist_ok=True)
                    
                    sql_content = cfg.get("query") or cfg.get("sql_query_preview") or f"-- dbt Model: {model_name}\nSELECT * FROM {{{{ ref('bronze_raw') }}}}"
                    model_path = os.path.join(target_dir, model_name)
                    with open(model_path, "w", encoding="utf-8") as f:
                        f.write(sql_content)
                    logger.info("Auto-synced dbt model file to %s", model_path)
                    synced_files.append(model_path)

                    # Collect for schema.yml tests
                    test_cfgs = cfg.get("test_configs") or cfg.get("quality_tests")
                    if not test_cfgs and any(k in cfg for k in ["unique_columns", "not_null_columns", "accepted_values", "enforce_contract"]):
                        test_cfgs = {
                            "unique_columns": cfg.get("unique_columns") or [],
                            "not_null_columns": cfg.get("not_null_columns") or [],
                            "accepted_values": cfg.get("accepted_values") or {},
                            "enforce_contract": cfg.get("enforce_contract") or False,
                        }
                    if test_cfgs:
                        model_entry = generate_model_schema_dict(model_name, test_cfgs, sql_content)
                        models_by_layer.setdefault(layer, []).append(model_entry)

            # 2. Auto-sync schema.yml per layer if test_configs are present
            for layer, m_list in models_by_layer.items():
                if not m_list:
                    continue
                schema_file = os.path.join(models_dir, layer, "schema.yml")
                existing_schema: Dict[str, Any] = {"version": 2, "models": []}
                if os.path.exists(schema_file):
                    try:
                        with open(schema_file, "r", encoding="utf-8") as sf:
                            existing_schema = yaml.safe_load(sf) or {"version": 2, "models": []}
                    except Exception:
                        pass
                
                existing_models = {m.get("name"): m for m in existing_schema.get("models", [])}
                for new_m in m_list:
                    existing_models[new_m["name"]] = new_m
                existing_schema["models"] = list(existing_models.values())

                with open(schema_file, "w", encoding="utf-8") as sf:
                    yaml.safe_dump(existing_schema, sf, sort_keys=False, default_flow_style=False)
                logger.info("Auto-synced dbt schema tests to %s", schema_file)
                synced_files.append(schema_file)

            # 3. Auto-sync sources.yml from source_db nodes
            sources_file = os.path.join(models_dir, "sources.yml")
            source_tables_by_schema = {}
            for node in nodes:
                node_type = node.get("data", {}).get("type", "")
                cfg = node.get("data", {}).get("config", {})
                if node_type in ["source_db", "database"]:
                    tbl = cfg.get("table_name", "")
                    if tbl:
                        parts = tbl.split(".")
                        s_name = parts[0] if len(parts) > 1 else "public"
                        t_name = parts[1] if len(parts) > 1 else parts[0]
                        source_tables_by_schema.setdefault(s_name, set()).add(t_name)

            if source_tables_by_schema and os.path.exists(sources_file):
                try:
                    with open(sources_file, "r", encoding="utf-8") as sf:
                        src_data = yaml.safe_load(sf) or {}
                    
                    sources_list = src_data.get("sources", [])
                    modified = False
                    for schema_name, tbl_set in source_tables_by_schema.items():
                        schema_entry = next((s for s in sources_list if s.get("schema") == schema_name), None)
                        if not schema_entry:
                            schema_entry = {"name": schema_name, "schema": schema_name, "tables": []}
                            sources_list.append(schema_entry)
                            modified = True
                        
                        existing_tables = {t.get("name") for t in schema_entry.get("tables", [])}
                        for t in tbl_set:
                            if t not in existing_tables:
                                schema_entry.setdefault("tables", []).append({
                                    "name": t,
                                    "description": f"Auto-registered source table '{t}' from canvas pipeline"
                                })
                                modified = True
                    
                    if modified:
                        src_data["sources"] = sources_list
                        with open(sources_file, "w", encoding="utf-8") as sf:
                            yaml.safe_dump(src_data, sf, sort_keys=False, default_flow_style=False)
                        logger.info("Auto-synced sources.yml at %s", sources_file)
                except Exception as se:
                    logger.warning("Auto-sync sources.yml warning: %s", se)

            # 4. Auto-sync bronze view models from source_db tables so downstream models can ref() them
            bronze_dir = os.path.join(models_dir, "bronze")
            os.makedirs(bronze_dir, exist_ok=True)
            for schema_name, tbl_set in source_tables_by_schema.items():
                for t_name in tbl_set:
                    m_name = f"bronze_{t_name}.sql" if not t_name.startswith("bronze_") else f"{t_name}.sql"
                    m_path = os.path.join(bronze_dir, m_name)
                    if not os.path.exists(m_path):
                        b_content = (
                            f"-- dbt Bronze View Model: {t_name}\n"
                            f"{{{{ config(materialized='view') }}}}\n\n"
                            f"SELECT * FROM {{{{ source('{schema_name}', '{t_name}') }}}}\n"
                        )
                        with open(m_path, "w", encoding="utf-8") as bf:
                            bf.write(b_content)
                        logger.info("Auto-synced bronze dbt model to %s", m_path)
                        synced_files.append(m_path)

        except Exception as e:
            logger.warning("dbt Model Sync Note: %s", e)
            
    return synced_files