core_engine.py

import sys
import re
import json
import os
import spacy
import networkx as nx
import torch
import torch.nn as nn
from transformers import AutoTokenizer, AutoModel
from sentence_transformers.cross_encoder import CrossEncoder


# =========================================================
# NEURAL SQL COMPOSER (HRM ARCHITECTURE)
# =========================================================

class HRMSQLComposer(nn.Module):
    def __init__(self, embed_dim=384, hidden_dim=256, num_actions=9):
        super().__init__()
        # MACRO ACTIONS: 0-8 (SET_DB, SET_PIPELINE, SELECT, FROM, WHERE, GROUP_BY, ORDER_BY, LIMIT, STOP)
        self.gru = nn.GRUCell(embed_dim, hidden_dim)
        
        self.h_module = nn.Sequential(
            nn.Linear(hidden_dim, hidden_dim),
            nn.ReLU(),
            nn.Linear(hidden_dim, num_actions)
        )
        
        self.l_query = nn.Linear(hidden_dim, embed_dim)

    def forward(self, current_input, state):
        """Processes a single step in the sequence"""
        state = self.gru(current_input, state)
        macro_logits = self.h_module(state)
        pointer_query = self.l_query(state)
        return macro_logits, pointer_query, state



#============================================================
# DB Class
#============================================================


from sentence_transformers import SentenceTransformer, util
import sqlite3



class DBValueLookup:
    def __init__(self, db_dir, model_name='all-MiniLM-L6-v2'):
        self.db_dir = db_dir
        self._cache = {}  # (db_id, table, col) -> [values]
        self.device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
        print(f"Loading Semantic DB Lookup model ({model_name})...")
        self.model = SentenceTransformer(model_name).to(self.device)

    def get_values(self, db_id, table, col):
        key = (db_id, table, col)
        if key not in self._cache:
            db_path = f"{self.db_dir}/{db_id}/{db_id}.sqlite"
            try:
                conn = sqlite3.connect(db_path)
                # Fetch distinct values, ignore nulls
                rows = conn.execute(
                    f"SELECT DISTINCT {col} FROM {table} WHERE {col} IS NOT NULL LIMIT 200"
                ).fetchall()
                conn.close()
                self._cache[key] = [str(r[0]) for r in rows if r[0]]
            except Exception:
                self._cache[key] = []
        return self._cache[key]

    def semantic_match(self, value_text, db_id, table, col, threshold=0.5):
        candidates = self.get_values(db_id, table, col)
        if not candidates:
            return value_text
            
        # Fast-path: Exact match saves us running the neural model
        v_lower = value_text.lower()
        for c in candidates:
            if v_lower == c.lower():
                return c
                
        # Semantic matching
        query_embedding = self.model.encode(value_text, convert_to_tensor=True, show_progress_bar=False, device=self.device)
        db_embeddings = self.model.encode(candidates, convert_to_tensor=True, show_progress_bar=False, device=self.device)
        
        cosine_scores = util.cos_sim(query_embedding, db_embeddings)
        best_score_val, best_idx = torch.max(cosine_scores, dim=1)
        best_score = best_score_val.item()
        
        if best_score >= threshold:
            return candidates[best_idx]
        return value_text # Fallback to original text if no good match



#==========================================================
# OPERATORCLISSIFIER CLASSS
# =========================================================



class OperatorClassifier:
    def __init__(self, model_path):
        self.device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
        self.tokenizer = AutoTokenizer.from_pretrained(model_path)
        
        with open(f"{model_path}/operators.json") as f:
            self.operators = json.load(f)
        
        backbone = AutoModel.from_pretrained('cross-encoder/ms-marco-MiniLM-L-6-v2')
        
        class OperatorModel(nn.Module):
            def __init__(self, backbone, n_classes):
                super().__init__()
                self.backbone = backbone
                self.classifier = nn.Linear(384, n_classes)
                self.dropout = nn.Dropout(0.1)
            
            def forward(self, input_ids, attention_mask):
                out = self.backbone(input_ids=input_ids, attention_mask=attention_mask)
                pooled = out.last_hidden_state[:, 0, :]
                return self.classifier(self.dropout(pooled))
        
        self.model = OperatorModel(backbone, len(self.operators)).to(self.device)
        self.model.load_state_dict(torch.load(f"{model_path}/model.pt", map_location=self.device))
        self.model.eval()

    def predict(self, context, is_negated=False):
        """is_negated is passed from bind_values — purely structural"""
        negation_token = " [NEGATED]" if is_negated else ""
        augmented_context = context + negation_token
        
        enc = self.tokenizer([augmented_context], padding=True, truncation=True,
                             max_length=64, return_tensors='pt')
        
        with torch.no_grad():
            logits = self.model(
                enc['input_ids'].to(self.device),
                enc['attention_mask'].to(self.device)
            )
        idx = int(logits.argmax(dim=1).item())
        return self.operators[idx]


#=========================================================
# 1. SCHEMA GRAPH
# =========================================================

class SchemaGraph:

    def __init__(self, schema_text=None, spider_schema=None):
        self.graph = nx.Graph()
        self.tables = {}
        self.primary_keys = {}
        if spider_schema:
            self._parse_spider(spider_schema)
        elif schema_text:
            self._parse_text(schema_text)
        # Only build graph if something was parsed
        # If neither provided, caller will set tables manually then call _build_graph
        if schema_text or spider_schema:
            self._build_graph()

    def _parse_text(self, text):
        table_pattern = re.compile(r'CREATE\s+TABLE\s+(\w+)\s*\((.*?)\);', re.I | re.S)
        for match in table_pattern.finditer(text):
            table = match.group(1).lower()
            body  = match.group(2)
            cols  = {}
            for line in body.split(','):
                line = line.strip()
                if not line:
                    continue
                parts    = line.split()
                col_name = parts[0].lower()
                
                # Dynamic Type Assignment
                col_type = "TEXT"
                if len(parts) > 1 and any(t in parts[1].upper() for t in ["INT", "FLOAT", "NUMBER"]):
                    col_type = "NUM"
                elif col_name.startswith('is_') or col_name.startswith('has_') or col_name.startswith('hlc_is_'):
                    col_type = "BOOL"
                    
                cols[col_name] = col_type
                
                if "PRIMARY KEY" in line.upper() or col_name in ('id', f"{table}_id") or col_name.endswith('id'):
                    self.primary_keys[table] = col_name
            self.tables[table] = cols
            self.graph.add_node(table)

    def _parse_spider(self, schema_json):
        table_names_orig = schema_json['table_names_original']
        table_names_norm = schema_json['table_names'] # The natural language table names
        col_names_orig   = schema_json['column_names_original']
        col_names_norm   = schema_json['column_names'] # The natural language column names
        col_types        = schema_json['column_types']
        
        # New dictionaries to hold the translation mapping
        self.norm_table_names = {}
        self.norm_column_names = {}

        for i, tbl in enumerate(table_names_orig):
            t_name = tbl.lower()
            self.tables[t_name] = {}
            self.graph.add_node(t_name)
            # Map original table name -> normalized table name
            self.norm_table_names[t_name] = table_names_norm[i].lower()

        for i, (tbl_idx, col_name) in enumerate(col_names_orig):
            if tbl_idx == -1:
                continue
            
            t_name = table_names_orig[tbl_idx].lower()
            c_name = col_name.lower()
            
            # Spider stores the normalized name at index 1 of the inner list
            norm_c_name = col_names_norm[i][1].lower() 
            
            # Dynamic Type Assignment
            c_type = "TEXT"
            if col_types[i] == "number":
                c_type = "NUM"
            elif c_name.startswith('is_') or c_name.startswith('has_') or c_name.startswith('hlc_is_'):
                c_type = "BOOL"
                
            self.tables[t_name][c_name] = c_type
            
            # Map (table, original_column) -> normalized column
            self.norm_column_names[(t_name, c_name)] = norm_c_name
            
            if c_name in ('id', f"{t_name}_id"):
                self.primary_keys[t_name] = c_name

        for (col_idx_1, col_idx_2) in schema_json['foreign_keys']:
            t1_idx, c1_name = col_names_orig[col_idx_1]
            t2_idx, c2_name = col_names_orig[col_idx_2]
            t1 = table_names_orig[t1_idx].lower()
            t2 = table_names_orig[t2_idx].lower()
            self.graph.add_edge(t1, t2, on=f"{t1}.{c1_name.lower()} = {t2}.{c2_name.lower()}")




    def _build_graph(self):
        """Build graph using ONLY relationships EXPLICITLY described or strongly implied in the system prompt.
           NO auto-inference from generic columns like owner_id. No Influx joins."""
        tables = list(self.tables.keys())
   
        # 1. Performance Schema <-> SQL Plan (Explicit FK in prompt)
        if 'performance_schema' in tables and 'sql_plan' in tables:
            self.graph.add_edge(
                'performance_schema', 'sql_plan',
                on="performance_schema.sql_plan_id = sql_plan.id"
            )
            
        # # 2. Target <-> Collector / dbactdc_instance (Stated in prompt logic)
        # # "hlc_server_name... Name of the HLC server managing this target"
        # if 'target' in tables and 'dbactdc_instance' in tables:
        #     self.graph.add_edge(
        #         'target', 'dbactdc_instance',
        #         on="target.hlc_server_name = dbactdc_instance.name"
        #     )
            
        # 3. Target <-> Cloud Pricing (Stated in prompt logic)
        if 'target' in tables and 'cloud_database_pricing' in tables:
            self.graph.add_edge(
                'target', 'cloud_database_pricing',
                on="target.cloud_region = cloud_database_pricing.region" 
                #AND target.hlc_database_type = cloud_database_pricing.cloud_database_provider"
            )
            
        # 4. Template <-> Charts Template Element (Explicit FK in prompt)
        if 'template' in tables and 'charts_template_element' in tables:
            self.graph.add_edge(
                'template', 'charts_template_element',
                on="template.id = charts_template_element.template_id"
            )

        # 5. Report <-> Report Template (Explicit FK in prompt)
        if 'report' in tables and 'report_template' in tables:
            self.graph.add_edge(
                'report', 'report_template',
                on="report.id = report_template.report_id"
            )

        # 6. Report Template <-> Report Element Alerts (Explicit FK in prompt)
        if 'report_template' in tables and 'report_template_element_alert' in tables:
            self.graph.add_edge(
                'report_template', 'report_template_element_alert',
                on="report_template.id = report_template_element_alert.report_template_id"
            )

        # === ALL OTHER TABLES ARE ISOLATED ===
        # Influx measurements are deliberately NOT connected
        for t in tables:
            if t not in self.graph:
                self.graph.add_node(t)
   
        print(f"[SchemaGraph] Built with {len(self.graph.edges)} explicit edge(s) based on prompt logic.")



def build_derby_schema():
        schema = SchemaGraph(schema_text=" ")
        schema.tables = {
            'target': {
                'name': 'TEXT', 'is_data_collector_deployed': 'BOOL',
                'hlc_server_name': 'TEXT', 'hlc_server_platform': 'TEXT',
                'hlc_server_type': 'TEXT', 'hlc_username': 'TEXT',
                'hlc_password': 'TEXT', 'hlc_is_private_key_location': 'BOOL',
                'hlc_private_key_location': 'TEXT', 'hlc_pass_phrase': 'TEXT',
                'hlc_database_type': 'TEXT', 'hlc_database_version': 'TEXT',
                'hlc_database_instance': 'TEXT', 'hlc_database_home': 'TEXT',
                'hlc_data_collector_home': 'TEXT', 'hlc_collector_configuration': 'TEXT',
                'hlc_environment_configuration': 'TEXT', 'hlc_is_local_collector': 'BOOL',
                'hlc_is_on_target_collector': 'BOOL', 'hlc_local_home': 'TEXT',
                'hlc_local_configuration': 'TEXT',
                'hlc_is_data_collector_deployment_configuration': 'BOOL',
                'hlc_database_jmx': 'TEXT', 'hlc_database_jmx_user': 'TEXT',
                'hlc_database_jmx_password': 'TEXT', 'hlc_logs_location': 'TEXT',
                'hlc_is_capture_file_location_translation': 'BOOL',
                'hlc_translation_source': 'TEXT', 'hlc_translation_destination': 'TEXT',
                'appdynamics_url': 'TEXT', 'appdynamics_user_name': 'TEXT',
                'appdynamics_account_name': 'TEXT', 'appdynamics_account_password': 'TEXT',
                'newrelic_api_key': 'TEXT', 'license': 'TEXT',
                'hlc_password2': 'TEXT', 'hlc_database_jmx_password2': 'TEXT',
                'is_java_collector': 'BOOL', 'linked_target_name': 'TEXT',
                'is_on_cloud': 'BOOL', 'hlc_credit_cost': 'NUM',
                'cloud_region': 'TEXT', 'cloud_pricing_lookup_code': 'TEXT',
                'hlc_database_auth_type': 'TEXT', 'hlc_database_auth_secret_name': 'TEXT',
                'is_demo_target': 'BOOL', 'db_target_status': 'TEXT',
                'datadog_api_key': 'TEXT', 'datadog_application_key': 'TEXT',
                'dtype': 'TEXT', 'owner_id': 'NUM',
            },
            'template': {
                'name': 'TEXT', 'charts_settings_json': 'TEXT',
                'is_dynamic': 'BOOL', 'created_timestamp': 'NUM',
                'template_type': 'TEXT', 'finance_time_range_settings_json': 'TEXT',
                'finance_targets_count': 'NUM', 'regular_finops_target_name_list': 'TEXT',
                'regular_finops_report_params_json': 'TEXT',
                'is_system_default': 'BOOL', 'owner_id': 'NUM',
            },
            'charts_template_element': {
                'target': 'TEXT', 'capture': 'TEXT', 'metric': 'TEXT',
                'measurement': 'TEXT', 'is_new': 'BOOL', 'sub_metrics': 'TEXT',
                'charts_settings_json': 'TEXT', 'template_id': 'NUM',
                'order_no': 'NUM', 'is_dynamic_element': 'BOOL',
            },
            'report': {
                'name': 'TEXT', 'title': 'TEXT', 'is_scheduled': 'BOOL',
                'scheduler_params_period': 'TEXT', 'scheduler_params_month': 'NUM',
                'scheduler_params_day_of_month': 'NUM',
                'scheduler_params_day_of_week': 'TEXT',
                'scheduler_params_hours': 'NUM', 'scheduler_params_minutes': 'NUM',
                'time_range_start_time': 'NUM', 'time_range_end_time': 'NUM',
                'is_email_html_report_enabled': 'BOOL',
                'is_email_pdf_report_enabled': 'BOOL',
                'is_cc_enabled': 'BOOL', 'cc_email_id': 'TEXT',
                'is_system_default': 'BOOL', 'owner_id': 'NUM',
            },
            'report_template': {
                'template_id': 'TEXT', 'template_name': 'TEXT',
                'report_id': 'NUM', 'name': 'TEXT', 'order_no': 'NUM',
                'template_type': 'TEXT', 'finops_alert_type': 'TEXT',
                'finops_alert_value': 'NUM', 'attached_report_params': 'TEXT',
            },
            'report_template_element_alert': {
                'template_element_id': 'TEXT', 'alert_type': 'TEXT',
                'alert_value': 'NUM', 'report_template_id': 'NUM',
            },
            'integration': {
                'integration_name': 'TEXT', 'title': 'TEXT',
                'enabled': 'BOOL', 'config_json': 'TEXT',
            },
            'dashlet_configuration': {
                'dashlet_type': 'TEXT', 'configuration': 'TEXT',
            },
            'cloud_database_pricing': {
                'cloud_database_provider': 'TEXT', 'credit_cost': 'NUM',
                'region': 'TEXT', 'lookup_code': 'TEXT',
            },
            'dbactdc_instance': {
                'name': 'TEXT', 'ssh_host': 'TEXT', 'ssh_port': 'NUM',
                'ssh_username': 'TEXT', 'ssh_password': 'TEXT',
                'home': 'TEXT', 'secret_key': 'TEXT', 'platform': 'TEXT',
                'service_port': 'NUM', 'is_remote': 'BOOL', 'status': 'TEXT',
                'owner_id': 'NUM', 'jre_executable_path': 'TEXT',
            },
            'performance_schema': {
                'digest': 'TEXT', 'digest_text': 'TEXT',
                'sql_plan_id': 'NUM', 'owner_id': 'NUM',
            },
            'sql_plan': {
                'id': 'NUM', 'sql_plan': 'TEXT', 'sql_plain_text_plan': 'TEXT',
            },
        }
        schema.primary_keys = {
            'performance_schema': 'digest',
            'sql_plan': 'id',
        }
        for table in schema.tables:
            schema.graph.add_node(table)
        schema._build_graph()
        return schema


def build_influx_schema():
        schema = SchemaGraph(schema_text=" ")
        schema.tables = {
            'target_based_time_series_data': {
                'target': 'TEXT', 'capture': 'TEXT', 'metric': 'TEXT',
                'changed_from_zero': 'BOOL', 'value': 'NUM',
            },
            'sys_target_alerts': {
                'target': 'TEXT', 'capture': 'TEXT', 'metric': 'TEXT',
                'measurement_value': 'TEXT', 'templatename': 'TEXT',
                'data_value_double': 'NUM', 'data_value_long': 'NUM',
                'confidencescore': 'NUM', 'is_demo_alert': 'BOOL',
                'owner_id': 'NUM',
            },
        }
        schema.primary_keys = {}
        for table in schema.tables:
            schema.graph.add_node(table)
        return schema



# =========================================================
# 2. LINGUISTIC ENGINE
# =========================================================




class LinguisticEngine:
    def __init__(self, schema_graph, table_model_path, column_model_path,
                 value_model_path, skeleton_model_path, db_id=None):
        self.schema = schema_graph
        self.db_id = db_id
        print("Loading spaCy...")
        self.nlp = spacy.load("en_core_web_sm")
        print("Loading DB value lookup...")
        self.db_lookup = DBValueLookup('.')
        print("Loading table linker...")
        self.table_linker  = CrossEncoder(table_model_path)
        print("Loading column linker...")
        self.column_linker = CrossEncoder(column_model_path)
        print("Loading value linker...")
        self.value_linker  = CrossEncoder(value_model_path)
        # print("Loading skeleton classifier...")
        # self.skeleton = SkeletonClassifier(skeleton_model_path)
        # print("Ready.")
        print("Loading operator classifier...")
        self.operator_clf = OperatorClassifier('./model_operator')
        self.boolean_cols = self._build_boolean_column_registry()
        


    def _serialize_table(self, table_name):
        cols = list(self.schema.tables.get(table_name, {}).keys())
        
        # Use mapped table name and mapped column names
        norm_t_name = getattr(self.schema, 'norm_table_names', {}).get(table_name, table_name)
        norm_cols = [
            getattr(self.schema, 'norm_column_names', {}).get((table_name, c), c) 
            for c in cols
        ]
        
        return f"{norm_t_name} | {', '.join(norm_cols)}"

    def _serialize_column(self, table_name, col_name):
        all_cols = list(self.schema.tables.get(table_name, {}).keys())
        col_type = self.schema.tables[table_name].get(col_name, 'text')
        
        # Safely fetch normalized names (falling back to raw names if not a Spider schema)
        norm_t_name = getattr(self.schema, 'norm_table_names', {}).get(table_name, table_name)
        norm_c_name = getattr(self.schema, 'norm_column_names', {}).get((table_name, col_name), col_name)
        
        # Normalize the context columns as well
        norm_context_list = [
            getattr(self.schema, 'norm_column_names', {}).get((table_name, c), c) 
            for c in all_cols if c != col_name
        ]
        context = ', '.join(norm_context_list)
        
        return f"{norm_t_name} | {norm_c_name} | {col_type} | context: {context}"

    def resolve_ordering(self, question, active_tables):
        q_lower = question.lower()
        doc = self.nlp(q_lower)
        target_spans = []
        direction = "DESC"

        dir_map = {
            'highest': 'DESC', 'largest': 'DESC', 'maximum': 'DESC', 'best': 'DESC', 'oldest': 'DESC', 'most': 'DESC',
            'lowest': 'ASC', 'smallest': 'ASC', 'minimum': 'ASC', 'worst': 'ASC', 'youngest': 'ASC', 'least': 'ASC'
        }

        for token in doc:
            if token.tag_ in ['JJS', 'RBS'] or token.text in dir_map:
                # Safely set direction based on the first superlative encountered
                direction = dir_map.get(token.text, "ASC" if any(x in token.text for x in ['least','fewest','smallest','lowest','worst']) else "DESC")
                
                # Hard-target known abstract concepts to prevent noun pollution
                if token.text in ['youngest', 'oldest']:
                    target_spans.append('age')
                elif token.text in ['most', 'least']:
                    return 'COUNT(*)', direction
                else:
                    head = token.head
                    if head.pos_ in ['NOUN', 'PROPN'] and head.text != token.text:
                        modifiers = [child.text for child in head.children if child.dep_ in ['amod', 'compound'] and child.text != token.text]
                        modifiers.append(head.text)
                        target_spans.append(" ".join(modifiers)) 
                    else:
                        target_spans.append(token.text)
                break

        if not target_spans:
            if any(k in q_lower for k in ["descending order","ascending order","order by","sorted by","ordered by"]):
                chunks = list(doc.noun_chunks)
                target_span = chunks[-1].text if chunks else question.split()[-1]
                target_spans.append(target_span)
                if "asc" in q_lower: direction = "ASC"

        if not target_spans: return None, None



        # --- Structural Domain Override for Time-Series Superlatives ---
        if getattr(self, 'db_id', None) == 'influx_system' and 'sys_target_alerts' in active_tables:
            # Protect aggregate queries! If they ask for a count, do NOT force the raw double value
            is_aggregate = any(w in question.lower() for w in ['count', 'total', 'how many', 'average', 'sum'])
            if not is_aggregate:
                if any(w in question.lower() for w in ['anomaly', 'bottleneck', 'spike']):
                    return 'data_value_double', direction

        candidates = [(t, c) for t in active_tables for c in self.schema.tables[t] if c not in ('id', f"{t}_id")]

        # candidates = [(t, c) for t in active_tables for c in self.schema.tables[t] if c not in ('id', f"{t}_id")]
        if not candidates: return None, None

        best_overall_score = -float('inf')
        best_overall_match = None

        for span in target_spans:
            pairs = [(span, self._serialize_column(t, c)) for t, c in candidates]
            scores = self.column_linker.predict(pairs, show_progress_bar=False)
            best_idx = int(scores.argmax())
            best_score = float(scores[best_idx])
            
            if best_score > best_overall_score:
                best_overall_score = best_score
                best_t, best_c = candidates[best_idx]
                best_overall_match = f"{best_t}.{best_c}"

        return best_overall_match, direction


    def detect_aggregation(self, question):
        q = question.lower()
        found_aggs = [] 

        checks = [
            ([r"\baverage\b", r"\bmean\b", r"\bavg\b"], "AVG"),
            ([r"\bmaximum\b", r"\bmax\b", r"\bhighest\b"], "MAX"),
            ([r"\bminimum\b", r"\bmin\b", r"\blowest\b"], "MIN"),
            ([r"\bhow many\b", r"\bcount the number\b"], "COUNT"), 
            ([r"\btotal sum\b", r"\bsum of\b"], "SUM")
        ]

        # Robust mathematical intent detection for "count"
        if "count" in q and any(w in q for w in ["per", "total", "of", "for"]):
            found_aggs.append((q.find("count"), "COUNT"))

        for keywords, agg_type in checks:
            for kw in keywords:
                match = re.search(kw, q)
                if match:
                    if agg_type == "SUM" and "total number" in q:
                        continue 
                        
                    found_aggs.append((match.start(), agg_type))
                    break

        

        found_aggs.sort(key=lambda x: x[0])
        if 'SUM' in [a for _, a in found_aggs] and 'COUNT' in [a for _, a in found_aggs]:
            if 'total number' in q or 'total count' in q or 'count per' in q:
                found_aggs = [(i, a) for i, a in found_aggs if a == 'COUNT']
        return [agg for _, agg in found_aggs]






    def bind_values(self, question, active_tables, window_size=5, debug=False):
        question = question.replace('\u2018', "'").replace('\u2019', "'").replace('\u201c', '"').replace('\u201d', '"')
       
        quoted = [(m.group(1), False, True) for m in re.finditer(r"['\"](.+?)['\"]", question)]
        top_n_positions = {m.start(1) for m in re.finditer(
            r'\b(?:top|bottom|worst|best)\s+(\d+)\b', question.lower())}
        numbered = []
        for m in re.finditer(r'\b(\d+(?:\.\d+)?)\b', question):
            if m.start() in top_n_positions:
                continue
            num_str = m.group(1)
            if not any(num_str in q[0] for q in quoted):
                numbered.append((num_str, True, False))
       
        doc = self.nlp(question)
        entity_texts = set(q[0].lower() for q in quoted)
        valid_ent_types = {'PERSON', 'ORG', 'GPE', 'LOC', 'FAC', 'PRODUCT', 'NORP', 'EVENT', 'WORK_OF_ART'}
       
        entities = []
        for ent in doc.ents:
            if ent.text.lower() in entity_texts or re.search(r'\d+', ent.text) or ent.label_ not in valid_ent_types:
                continue
            entities.append((ent.text, False, False))
        if self.db_id == 'derby_system' or self.db_id == 'influx_system':
            known_targets = self.db_lookup.get_values('derby_system', 'target', 'name')
            for kt in known_targets:
                if str(kt).lower() in question.lower() and len(str(kt)) > 3:
                    entities.append((str(kt), False, False))
        search_tables = set(active_tables)
        for t in active_tables:
            search_tables.update(self.schema.graph.neighbors(t))
        if self.db_id == 'influx_system':
            search_tables = {'sys_target_alerts'} # isolated

        if self.db_id:
            complex_names = re.findall(r'\b[A-Za-z0-9]+(?:_[A-Za-z0-9_]+|-[\w\-]+)+\b', question)
            numbered = [num for num in numbered if not any(num[0] in c for c in complex_names)]
           
            seen_texts = entity_texts | set(v[0].lower() for v in numbered)
            for name in complex_names:
                if name.lower() not in seen_texts and not any(name.lower() in st for st in seen_texts):
                    entities.append((name, False, False))
                    seen_texts.add(name.lower())
            schema_keywords = set()
            for t in self.schema.tables.keys():
                schema_keywords.add(t.lower())
                schema_keywords.add(t.lower() + 's')
                schema_keywords.add(getattr(self.schema, 'norm_table_names', {}).get(t, t).lower())
            for token in doc:
                if token.pos_ not in ['NOUN', 'PROPN', 'X']: continue
                for txt in [token.text, token.lemma_]:
                    if txt.lower() in seen_texts or len(txt) < 3 or txt.lower() in schema_keywords:
                        continue
                       
                    for table in search_tables:
                        for col, col_type in self.schema.tables.get(table, {}).items():
                            if col_type != 'TEXT': continue
                            db_vals = self.db_lookup.get_values(self.db_id, table, col)
                            if any(txt.lower() == str(v).lower() for v in db_vals):
                                entities.append((txt, False, False))
                                seen_texts.add(txt.lower())
                                seen_texts.add(token.text.lower())
                                break
                        else: continue
                        break
       

        unique_vals = []
        seen = set()
        for v in quoted + numbered + entities:
            if v[0].lower() not in seen:
                seen.add(v[0].lower())
                unique_vals.append(v)
        
        # =========================================================================
        # 🛡️ THE GENERALIZED CONTEXT SHIELD (Your Fix)
        # If the extracted value is a KNOWN Target or Digest, kill it from the 
        # neural pipeline. It will be handled natively via extract_universal_context.
        # =========================================================================
        all_values = []
        for v in unique_vals:
            # Strip quotes for clean lookup
            val_clean = v[0].replace("'", "").replace('"', '').upper()
            
            # Check against the global dictionary
            if val_clean in GLOBAL_REGISTRY["targets"] or val_clean in GLOBAL_REGISTRY["digests"]:
                if debug: print(f"[BIND BYPASS] Skipping '{v[0]}' - Handled natively as global Context Entity.")
                continue # Do not pass to neural net
                
            all_values.append(v)
            
        if not all_values: return []
        # =========================================================================

        candidates = [(table, col, col_type) for table in search_tables for col, col_type in self.schema.tables.get(table, {}).items() if col != 'id' and not col.endswith('_id')]
        if not candidates: return []

        
        filters = []
        skip_vals = set()
       
        for val_text, is_numeric, is_quoted in all_values:
            if val_text in skip_vals: continue
            val_pos = question.lower().find(val_text.lower())
            if val_pos == -1: continue
            before = question[:val_pos].split()[-window_size:]
            after = question[val_pos + len(val_text):].split()[:window_size]
            context = ' '.join(before + [val_text] + after)
            
            # =====================================================================
            # STRUCTURAL NEGATION DETECTION (no regex patterns)
            # =====================================================================
            negation_words = {'not', "n't", 'never', 'without', 'excluding', 'except', 'non-', 
                              'disable', 'inactive', 'stopped', 'off', 'false'}
            # is_negated = any(w in context.lower() for w in negation_words) and \
            #              not any(phrase in context.lower() for phrase in ['not only', 'not just', 'not even'])
            is_negated = any(re.search(rf'\b{w}\b', context.lower()) for w in ['not', "n't", 'never', 'without', 'excluding', 'except', 'non'])


            
            valid_candidates = candidates
            is_complex_identifier = bool(re.match(r'^[A-Za-z0-9]+(?:_[A-Za-z0-9_]+|-[\w\-]+)+$', val_text))

            # =========================================================================
            # THE HARD-BINDING (DELEXICALIZATION) BYPASS
            # =========================================================================
            if self.db_id and not is_numeric and not is_quoted:
                exact_matches = []
                for t, c, ct in candidates:
                    if ct == 'TEXT':
                        db_vals = self.db_lookup.get_values(self.db_id, t, c)
                        if val_text.lower() in [str(v).lower() for v in db_vals]:
                            exact_matches.append((t, c, ct))
                            
                if exact_matches:
                    # <<< SHIELD TARGET TABLE PRIORITY >>>
                    target_matches = [m for m in exact_matches if m[0] == 'target']
                    if len(target_matches) > 0:
                        exact_matches = target_matches

                    # <<< SHIELD IDENTIFIER PRIORITY >>>
                    id_matches = [m for m in exact_matches if m[1] in ('name', 'target', 'integration_name')]
                    
                    if id_matches:
                        best_t, best_c, _ = id_matches[0] # Force lock to identifier
                    elif len(exact_matches) == 1:
                        best_t, best_c, _ = exact_matches[0]
                    else:
                        valid_candidates = exact_matches
                        continue
                        
                    grounded_val = self.db_lookup.semantic_match(val_text, self.db_id, best_t, best_c)
                    is_neg_bypass = any(re.search(rf'\b{w}\b', context.lower()) for w in ['not', "n't", 'never', 'without', 'excluding', 'except', 'non'])
                    op = '!=' if is_neg_bypass else '='
                    intent = 'EXCLUDE' if is_neg_bypass else 'INCLUDE'
                    
                    like_triggers = ['have', 'having', 'contain', 'containing', 'start', 'end', 'like', 'starts with', 'ends with']
                    if any(re.search(rf'\b{k}\b', context.lower()) for k in like_triggers) and not is_neg_bypass:
                        op = 'LIKE'
                        grounded_val = f"%{grounded_val}%"
                    
                    filters.append((best_t, best_c, op, f"'{grounded_val}'", intent))
                    if debug: print(f"[HARD-BIND BYPASS] Locked '{val_text}' -> {best_t}.{best_c}. Skipping neural.")
                    continue
            
            # --- FALLBACK: RUN THE NEURAL NETWORK (CROSS-ENCODER) ---
            pairs = [(context, self._serialize_column(t, c)) for t, c, ct in valid_candidates]
            scores = self.value_linker.predict(pairs, show_progress_bar=False)
            sorted_indices = scores.argsort()[::-1]
            for idx in sorted_indices:
                best_score = float(scores[idx])
                t, c, ct = valid_candidates[idx]
               
                # --- TYPE SHIELDING: BOOLEAN BYPASS ---
                if ct == 'BOOL':
                    core_concept = c.replace('is_', '').replace('has_', '').replace('hlc_is_', '').replace('_', ' ')
                    if core_concept not in question.lower():
                        continue
                    val = "'false'" if is_negated else "'true'"
                    filters.append((t, c, '=', val, 'INCLUDE'))
                    break
                
                col_clean = c.lower().replace("_", "")
                q_clean = question.lower().replace("_", "").replace(" ", "")
                is_lexical = col_clean in q_clean
               
                threshold = -2.0 if is_numeric else -5.0
               
                if best_score <= threshold and not is_lexical:
                    continue
                if ct == 'NUM' and not is_numeric: continue
               
                # <<< KEY CHANGE: Pass negation flag to OperatorClassifier >>>
                op = self.operator_clf.predict(context, is_negated=is_negated)
                
                like_triggers = ['have', 'having', 'contain', 'containing', 'start', 'end', 'like', 'starts with', 'ends with']
                if is_quoted and any(k in context.lower() for k in like_triggers):
                    op = 'LIKE'
                    val = f"'%{val_text}%'"
                else:
                    if not is_numeric and not is_quoted and ct == 'TEXT' and self.db_id:
                        grounded_val = self.db_lookup.semantic_match(val_text, self.db_id, t, c)
                        val = f"'{grounded_val}'"
                    else:
                        val = val_text if is_numeric else f"'{val_text}'"
               
                if op == 'BETWEEN':
                    remaining = question[question.lower().find(val_text.lower()) + len(val_text):]
                    next_num = re.search(r'\b(\d+(?:\.\d+)?)\b', remaining)
                    if next_num:
                        next_val = next_num.group(1)
                        val = f"{val_text} AND {next_val}"
                        skip_vals.add(next_val)
                    else:
                        op = '='
                
                intent = 'EXCLUDE' if (op in ('!=', 'NOT LIKE') or is_negated) and ct == 'TEXT' else \
                         ('INCLUDE' if op in ('=', 'LIKE') and ct == 'TEXT' else 'COMPARE')
                
                filters.append((t, c, op, val, intent))
                break
               
        return filters




    def _build_boolean_column_registry(self):

        boolean_cols = set()
    
        for table_name, columns in self.schema.tables.items():

            for col_name, col_info in columns.items():
            # col_info is stored as a plain string type tag e.g. "TEXT", "NUM", "BOOL"
                col_type = col_info.upper() if isinstance(col_info, str) else ''
            
                # Strategy 1: explicit BOOL type tag (set during _parse_text / _parse_spider)
                if col_type == 'BOOL':
                    boolean_cols.add((table_name, col_name))
            
                # Strategy 2: naming convention patterns
                elif re.match(r'^is_|^has_|^hlc_is_', col_name):
                    boolean_cols.add((table_name, col_name))
            
                # Strategy 3: enum-like columns — low cardinality status/type/platform cols
                # These should filter, not display
                elif re.match(r'.+_(status|type|platform|period|provider)$', col_name):
                    boolean_cols.add((table_name, col_name))
            
                # Strategy 4: day-of-week, enabled flag
                elif col_name in ('enabled', 'status', 'platform'):
                    boolean_cols.add((table_name, col_name))
    
        return boolean_cols



    def extract_intent(self, question, top_k_tables=3, top_k_cols=6, table_margin=2.0, col_margin=2.0, debug=False):
        # PATCH: Derby execution-plan queries must stay in performance_schema + sql_plan only.
        # With 12 tables in derby_system, the report/template tables score within margin
        # for phrasing like "get the execution plan for digest X" because words like
        # "execution", "plan", "time" match scheduler/report column descriptions.
        # We short-circuit the neural table linker entirely for this intent.
        q_lower_pre = question.lower()

        _force_tables = None
        _injected_digest_filter = None
        if getattr(self, 'db_id', None) == 'derby_system':
            plan_kws = ['execution plan', 'sql plan', 'explain plan', 'query plan',
                        'plain text plan', 'plan for digest', 'plan for the digest']
            has_plan = any(kw in q_lower_pre for kw in plan_kws)
            has_digest_ref = bool(re.search(r'\bdigest\b', q_lower_pre)) and 'plan' in q_lower_pre
            if has_plan or has_digest_ref:
                _force_tables = ['performance_schema', 'sql_plan']
                # Also extract digest value if present but unquoted, so value binder
                # doesn't miss it (bind_values only catches quoted or numeric tokens).
                digest_val_match = re.search(
                    r"\bdigest\s+['\"]?([\w\-]+--[\w\-]+)['\"]?", q_lower_pre
                )
                if digest_val_match:
                    _injected_digest_filter = digest_val_match.group(1)

        all_tables = list(self.schema.tables.keys())
        table_pairs  = [(question, self._serialize_table(t)) for t in all_tables]

        all_tables = [t for t in self.schema.tables.keys()]
        if self.db_id == 'influx_system':
            # If we are in Influx, the ONLY valid tables are metrics
            valid_tables = ['sys_target_alerts', 'target_based_time_series_data']
        else:
            # If we are in Derby, exclude the Influx-specific tables
            valid_tables = [t for t in self.schema.tables.keys() if t not in ['sys_target_alerts', 'target_based_time_series_data']]
            
        # Filter the neural network's search space
        table_pairs = [(question, self._serialize_table(t)) for t in valid_tables]


        
        table_scores = self.table_linker.predict(table_pairs, show_progress_bar=False)

        sorted_table_indices = table_scores.argsort()[::-1]
        best_table_score = table_scores[sorted_table_indices[0]]

        active_tables = []
        for idx in sorted_table_indices[:top_k_tables]:
            score = table_scores[idx]
            if score >= (best_table_score - table_margin):
                # active_tables.append(all_tables[idx])
                active_tables.append(valid_tables[idx])

        # Apply force-lock after neural scoring (preserves scores for col linker)
        # if _force_tables is not None:
        #     active_tables = _force_tables

        # Apply force-lock after neural scoring (preserves scores for col linker)
        if _force_tables is not None:
            active_tables = _force_tables

        # Store injected digest for generate_sql to consume via engine attribute
        self._injected_digest_filter = getattr(self, '_injected_digest_filter', None) \
            if not hasattr(self, '_injected_digest_filter') else _injected_digest_filter
        self._injected_digest_filter = _injected_digest_filter

        # ── 1. The Smart NLP Lexical Anchor Extraction ──
        doc = self.nlp(question.lower())
        
        # Extract base lemmas ONLY for meaningful parts of speech. 
        meaningful_lemmas = {
            token.lemma_ for token in doc 
            if token.pos_ in ['NOUN', 'PROPN', 'ADJ']
        }
        
        # Extract noun chunks for multi-word concepts
        noun_chunks = {chunk.text for chunk in doc.noun_chunks}

        # ── 2. Smart Table Rescue ──
        # Force-include tables if their clean lemma is explicitly a noun in the question.
        for t in all_tables:
            # Check against the raw table name AND Spider's normalized English name
            norm_t = getattr(self.schema, 'norm_table_names', {}).get(t, t).lower()
            
            if t in meaningful_lemmas or norm_t in meaningful_lemmas:
                if t not in active_tables:
                    active_tables.append(t)

        if debug:
            print(f"[extract_intent] active_tables={active_tables}")

        col_hits = []
        for table in active_tables:
            cols = [c for c in self.schema.tables[table] if c not in ('id', f"{table}_id")]
            if not cols:
                continue

            col_pairs  = [(question, self._serialize_column(table, c)) for c in cols]
            col_scores = self.column_linker.predict(col_pairs, show_progress_bar=False)

            sorted_col_indices = col_scores.argsort()[::-1]
            best_col_score = col_scores[sorted_col_indices[0]]

            for idx in sorted_col_indices:
                score = col_scores[idx]
                if score >= (best_col_score - col_margin):
                    col_hits.append((table, cols[idx], float(score)))



        # ── 3. Smart Column Rescue ──
        already_included = {(t, c) for t, c, s in col_hits}
        
        for table in active_tables:
            for col in self.schema.tables[table]:
                if col in ('id',) or col.endswith('_id'):
                    continue
                
                # Fetch Spider's normalized name (e.g., "fname" -> "first name")
                norm_c = getattr(self.schema, 'norm_column_names', {}).get((table, col), col).lower()
                
                # We rescue the column if its normalized name is a direct lemma 
                # OR if it's explicitly contained inside a noun chunk.
                is_lemma_match = norm_c in meaningful_lemmas
                is_chunk_match = any(norm_c in chunk for chunk in noun_chunks)
                
                if is_lemma_match or is_chunk_match:
                    if (table, col) not in already_included:
                        # Append with a score of 0.0 so it survives pruning but doesn't override neural top-picks
                        col_hits.append((table, col, 0.0)) 

        rescue_triggers = {
            'enabled': ['disabled', 'inactive', 'off', 'disable'],
            'status': ['running', 'stopped', 'failed', 'error', 'active', 'inactive']
        }
        
        for col_name, triggers in rescue_triggers.items():

            if any(w in question.lower() for w in triggers):
                for t in active_tables:
                    if col_name in self.schema.tables.get(t, {}):
                        # Check if the column is already in the list of tuples
                        if not any(c == col_name for _, c, _ in col_hits):
                            col_hits.append((t, col_name, 0.0))

        if debug:
            print(f"[extract_intent] col_hits={col_hits}")


        filters = self.bind_values(question, active_tables, debug=debug)

        bound_as_target_name = any(f[0] == 'target' and f[1] == 'name' for f in filters)
        if bound_as_target_name and getattr(self, 'db_id', None) == 'derby_system':
            plan_kws = ['digest', 'sql', 'plan', 'query', 'execution']
            if not any(kw in q_lower_pre for kw in plan_kws):
                active_tables = [t for t in active_tables if t not in ('performance_schema', 'sql_plan')]
        
        
        # Rescue any tables discovered by the DB value matcher back into active_tables
        for t, c, op, val, intent in filters:
            if t not in active_tables:
                active_tables.append(t)
                
        if debug:
            print(f"[extract_intent] filters={filters}")

        return active_tables, col_hits, filters


# =========================================================
# 3. SQL COMPOSER (HRM WRAPPER)
# =========================================================

import torch
import re
import networkx as nx

MACRO_MAP = {0: "SET_DB", 1: "SET_PIPELINE", 2: "SELECT", 3: "FROM", 4: "WHERE", 5: "GROUP_BY", 6: "ORDER_BY", 7: "LIMIT", 8: "STOP"}




DOMAIN_LEXICON = {
    "table_mappings": {
        "database": "target",
        "databases": "target",
        "overview": "target",
        "system": "target",
        "systems": "target",
        "collector": "dbactdc_instance",
        "collectors": "dbactdc_instance",
        "agent": "dbactdc_instance",
        "agents": "dbactdc_instance",
        "integration": "integration",
        "integrations": "integration",
        "third party": "integration",
        "template": "template",
        "templates": "template",
        "dashboard": "template",
        "dashboards": "template",
        "report": "report",
        "reports": "report",
        "target": "target",
        "targets": "target",
        "anomaly": "sys_target_alerts",
        "anomalies": "sys_target_alerts",
        "spike": "sys_target_alerts",
        "spikes": "sys_target_alerts",
        "bottleneck": "sys_target_alerts",
        "bottlenecks": "sys_target_alerts",
        "alert": "sys_target_alerts",
        "alerts": "sys_target_alerts",
        "issue": "sys_target_alerts",
        "issues": "sys_target_alerts",
        "execution plan": "sql_plan",
        "sql text": "performance_schema",
        "query text": "performance_schema"
    },
    "column_mappings": {
        "linux": ("target", "hlc_server_platform", "'Linux'"),
        "windows": ("target", "hlc_server_platform", "'Windows'"),
        "aix": ("target", "hlc_server_platform", "'AIX'"),
        "cloud": ("target", "is_on_cloud", "'true'"),
        "cloud-hosted": ("target", "is_on_cloud", "'true'"),
        "on-prem": ("target", "cloud_region", "'on-prem'"),
        "demo": ("target", "is_demo_target", "'true'"),
        "stopped": ("target", "db_target_status", "'Stopped'"),
        "down": ("target", "db_target_status", "'Stopped'"),
        "aws": ("cloud_database_pricing", "cloud_database_provider", "'AWS'"),
        "inactive": ("dbactdc_instance", "status", "'Inactive'"),
        "active": ("dbactdc_instance", "status", "'Active'"),
        "remote": ("dbactdc_instance", "is_remote", "'true'"),
        "pdf": ("report", "is_email_pdf_report_enabled", "'true'"),
        "html": ("report", "is_email_html_report_enabled", "'true'"),
        "daily": ("report", "scheduler_params_period", "'DAILY'"),
        "monday": ("report", "scheduler_params_day_of_week", "'MONDAY'"),
        "disabled": ("integration", "enabled", "'false'"),
        "enabled": ("integration", "enabled", "'true'"),
        "mysql": ("target", "hlc_database_type", "'MySQL'"),
        "oracle": ("target", "hlc_database_type", "'Oracle'"),
        "sqlserver": ("target", "hlc_database_type", "'SQLServer'"),
        "sql server": ("target", "hlc_database_type", "'SQLServer'"),
        "postgres": ("target", "hlc_database_type", "'PostgreSQL'"),
        "postgresql": ("target", "hlc_database_type", "'PostgreSQL'"),
        "dynamic": ("template", "is_dynamic", "'true'"),
        "scheduled": ("report", "is_scheduled", "'true'")
    }
}




def enrich_and_inject_nuance(question, active_tables, col_hits, raw_filters, engine):
    filters = list(raw_filters)
    q_lower = question.lower()
    
    # ── 0. UNIFIED CONTEXT EXTRACTION ──
    # Replace the old extract_target_context with our new universal extractor
    ctx = extract_universal_context(question)
    t_name = ctx["target"]
    db_type = ctx["type"]
    engine_id = getattr(engine, 'db_id', None)

    # --- 1. LEXICON TABLE OVERRIDE ---
    lexicon_tables = []
    for term, tbl in DOMAIN_LEXICON["table_mappings"].items():
        if re.search(rf'\b{term}\b', q_lower) and tbl in engine.schema.tables:
            if tbl not in lexicon_tables: lexicon_tables.append(tbl)
            
    # Force explicit lexicon noun tables to be the ROOT nodes
    for tbl in reversed(lexicon_tables):
        if tbl in active_tables: active_tables.remove(tbl)
        active_tables.insert(0, tbl)
                
    # --- 2. LEXICON FILTER INJECTION (Engine Aware) ---
    for term, (tbl, col, val) in DOMAIN_LEXICON["column_mappings"].items():
        if re.search(rf'\b{term}\b', q_lower):
            
            # 🧠 THE FIX: Rely on the actual engine processing the request, not the global route
            if engine_id == 'influx_system' and tbl == 'target':
                tbl = 'sys_target_alerts'
                if col == 'name': col = 'target'
                if col == 'is_demo_target': col = 'is_demo_alert'

            if tbl in engine.schema.tables:
                # Check 25 characters before the matched term for negations
                match = re.search(rf'\b{term}\b', q_lower)
                window = q_lower[max(0, match.start() - 25):match.start()]
                is_negated_lex = any(re.search(rf'\b{w}\b', window) for w in ['not', 'except', 'excluding', 'without', 'non', 'neither'])
                
                lex_op = '!=' if is_negated_lex else '='
                lex_intent = 'EXCLUDE' if is_negated_lex else 'INCLUDE'
                
                LEXICON_OVERRIDE_COLUMNS = {'is_on_cloud', 'is_demo_target', 'is_data_collector_deployed', 
                             'db_target_status', 'hlc_server_type'}
                already_bound = any(fc == col for _, fc, _, _, _ in filters)
                if not already_bound:
                    filters.append((tbl, col, lex_op, val, lex_intent))
                elif col in LEXICON_OVERRIDE_COLUMNS:
                    # Lexicon wins over bind_values for these columns — remove the neural binding and replace
                    filters = [f for f in filters if f[1] != col]
                    filters.append((tbl, col, lex_op, val, lex_intent))

            

    # --- 3. SURGICAL DOMAIN LOGIC ---
    if t_name:
        tbl = 'sys_target_alerts' if engine_id == 'influx_system' else 'target'
        col = 'target' if tbl == 'sys_target_alerts' else 'name'
        
        # Check for negation around the target name
        match = re.search(rf'\b{re.escape(t_name.lower())}\b', q_lower)
        is_neg_target = False
        if match:
            window = q_lower[max(0, match.start() - 30):match.start()]
            is_neg_target = any(re.search(rf'\b{w}\b', window) for w in ['not', 'except', 'excluding', 'without', 'non', 'neither'])
            
        t_op = '!=' if is_neg_target else '='
        t_intent = 'EXCLUDE' if is_neg_target else 'INCLUDE'
        
        if not any(fc == col for _, fc, _, _, _ in filters):
            filters.append((tbl, col, t_op, f"'{t_name}'", t_intent))
        if tbl not in active_tables: active_tables.append(tbl)

    db_type_keywords = set(DOMAIN_LEXICON["column_mappings"].keys()) & {
        'mysql', 'oracle', 'sqlserver', 'sql server', 'postgres', 'postgresql'
    }
    user_mentioned_db_type = any(re.search(rf'\b{kw}\b', q_lower) for kw in db_type_keywords)
    
    # 🧠 THE FIX: Only inject Derby-specific database types if we are actually scanning Derby
    if db_type and engine_id == 'derby_system' and user_mentioned_db_type:
        if not any(fc == 'hlc_database_type' for _, fc, _, _, _ in filters):
            filters.append(('target', 'hlc_database_type', '=', f"'{db_type}'", 'INCLUDE'))
        if 'target' not in active_tables: active_tables.append('target')

    # 🧠 THE FIX: Only inject Influx-specific demo guards if we are in Influx
    if engine_id == 'influx_system':
        if not any(fc == 'is_demo_alert' for _, fc, _, _, _ in filters):
            filters.append(('sys_target_alerts', 'is_demo_alert', '!=', "'true'", 'INCLUDE'))
        if 'sys_target_alerts' not in active_tables: active_tables.append('sys_target_alerts')
        
        # Kill data_value_long hallucination
        col_hits = [ch for ch in col_hits if ch[1] != 'data_value_long']



    # --- 4. SCRUBBER (Neural-Safe) ---
    cleaned_filters = []
    table_names = [tbl.lower() for tbl in engine.schema.tables.keys()]
    # Get a flat set of every column name in the database
    all_col_names = {col.lower() for tbl_cols in engine.schema.tables.values() for col in tbl_cols}

    for f in filters:
        ft, fc, fop, fval, fintent = f
        # Strip quotes for comparison
        val_clean = fval.replace("'", "").lower()
    
        # ❌ RULE 1: If the value is literally a TABLE name, it's structural noise.
        if val_clean in table_names:
            continue

        # ❌ RULE 2: If the value is literally a COLUMN name, it's a hallucination.
        if val_clean in all_col_names:
            continue

        # ❌ RULE 3: If the value is the database ID (e.g., 'derby_system'), kill it.
        if val_clean == engine.db_id:
            continue
        
        cleaned_filters.append(f)

    unique_filters = []
    seen = set()
    for f in cleaned_filters:
        # Force lower casing for deduplication matching
        f_str = f"{f[0]}.{f[1]} {f[2]} {f[3]}".lower()
        if f_str not in seen:
            seen.add(f_str)
            unique_filters.append(f)
            
    return active_tables, col_hits, unique_filters




# def generate_sql_hrm(question, engine, hrm_model, embedder, device, debug=False, injected_tables=None, injected_col_hits=None, injected_filters=None):
def generate_sql_hrm(question, engine, hrm_model, embedder, device, debug=False, injected_tables=None, injected_col_hits=None, injected_filters=None, target_db=None):
    q_lower = question.lower()
    true_db = target_db if target_db is not None else route_query(question)  
    if true_db == 'multi_step':
        true_db = 'influx_system'
    # true_db = route_query(question)
    
    # --- MULTI-STEP ROUTING HEURISTIC ---
    needs_sql_text = bool(re.search(r'\b(sql quer(?:y|ies)|execution plans?|actual sql|actual quer(?:y|ies)|quer(?:y|ies) (?:behind|causing)|what sql|what quer(?:y|ies)|look like)\b', q_lower))
    needs_anomalies = bool(re.search(r'\b(worst|bottlenecks?|anomal(?:y|ies)|spikes?|issues?|problems?|most anomalous|severe)\b', q_lower))
    pipeline = "MULTI_STEP_PIPELINE" if (needs_sql_text and needs_anomalies) else "STANDARD_QUERY"

    # 1. Use INJECTED arrays if provided by Orchestrator
    if injected_tables is not None and injected_filters is not None:
        active_tables = list(injected_tables)
        col_hits = list(injected_col_hits) if injected_col_hits else []
        filters = list(injected_filters)
    else:
        # Otherwise, run raw extraction
        active_tables, col_hits, raw_filters = engine.extract_intent(question, debug=debug)
        active_tables, col_hits, filters = enrich_and_inject_nuance(question, active_tables, col_hits, raw_filters, engine)


        # =========================================================
        # DETERMINISTIC PLATFORM INJECTION
        # =========================================================
        platforms = ['Linux', 'Windows', 'AIX', 'Solaris', 'HP-UX', 'Ubuntu', 'RHEL'] 
        found_platforms = [p for p in platforms if re.search(rf'\b{p.lower()}\b', q_lower)]
        
        if found_platforms and true_db == 'derby_system':
            filters = [f for f in filters if f[1] not in ('hlc_server_platform', 'platform')]
            for p_name in found_platforms:
                match = re.search(rf'\b{p_name.lower()}\b', q_lower)
                is_negated = False
                if match:
                    window = q_lower[max(0, match.start() - 35):match.start()]
                    is_negated = any(re.search(rf'\b{w}\b', window) for w in ['not', 'except', 'excluding', 'without', 'non', 'neither'])
                    
                op = '!=' if is_negated else '='
                intent_val = 'EXCLUDE' if is_negated else 'INCLUDE'
                filters.append(('target', 'hlc_server_platform', op, f"'{p_name}'", intent_val))
            if 'target' not in active_tables: active_tables.append('target')

        # =========================================================
        # DYNAMIC TABLE PRUNING (Principled Graph Resolution)
        # =========================================================
        explicit_tables = set()
        for term, tbl in DOMAIN_LEXICON["table_mappings"].items():
            if re.search(rf'\b{term}\b', q_lower):
                explicit_tables.add(tbl)

        # Prune peripheral tables if target is the star
        if 'target' in active_tables:
            peripheral_tables = ['dbactdc_instance', 'cloud_database_pricing', 'integration']
            for p_table in peripheral_tables:
                if p_table in active_tables:
                    is_explicit = p_table in explicit_tables
                    has_hard_filter = any(f[0] == p_table for f in filters)
                    if not is_explicit and not has_hard_filter:
                        if debug: print(f"[GENERATOR] ✂️ Pruning neural noise: '{p_table}'")
                        active_tables.remove(p_table)
                        
        # Prune target if a peripheral table is the star and target is noise (Fixes Q62)
        if 'target' in active_tables and 'target' not in explicit_tables:
            has_target_filter = any(f[0] == 'target' for f in filters)
            if not has_target_filter and len(explicit_tables) > 0:
                if debug: print(f"[GENERATOR] ✂️ Pruning neural noise: 'target'")
                active_tables.remove('target')

    aggs = engine.detect_aggregation(question)    
    sort_col_raw, sort_dir = engine.resolve_ordering(question, active_tables)
    top_n_match = re.search(r'\b(?:top|worst|best|bottom)\s+(\d+)\b', q_lower)
    forced_limit = int(top_n_match.group(1)) if top_n_match else None    
    # true_db = route_query(question)

    # ── Fix 3: Influx isolation (NOW SAFE — filters exists)
    if true_db == 'influx_system':
        if 'capture' not in q_lower and 'time series' not in q_lower:
            active_tables = ['sys_target_alerts']
            if debug:
                print("[INFLUX GUARD] Forced sys_target_alerts only (no target_based_time_series_data join)")

    
    # ── Fix 1: Single-filter collapse (NOW SAFE)
    all_filter_tables = set(t for t, c, op, val, intent in filters) if filters else set()
    
    for ft in all_filter_tables:
        if ft not in active_tables and ft in engine.schema.tables:
            active_tables.append(ft)

    # -------------------------------------------------------------------------
    # 🚨 GRAPH ALIGNMENT FIX: Prevent Ghost Aliases + PROMPT-COMPLIANT ISOLATION
    # -------------------------------------------------------------------------
    valid_tables = []
    path_nodes = set()

    if len(active_tables) > 1:
        root = active_tables[0]
        valid_tables.append(root)
        path_nodes.add(root)
   
        for tgt in active_tables[1:]:
            try:
                path = nx.shortest_path(engine.schema.graph, root, tgt)
                path_nodes.update(path)
                valid_tables.append(tgt)
            except nx.NetworkXNoPath:
                pass
           
        ordered_nodes = []
        for t in valid_tables:
            if t not in ordered_nodes: ordered_nodes.append(t)
        for t in path_nodes:
            if t not in ordered_nodes: ordered_nodes.append(t)
    else:
        ordered_nodes = active_tables


    
    start_node = ordered_nodes[0] if ordered_nodes else active_tables[0]
    table_aliases = {tbl: f"T{i+1}" for i, tbl in enumerate(ordered_nodes)}
    subgraph = engine.schema.graph.subgraph(ordered_nodes)
    use_aliases = len(ordered_nodes) > 1

    if use_aliases:
        join_clauses = [f"{start_node} AS {table_aliases[start_node]}"]
        actually_joined_tables = {start_node} # Track what ACTUALLY gets joined


        try:
            for u, v in nx.bfs_edges(subgraph, source=start_node):
                edge_data = engine.schema.graph.get_edge_data(u, v)
                if edge_data:
                    on_cond = edge_data['on']
                    # Handle complex edges with multiple conditions (AND)
                    conditions = on_cond.split(' AND ')
                    parsed_conds = []
                    
                    for cond in conditions:
                        left_part, right_part = cond.split(' = ')
                        t1, c1 = left_part.strip().split('.')
                        t2, c2 = right_part.strip().split('.')
                        u_col, v_col = (c1, c2) if t1 == u else (c2, c1)
                        parsed_conds.append(f"{table_aliases[u]}.{u_col} = {table_aliases[v]}.{v_col}")
                        
                    join_clauses.append(f"JOIN {v} AS {table_aliases[v]} ON {' AND '.join(parsed_conds)}")
                    actually_joined_tables.add(v)
        except Exception as e: 
            if debug: print(f"[DEBUG] Join builder skipped an edge due to error: {e}")
            pass
        joins_sql = ' ' + ' '.join(join_clauses)
        
        # 🚨 THE GHOST BUSTER: Prune ordered_nodes to strictly match the FROM clause
        ordered_nodes = [t for t in ordered_nodes if t in actually_joined_tables]
        
        # Re-evaluate aliases if we pruned the list down to 1 table
        use_aliases = len(ordered_nodes) > 1
        if not use_aliases:
             joins_sql = f" {ordered_nodes[0]}"
             table_aliases = {}
    else:
        joins_sql = f" {start_node}"

    def apply_alias(col_str):
        if "." in col_str:
            t_name, c_name = col_str.split('.')
            if t_name not in ordered_nodes: return c_name # Safety bypass
            if not use_aliases: return c_name
            if t_name in table_aliases: return f"{table_aliases[t_name]}.{c_name}"
        return col_str


    db_options = [f"[DB] {true_db}", "[DB] influx_system" if true_db == "derby_system" else "[DB] derby_system"]
    inventory = db_options + ["[PIPELINE] STANDARD_QUERY", "[PIPELINE] MULTI_STEP_PIPELINE"]
    type_masks = {"TABLE": [], "COLUMN": [], "FILTER": [], "DB": [0, 1], "PIPELINE": [2, 3], "ORDER_BY": []}
    
    tbl_str = f"[TABLE] {joins_sql.strip()}"
    inventory.append(tbl_str)
    type_masks["TABLE"].append(inventory.index(tbl_str))
    
    valid_col_hits = [ch for ch in col_hits if ch[0] in ordered_nodes]
    
    for t, c, score in valid_col_hits[:8]:
        col_str = apply_alias(f"{t}.{c}") if use_aliases else c
        item = f"[COLUMN] {col_str}"
        if item not in inventory: inventory.append(item)
        type_masks["COLUMN"].append(inventory.index(item))
        
        col_type = engine.schema.tables.get(t, {}).get(c, 'TEXT')
        has_sort_keyword = any(w in q_lower for w in ['top', 'bottom', 'worst', 'best', 'highest', 'lowest', 'order', 'sort'])
        
        if col_type == 'NUM' or c == 'time' or has_sort_keyword:
            inv_asc = f"[ORDER_ASC] {col_str}"
            inv_desc = f"[ORDER_DESC] {col_str}"
            if inv_asc not in inventory:
                inventory.extend([inv_asc, inv_desc])
                type_masks["ORDER_BY"].extend([inventory.index(inv_asc), inventory.index(inv_desc)])
    
    valid_filters = [f for f in filters if f[0] in ordered_nodes]
    
    for t, c, op, val, intent in valid_filters:
        col_str = apply_alias(f"{t}.{c}") if use_aliases else c
        item = f"[FILTER] {col_str} {op} {val}"
        if item not in inventory: inventory.append(item)
        type_masks["FILTER"].append(inventory.index(item))
        
    fallbacks = ["*"]
    if aggs and "COUNT" in aggs:
        fallbacks.append("count(*)")
        
    for fallback in fallbacks:
        item = f"[COLUMN] {fallback}"
        if item not in inventory: inventory.append(item)
        type_masks["COLUMN"].append(inventory.index(item))

    if debug: print(f"📦 V3 ALIGNED INVENTORY ({len(inventory)} items): {inventory}")

    with torch.no_grad():
        q_emb = embedder.encode(question, convert_to_tensor=True, device=device).unsqueeze(0)
        inv_embs = embedder.encode(inventory, convert_to_tensor=True, device=device)
        
        state = torch.zeros(1, 256).to(device)
        current_input = q_emb
        
        slots = {"SELECT": [], "FROM": [], "WHERE": [], "GROUP_BY": [], "ORDER_BY": [], "LIMIT": []}
        target_db = true_db 
        
        for step in range(15): 
            macro_logits, pointer_query, state = hrm_model(current_input, state)
            macro_idx = torch.argmax(macro_logits, dim=-1).item()
            macro_action = MACRO_MAP[macro_idx]
            
            if macro_action == "STOP": break
                
            pointer_logits = torch.matmul(pointer_query, inv_embs.T)
            
            valid_indices = None
            if macro_action == "SET_DB": valid_indices = type_masks["DB"]
            elif macro_action == "SET_PIPELINE": valid_indices = type_masks["PIPELINE"]
            elif macro_action == "SELECT": valid_indices = type_masks["COLUMN"]
            elif macro_action == "FROM": valid_indices = type_masks["TABLE"]
            elif macro_action == "WHERE": valid_indices = type_masks["FILTER"]
            elif macro_action == "ORDER_BY": valid_indices = type_masks.get("ORDER_BY", type_masks["COLUMN"])
            elif macro_action == "GROUP_BY": valid_indices = type_masks["COLUMN"]
            
            if valid_indices is not None:
                if len(valid_indices) > 0:
                    mask = torch.full_like(pointer_logits, -1e9)
                    mask[0, valid_indices] = pointer_logits[0, valid_indices]
                    pointer_logits = mask
                else:
                    continue

            micro_idx = torch.argmax(pointer_logits, dim=-1).item()
            chosen_item = inventory[micro_idx]
            clean_item = re.sub(r'^\[.*?\]\s*', '', chosen_item)
            
            if macro_action == "ORDER_BY":
                if "[ORDER_ASC]" in chosen_item:
                    clean_item = f"{clean_item} ASC"
                elif "[ORDER_DESC]" in chosen_item:
                    clean_item = f"{clean_item} DESC"
            
            # 🚨 DATABASE ROUTING LOCK FIX
            if macro_action == "SET_DB": 
                pass # LOCKED: Prevents "no such table" errors
            elif macro_action == "SET_PIPELINE": 
                pass # LOCKED: The regex heuristic controls this now!
            elif macro_action == "LIMIT":
                top_n = re.search(r'\b(?:top|worst|best|bottom)\s+(\d+)\b', question.lower())
                slots["LIMIT"] = [f"LIMIT {top_n.group(1)}" if top_n else "LIMIT 5"]
            else:
                if clean_item not in slots[macro_action]:
                    slots[macro_action].append(clean_item)
            
            current_input = inv_embs[micro_idx].unsqueeze(0)



    

    # =====================================================================
    # 5. INTEGRATE HELPER LOGIC & CONVERSATIONAL BYPASS (NEURAL/STATE DRIVEN)
    # =====================================================================
    target_db = true_db 
    
    # --- 1. SELECT CLAUSE RESCUE (Trusting the Neural Hits) ---
    # Clean up redundant '*' 
    if "*" in slots["SELECT"] and len(slots["SELECT"]) > 1:
        if not any("count" in s.lower() for s in slots["SELECT"]):
            slots["SELECT"] = [c for c in slots["SELECT"] if c != "*"]

    where_cols_clean = []
    for w in slots["WHERE"]:
        col_part = str(w).split(' ')[0] 
        where_cols_clean.append(col_part.split('.')[-1]) 

    # Collect high-scoring columns not used in the WHERE clause
    high_score_cols = []
    
    # ✅ FIX: Iterate over valid_col_hits, or check `if t in ordered_nodes:`
    for t, c, score in col_hits:
        if t not in ordered_nodes:
            continue # Skip columns for tables that aren't in our FROM/JOIN clause

        # Score > 10 means the user explicitly asked for this field
        if score > 10.0 and c not in where_cols_clean and c not in ['is_demo_alert', 'confidencescore', 'data_value_long', 'owner_id']:
            if engine.schema.tables.get(t, {}).get(c) is not None:
                col_format = apply_alias(f"{t}.{c}") if use_aliases else c
                high_score_cols.append(col_format)

    if high_score_cols:
        # If the neural net lazily picked '*', overwrite it with the explicit columns!
        if slots["SELECT"] == ["*"]:
            slots["SELECT"] = []
            
        for col_format in reversed(high_score_cols): # Reverse to keep order when inserting at 0
            if col_format not in slots["SELECT"]:
                slots["SELECT"].insert(0, col_format)
                
        # Quality of Life: If we are querying 'target', always include the name so results make sense
        if 'target' in active_tables:
            name_col = apply_alias("target.name") if use_aliases else "name"
            if name_col not in slots["SELECT"]:
                 slots["SELECT"].insert(0, name_col)

    # Fallback if STILL empty
    if not slots["SELECT"] and not aggs:
        slots["SELECT"] = ["*"]

    # --- 2. WHERE CLAUSE GROUPING ---
    exclude_groups = {}
    include_groups = {}

    for t, c, op, val, intent in filters:
        if t in ordered_nodes:
            col_str = apply_alias(f"{t}.{c}") if use_aliases else c
            
            # Group excludes together for NOT IN(...)
            if intent == 'EXCLUDE' or op == '!=':
                key = col_str
                if key not in exclude_groups: exclude_groups[key] = []
                exclude_groups[key].append(val)
                
            # 🛡️ THE FIX: Group includes together for IN(...)
            elif intent == 'INCLUDE' and op == '=':
                key = col_str
                if key not in include_groups: include_groups[key] = []
                include_groups[key].append(val)
                
            else:
                f_str = f"{col_str} {op} {val}"
                if not any(f_str in w for w in slots["WHERE"]):
                    slots["WHERE"].append(f_str)

    # Process grouped exclusions
    for col_str, vals in exclude_groups.items():
        unique_vals = list(set(vals))
        if len(unique_vals) > 1:
            f_str = f"{col_str} NOT IN ({', '.join(unique_vals)})"
        else:
            f_str = f"{col_str} != {unique_vals[0]}"
        if not any(f_str in w for w in slots["WHERE"]):
            slots["WHERE"].append(f_str)
            
    # Process grouped inclusions
    for col_str, vals in include_groups.items():
        unique_vals = list(set(vals))
        if len(unique_vals) > 1:
            f_str = f"{col_str} IN ({', '.join(unique_vals)})"
        else:
            f_str = f"{col_str} = {unique_vals[0]}"
        if not any(f_str in w for w in slots["WHERE"]):
            slots["WHERE"].append(f_str)

    # --- 3. SORTING AND LIMITS ---
    has_sorting_intent = bool(sort_col_raw) or bool(forced_limit)
    
    if not has_sorting_intent:
        slots["ORDER_BY"] = []
        slots["LIMIT"] = []
    else:
        if forced_limit:
            slots["LIMIT"] = [f"LIMIT {forced_limit}"]
        elif has_sorting_intent and not slots["LIMIT"]:
            slots["LIMIT"] = ["LIMIT 5"] # Default fallback limit
            
        if target_db == 'influx_system':
            # DOMAIN LOGIC: For anomalies, "worst" = HIGHEST score (DESC). 
            # Only use ASC if they explicitly ask for the best/lowest.
            if debug: print(f"[INFLUX FORMAT] aggs={aggs}, forced_limit={forced_limit}, has_sorting_intent={has_sorting_intent}")
            is_asking_for_low = any(w in q_lower for w in ['best', 'lowest', 'smallest', 'least', 'bottom'])
            influx_dir = "ASC" if is_asking_for_low else "DESC"
            slots["ORDER_BY"] = [f"data_value_double {influx_dir}"]
            
        elif target_db == 'derby_system' and sort_col_raw and not slots["ORDER_BY"]:
            slots["ORDER_BY"] = [f"{sort_col_raw.split('.')[-1]} {sort_dir}"]

    # --- 4. INFLUXDB SPECIFIC FORMATTING ---
    if target_db == 'influx_system':
        if debug: print(f"[INFLUX FORMAT] aggs={aggs}, forced_limit={forced_limit}, has_sorting_intent={has_sorting_intent}")
        if not any('is_demo_alert' in str(w) for w in slots["WHERE"]):
            slots["WHERE"].append("is_demo_alert != 'true'")
            
        if forced_limit or has_sorting_intent:
            slots["SELECT"] = ['target', 'metric', 'data_value_double', 'time']
            slots["GROUP_BY"] = [] 
        elif aggs:
            if 'COUNT' in aggs:
                slots["SELECT"], slots["GROUP_BY"] = ["target", "count(*)"], ["target"]
            elif 'AVG' in aggs:
                slots["SELECT"], slots["GROUP_BY"] = ["target", "avg(data_value_double)"], ["target"]
            elif 'MIN' in aggs:
                slots["SELECT"], slots["GROUP_BY"] = ["target", "min(data_value_double)"], ["target"]
        else:
            # Rely on active tables instead of raw string checks
            if 'capture' in [c[1] for c in valid_col_hits if c[2] > 0]:
                slots["SELECT"] = ['target', 'capture', 'data_value_double']
            else:
                slots["SELECT"] = ['target', 'metric', 'data_value_double', 'time']

    # --- 5. DERBY SPECIFIC FORMATTING ---
    elif target_db == 'derby_system':
        if debug: 
            print(f"[AGG DEBUG] slots SELECT = {slots['SELECT']}")
            print(f"[AGG DEBUG] valid_col_hits = {valid_col_hits}")
            print(f"[AGG DEBUG] active_tables = {active_tables}")
            for t, c, s in valid_col_hits:
                dtype = engine.schema.tables.get(t, {}).get(c, 'NOT FOUND')
                print(f"[AGG DEBUG] {t}.{c} dtype={dtype}")

        if aggs:
            agg_func = aggs[0].lower()
            if agg_func == 'count':
                slots["SELECT"] = ["count(*)"]
            else:
                # Only NUM columns are valid for avg/sum/min/max
                num_cols = [c for c in slots["SELECT"] 
                            if c != "*" and
                            any(engine.schema.tables.get(t, {}).get(c.split('.')[-1], '') == 'NUM'
                                for t in active_tables)]
                if num_cols:
                    slots["SELECT"] = [f"{agg_func}({num_cols[0]})"]
                else:
                    # Fall back to highest scoring NUM column from cross-encoder hits
                    num_hits = [(t, c, s) for t, c, s in valid_col_hits
                                if engine.schema.tables.get(t, {}).get(c, '') == 'NUM']
                    if num_hits:
                        top_t, top_c, _ = num_hits[0]
                        col_str = apply_alias(f"{top_t}.{top_c}") if use_aliases else top_c
                        slots["SELECT"] = [f"{agg_func}({col_str})"]
                    elif valid_col_hits:
                        # Last resort — take top hit regardless of type
                        top_table, top_col_name, _ = valid_col_hits[0]
                        col_str = apply_alias(f"{top_table}.{top_col_name}") if use_aliases else top_col_name
                        slots["SELECT"] = [f"{agg_func}({col_str})"]

        # --- DB2 / SQL Plan Join Logic (State-Driven) ---
        is_plan_query = 'sql_plan' in active_tables
        has_digest_filter = 'performance_schema' in active_tables and slots["WHERE"]
        
        if is_plan_query and has_digest_filter:
            slots["FROM"] = ["performance_schema AS T1 JOIN sql_plan AS T2 ON T1.sql_plan_id = T2.id"]
            
            # NLP/Hit based select
            if any('plain_text' in c for c in slots["SELECT"]) or 'plain text' in q_lower: 
                slots["SELECT"] = ["T2.sql_plain_text_plan"]
            elif not is_plan_query:
                slots["FROM"], slots["SELECT"] = ["performance_schema"], ["digest_text"]
            else: 
                slots["SELECT"] = ["T2.sql_plan"]
            
            # Alias cleanup
            new_where = []
            for w in slots["WHERE"]:
                w_clean = re.sub(r'\b(?:T\d+|performance_schema|sql_plan)\.digest\b', 'T1.digest', str(w))
                w_clean = re.sub(r'(?<!\.)\bdigest\b', 'T1.digest', w_clean)
                new_where.append(w_clean)
            slots["WHERE"] = new_where

    # 6. ASSEMBLE SQL
    sel = "SELECT " + (", ".join(slots["SELECT"]) if slots["SELECT"] else "*")
    frm = " FROM " + (", ".join(slots["FROM"]) if slots["FROM"] else joins_sql.strip())
    whr = ""
    if slots["WHERE"]: whr = " WHERE " + " AND ".join([str(w) for w in slots["WHERE"]])
    grp = " GROUP BY " + ", ".join(slots["GROUP_BY"]) if slots["GROUP_BY"] else ""
    ordr = " ORDER BY " + ", ".join(slots["ORDER_BY"]) if slots["ORDER_BY"] else ""
    lmt = " " + slots["LIMIT"][-1] if slots["LIMIT"] else ""
    
    final_sql = f"{sel}{frm}{whr}{grp}{ordr}{lmt}".strip()
    
    return {
        "target_db": target_db,
        "pipeline_type": pipeline,
        "sql": final_sql
    }

# =========================================================
# ROUTING & ORCHESTRATION
# =========================================================

print("\nInitializing Semantic Router...")
router_model = SentenceTransformer('all-MiniLM-L6-v2')

faq_intents = {

    # ── SCHEMA META ──
    "LIST_DERBY_TABLES": {
        "anchors": [
            "What tables are in the database?", "Show me the Derby database schema.",
            "List all the tables you have.", "What configuration tables exist?",
            "Show me the metadata tables.", "what tables do you have",
            "show me your schema", "what data do you store", "list tables",
            "what's in the derby database", "database structure"
        ],
        "clarification": "Did you mean to list all configuration and metadata tables in the Derby system?",
        "target_db": "derby_system",
        "sql": "SELECT name FROM sqlite_master WHERE type='table';"
    },
    "LIST_INFLUX_TABLES": {
        "anchors": [
            "What tables are in InfluxDB?", "Show me the time series measurements.",
            "Where are the anomalies stored?", "What does the Influx database contain?",
            "what metrics do you store", "show influx schema", "what time series data exists",
            "where is performance data stored", "influx measurements"
        ],
        "clarification": "Did you mean to list the time-series metric tables in InfluxDB?",
        "target_db": "influx_system",
        "sql": "SELECT name FROM sqlite_master WHERE type='table';"
    },

    # ── TARGETS ──
    "LIST_TARGETS": {
        "anchors": [
            "show me all targets", "list my targets", "what targets do I have", "all my databases", "what am I monitoring",
            "list all monitored databases", "show targets", "get all targets",
            "what databases are being monitored", "give me the target list",
            "show me everything you monitor"
        ],
        "clarification": "Did you mean to list all monitored database targets?",
        "target_db": "derby_system",
        "sql": "SELECT name, hlc_database_type, db_target_status FROM target;"
    },

    "TARGET_DETAILS": {
        "anchors": [
            "give me target detail", "show target details", "target profile", 
            "get details for target", "configuration for target", "what is the target detail",
            "show me everything about target"
        ],
        "clarification": "Did you mean to view the full configuration details for this specific target?",
        "target_db": "derby_system",
        "sql": "SELECT * FROM target" 
    },
    

    "LIST_SUPPORTED_DATABASES": {
        "anchors": [
            "What are the different databases you have?", "Which database engines are supported?",
            "List the target database types.", "Do you support Oracle and MySQL?",
            "what database types do you monitor", "which engines are monitored",
            "do you support postgres", "what kind of databases", "supported db types",
            "show me database vendors", "which database flavors"
        ],
        "clarification": "Did you mean to ask what database engine types are currently monitored?",
        "target_db": "derby_system",
        "sql": "SELECT DISTINCT hlc_database_type FROM target WHERE hlc_database_type IS NOT NULL;"
    },

    "TARGET_COST": {
        "anchors": [
            "how much do targets cost", "show target credit costs", "what is the monitoring cost",
            "show me credit usage per target", "target billing", "how much is each target costing"
        ],
        "clarification": "Did you mean to view the credit cost assigned to each monitored target?",
        "target_db": "derby_system",
        "sql": "SELECT name, hlc_credit_cost, cloud_region FROM target WHERE hlc_credit_cost > 0;"
    },

    # ── CLOUD PRICING ──
    "CLOUD_PRICING_INFO": {
        "anchors": [
            "How does cloud pricing work?", "Show me the cost of cloud databases.",
            "What are the pricing lookup codes?", "Show me billing and credit costs for AWS.",
            "cloud pricing table", "show pricing info", "what does cloud monitoring cost",
            "AWS pricing", "Azure pricing", "GCP pricing", "credit cost table",
            "show me the pricing chart", "cloud cost breakdown"
        ],
        "clarification": "Did you mean to view the cloud database pricing chart?",
        "target_db": "derby_system",
        "sql": "SELECT DISTINCT cloud_database_provider, region, credit_cost FROM cloud_database_pricing;"
    },
    "CLOUD_PRICING_BY_PROVIDER": {
        "anchors": [
            "show AWS prices", "what does Azure cost", "GCP credit cost",
            "pricing for Amazon RDS", "show me costs by cloud provider"
        ],
        "clarification": "Did you mean to see pricing broken down by cloud provider?",
        "target_db": "derby_system",
        "sql": "SELECT cloud_database_provider, region, credit_cost, lookup_code FROM cloud_database_pricing ORDER BY cloud_database_provider;"
    },

    # ── COLLECTORS (DBACTDC) ──
    "COLLECTOR_INFO": {
        "anchors": [
            "What are collector agents?", "Show me the dbactdc instances.",
            "List all the monitoring guards.", "Where are the collectors installed?",
            "What are the SSH hosts for the collectors?", "show me collectors",
            "list my collectors", "which collectors are running", "are my collectors active",
            "collector status", "dbactdc status", "is the collector up",
            "show collector agents", "list monitoring agents", "show dbactdc",
            "get collector list", "collector health"
        ],
        "clarification": "Did you mean to list the active collector agents (DBACTDC instances)?",
        "target_db": "derby_system",
        "sql": "SELECT name, status, platform, ssh_host FROM dbactdc_instance;"
    },

    
    "COLLECTOR_PLATFORM": {
        "anchors": [
            "which collectors are on windows", "linux collectors", "show collector platforms",
            "what OS are collectors on", "remote collectors", "local collectors"
        ],
        "clarification": "Did you mean to see what platform each collector runs on?",
        "target_db": "derby_system",
        "sql": "SELECT name, platform, is_remote, ssh_host FROM dbactdc_instance;"
    },

    # ── INTEGRATIONS ──
    "EXTERNAL_INTEGRATIONS": {
        "anchors": [
            "What external integrations are supported?", "Do you connect to Datadog or AppDynamics?",
            "Show me the third party tools.", "What alerts can be sent to New Relic?",
            "show integrations", "list integrations", "what tools are connected",
            "show third party connections", "what monitoring tools are linked",
            "do you integrate with anything", "show me all integrations",
            "which integrations are enabled", "external connections"
        ],
        "clarification": "Did you mean to view the configured external integrations?",
        "target_db": "derby_system",
        "sql": "SELECT integration_name, title, enabled FROM integration;"
    },


    # ── REPORTS ──
    "REPORTING_CAPABILITIES": {
        "anchors": [
            "What types of reports exist?", "Show me the report scheduling options.",
            "Can I get PDF emails?", "List the automated reports.",
            "show me reports", "list reports", "what reports do I have",
            "show scheduled reports", "report list", "what is being reported",
            "show me all reports", "get report list", "reporting schedule"
        ],
        "clarification": "Did you mean to ask about the automated report schedules and formats?",
        "target_db": "derby_system",
        "sql": "SELECT name, title, scheduler_params_period, is_email_pdf_report_enabled FROM report;"
    },

    "PDF_REPORTS": {
        "anchors": [
            "which reports send PDF emails", "show PDF report list",
            "what reports email PDFs", "PDF email reports"
        ],
        "clarification": "Did you mean to list reports that send PDF emails?",
        "target_db": "derby_system",
        "sql": "SELECT name, title FROM report WHERE is_email_pdf_report_enabled = 'true';"
    },

    # ── TEMPLATES & DASHBOARDS ──
    "DASHBOARD_TEMPLATES": {
        "anchors": [
            "What dashboard templates do you have?", "Show me the UI dashlets.",
            "What is a FinOps template?", "How are charts configured?",
            "show templates", "list templates", "what templates exist",
            "show me dashboards", "chart templates", "list dashboards",
            "what dashboards are available", "show all templates",
            "show me FinOps templates", "show dynamic templates"
        ],
        "clarification": "Did you mean to list the available dashboard and chart templates?",
        "target_db": "derby_system",
        "sql": "SELECT name, template_type, is_dynamic FROM template;"
    },

    "DYNAMIC_TEMPLATES": {
        "anchors": [
            "which templates are dynamic", "show dynamic dashboards",
            "list adaptive templates", "templates that auto update"
        ],
        "clarification": "Did you mean to list templates that dynamically adapt to available metrics?",
        "target_db": "derby_system",
        "sql": "SELECT name, template_type FROM template WHERE is_dynamic = 'true';"
    },
    "SYSTEM_DEFAULT_TEMPLATES": {
        "anchors": [
            "which templates are system defaults", "show default templates",
            "built in templates", "out of the box templates"
        ],
        "clarification": "Did you mean to list system-provided default templates?",
        "target_db": "derby_system",
        "sql": "SELECT name, template_type FROM template WHERE is_system_default = 'true';"
    },

    # ── PERFORMANCE SCHEMA / SQL DIGESTS ──
    "PERFORMANCE_SCHEMA_INFO": {
        "anchors": [
            "Where are the execution plans stored?", "How do you track SQL digests?",
            "Show me how SQL text is mapped.", "What is the performance schema?",
            "show sql digests", "list recent sql queries", "show me captured sql",
            "what sql has been captured", "show performance schema",
            "show digest table", "list sql digests", "recent query digests"
        ],
        "clarification": "Did you mean to ask how SQL queries and execution plans are stored?",
        "target_db": "derby_system",
        "sql": "SELECT digest, substr(digest_text, 1, 80) as sql_preview FROM performance_schema LIMIT 10;"
    },
    "SQL_PLANS": {
        "anchors": [
            "show sql execution plans", "list query plans", "show me query plans",
            "what execution plans are stored", "show all plans", "sql plan table"
        ],
        "clarification": "Did you mean to view stored SQL execution plans?",
        "target_db": "derby_system",
        "sql": "SELECT id, substr(sql_plain_text_plan, 1, 100) as plan_preview FROM sql_plan LIMIT 10;"
    },



    # ── OUT OF SCOPE / REJECT ──
    "OUT_OF_SCOPE": {
        "anchors": [
            "what is the weather today", "tell me a joke", "who won the football game",
            "write me a poem", "what is 2 plus 2", "how do I cook pasta",
            "what is the capital of France", "who is the president",
            "recommend a movie", "what time is it", "tell me something funny",
            "what is the meaning of life", "help me write an email",
            "translate this to spanish", "what stocks should I buy"
        ],
        "clarification": "I can only answer questions about your monitored database targets, performance metrics, collectors, reports, and SQL analysis. Could you rephrase?",
        "target_db": None,
        "sql": None
    }
}




# ---------------------------------------------------------
# EXHAUSTIVE TARGET DIALECT ROUTING & EXTRACTION
# ---------------------------------------------------------
TARGET_DIALECTS = {
    "SHOW_TABLES": {
        "mysql": "SELECT table_name FROM information_schema.tables WHERE table_schema = DATABASE();",
        "oracle": "SELECT table_name FROM all_tables;",
        "postgres": "SELECT tablename FROM pg_tables WHERE schemaname NOT IN ('pg_catalog', 'information_schema');",
        "sqlserver": "SELECT name AS table_name FROM sys.tables;",
        "snowflake": "SELECT table_name FROM INFORMATION_SCHEMA.TABLES;",
        "mongodb": "db.getCollectionNames();", 
        "cassandra": "SELECT keyspace_name, table_name FROM system_schema.tables;",
        "redshift": "SELECT tablename FROM pg_catalog.svv_table_info;",
        "db2": "SELECT tabname FROM SYSCAT.TABLES WHERE tabschema NOT LIKE 'SYS%';",
        "sybase": "SELECT name FROM sysobjects WHERE type = 'U';",
        "clickhouse": "SELECT name FROM system.tables WHERE database != 'system';"
    },
    "SHOW_COLUMNS": {
        "mysql": "SELECT table_name, column_name, data_type FROM information_schema.columns WHERE table_schema = DATABASE();",
        "oracle": "SELECT table_name, column_name, data_type FROM all_tab_columns;",
        "postgres": "SELECT relname AS table_name, relkind FROM pg_class WHERE relkind = 'r';",
        "sqlserver": "SELECT object_name(object_id) AS table_name, name AS column_name FROM sys.columns;",
        "snowflake": "SELECT table_name, column_name, data_type FROM INFORMATION_SCHEMA.COLUMNS;",
        "mongodb": "db.collection.findOne();", 
        "cassandra": "SELECT keyspace_name, table_name, column_name, type FROM system_schema.columns;",
        "redshift": "SELECT tablename, column, type FROM pg_catalog.pg_table_def;",
        "db2": "SELECT tabname, colname, typename FROM SYSCAT.COLUMNS;",
        "sybase": "SELECT object_name(id) AS table_name, name AS column_name FROM syscolumns;",
        "clickhouse": "SELECT table, name, type FROM system.columns WHERE database != 'system';"
    },
    "SHOW_INDEXES": {
        "mysql": "SELECT index_name, table_name FROM information_schema.statistics;",
        "oracle": "SELECT index_name, table_name FROM all_indexes;",
        "postgres": "SELECT indexname, tablename FROM pg_indexes;",
        "sqlserver": "SELECT object_name(object_id) AS table_name, name AS index_name FROM sys.indexes;",
        "snowflake": "SHOW INDEXES;", 
        "mongodb": "db.collection.getIndexes();",
        "cassandra": "SELECT keyspace_name, table_name, index_name FROM system_schema.indexes;",
        "redshift": "SELECT tablename, sortkey1 FROM pg_catalog.svv_table_info;",
        "db2": "SELECT indname, tabname FROM SYSCAT.INDEXES;",
        "sybase": "SELECT name FROM sysindexes;",
        "clickhouse": "SELECT name, table FROM system.parts;"
    },
    
    # ─── NEW: ACTIVE SESSIONS (Processlist/Activity) ───
    "SHOW_ACTIVE_SESSIONS": {
        "mysql": "SELECT * FROM information_schema.processlist WHERE command != 'Sleep';",
        "oracle": "SELECT sid, serial#, status, osuser, machine, sql_id FROM V$SESSION WHERE type != 'BACKGROUND';",
        "postgres": "SELECT pid, usename, state, query, backend_start FROM pg_stat_activity WHERE state != 'idle';",
        "sqlserver": "SELECT session_id, status, host_name, program_name FROM sys.dm_exec_sessions WHERE is_user_process = 1;",
        "snowflake": "SELECT * FROM TABLE(INFORMATION_SCHEMA.LOGIN_HISTORY()) LIMIT 100;",
        "mongodb": "db.currentOp();",
        "cassandra": "SELECT * FROM system.clients;",
        "redshift": "SELECT pid, user_name, starttime, query FROM pg_catalog.stv_rec;",
        "db2": "SELECT application_handle, session_auth_id, application_name FROM TABLE(MON_GET_CONNECTION(NULL, -1));",
        "sybase": "SELECT spid, status, loginame, hostname FROM master..sysprocesses;",
        "clickhouse": "SELECT query_id, user, query, elapsed FROM system.processes;"
    },

    # ─── NEW: SYSTEM CONFIGURATION (Parameters/Settings) ───
    "SHOW_SYSTEM_CONFIG": {
        "mysql": "SELECT * FROM performance_schema.global_variables;",
        "oracle": "SELECT name, value, display_value FROM V$SYSTEM_PARAMETER;",
        "postgres": "SELECT name, setting, short_desc FROM pg_settings;",
        "sqlserver": "SELECT name, value, value_in_use, description FROM sys.configurations;",
        "snowflake": "SHOW PARAMETERS;",
        "mongodb": "db.serverCmdLineOpts();",
        "cassandra": "SELECT * FROM system.local;",
        "redshift": "SHOW ALL;",
        "db2": "SELECT name, value FROM SYSIBMADM.DBMCFG;",
        "sybase": "sp_configure;",
        "clickhouse": "SELECT name, value FROM system.settings;"
    },

    # ─── NEW: TABLE STATISTICS (Optimizer/Row counts) ───
    "SHOW_TABLE_STATS": {
        "mysql": "SELECT table_name, table_rows, avg_row_length FROM information_schema.tables WHERE table_schema = DATABASE();",
        "oracle": "SELECT table_name, num_rows, last_analyzed FROM all_tables;",
        "postgres": "SELECT relname, n_live_tup, last_analyze, last_autoanalyze FROM pg_stat_user_tables;",
        "sqlserver": "SELECT object_name(object_id) as table_name, row_count FROM sys.dm_db_partition_stats;",
        "snowflake": "SELECT table_name, row_count, bytes FROM INFORMATION_SCHEMA.TABLES;",
        "mongodb": "db.collection.stats();",
        "cassandra": "SELECT keyspace_name, table_name, comment FROM system_schema.tables;",
        "redshift": "SELECT \"table\", tbl_rows, size FROM pg_catalog.svv_table_info;",
        "db2": "SELECT tabname, card as row_count, stats_time FROM SYSCAT.TABLES;",
        "sybase": "SELECT object_name(id), rowcnt FROM systabstats;",
        "clickhouse": "SELECT table, total_rows, total_bytes FROM system.tables;"
    },

    "EXPLAIN_TEMPLATE": {
        "mysql": "EXPLAIN ANALYZE {query};",
        "oracle": "EXPLAIN PLAN FOR {query};",
        "postgres": "EXPLAIN (ANALYZE, BUFFERS) {query};",
        "sqlserver": "SET SHOWPLAN_TEXT ON; {query};",
        "snowflake": "EXPLAIN USING TABULAR {query};",
        "mongodb": "db.collection.find({query}).explain();",
        "clickhouse": "EXPLAIN {query};",
        "redshift": "EXPLAIN {query};",
        "db2": "EXPLAIN PLAN SELECTION FOR {query};",
        "sybase": "SET SHOWPLAN ON; {query};"
    }
}



dba_intents = {
    "TARGET_SHOW_TABLES": {
        "anchors": ["show me the tables for", "list tables in target", "what tables exist in", "show all tables", "database structure for", "what are the tables in"],
        "intent_key": "SHOW_TABLES"
    },
    "TARGET_SHOW_COLUMNS": {
        "anchors": ["show columns for", "what are the columns in", "list data types for", "show table schema in", "what fields are in"],
        "intent_key": "SHOW_COLUMNS"
    },
    "TARGET_SHOW_INDEXES": {
        "anchors": ["show indexes for", "list the indexes in", "what indexes are on", "show index statistics"],
        "intent_key": "SHOW_INDEXES"
    },
    # ─── NEW EXPANDED DIAGNOSTIC INTENTS ───
    "TARGET_ACTIVE_SESSIONS": {
        "anchors": ["show active sessions", "who is connected to", "current processes", "running sessions", "process list", "what is running right now", "database activity"],
        "intent_key": "SHOW_ACTIVE_SESSIONS"
    },
    "TARGET_SYSTEM_CONFIG": {
        "anchors": ["show database configuration", "system parameters", "database settings for", "show configs", "memory settings", "resource settings", "how is it configured"],
        "intent_key": "SHOW_SYSTEM_CONFIG"
    },
    "TARGET_TABLE_STATS": {
        "anchors": ["table statistics", "show table stats", "optimizer statistics", "when were tables analyzed", "stats collection", "row counts"],
        "intent_key": "SHOW_TABLE_STATS"
    },
    # ───────────────────────────────────────
    "TARGET_EXPLAIN": {
        "anchors": ["explain why this is slow", "get execution plan for", "run explain analyze on", "show plan for", "how is this query running", "explain why the query with digest is slow", "get execution plan for digest", "explain digest"],
        "intent_key": "EXPLAIN_TEMPLATE"
    }
}



def normalize_db_type(raw_type):
    """Maps Derby's hlc_database_type to our dictionary keys."""
    if not raw_type: return None
    raw = raw_type.lower()
    if 'mysql' in raw: return 'mysql'
    if 'postgres' in raw or 'psql' in raw: return 'postgres'
    if 'oracle' in raw: return 'oracle'
    if 'sqlserver' in raw or 'mssql' in raw: return 'sqlserver'
    if 'snowflake' in raw: return 'snowflake'
    if 'mongo' in raw: return 'mongodb'
    if 'cassandra' in raw: return 'cassandra'
    if 'redshift' in raw: return 'redshift'
    if 'db2' in raw: return 'db2'
    return raw




def build_universal_registry(derby_db_path="./derby_system/derby_system.sqlite"):
    registry = {
        "targets": set(),
        "types": {},      # name -> engine_type (e.g., 'DB_01' -> 'mysql')
        "regions": set(),
        "platforms": set(),
        "db_engines": set(), # categorical types like 'Oracle', 'MySQL'
        "digests": set()
    }
    try:
        conn = sqlite3.connect(derby_db_path)
        # 1. Sweep the target table for all known identifiers
        cursor = conn.execute("SELECT name, hlc_database_type, cloud_region, hlc_server_platform FROM target")
        for name, db_type, region, platform in cursor.fetchall():
            if name: 
                registry["targets"].add(name.upper())
                registry["types"][name.upper()] = normalize_db_type(db_type)
            if db_type: registry["db_engines"].add(db_type.lower())
            if region: registry["regions"].add(region.lower())
            if platform: registry["platforms"].add(platform.lower())
            
        # 2. Sweep performance schema for digests
        cursor = conn.execute("SELECT DISTINCT digest FROM performance_schema WHERE digest IS NOT NULL")
        registry["digests"] = {row[0].upper() for row in cursor.fetchall()}
        
        conn.close()
    except Exception as e:
        print(f"⚠️ Registry build failed: {e}")
    return registry

GLOBAL_REGISTRY = build_universal_registry()


def extract_universal_context(question):
    q_low = question.lower()
    found = {
        "target": None, 
        "type": None, 
        "digests": [], 
        "filters": [] # List of (column, operator, value)
    }

    # 1. Deterministic Target & Type (The "common" link)
    for name in GLOBAL_REGISTRY["targets"]:
        if re.search(rf'\b{re.escape(name.lower())}\b', q_low):
            found["target"] = name
            found["type"] = GLOBAL_REGISTRY["types"].get(name)
            break

    # 2. Categorical Metadata Extraction (Region, Platform, Engine)
    # Mapping lexicon for matching
    meta_map = [
        (GLOBAL_REGISTRY["regions"], "cloud_region"),
        (GLOBAL_REGISTRY["platforms"], "hlc_server_platform"),
        (GLOBAL_REGISTRY["db_engines"], "hlc_database_type")
    ]

    for val_set, col_name in meta_map:
        for val in val_set:
            if re.search(rf'\b{re.escape(val)}\b', q_low):
                # Context-aware Negation Check
                match_idx = q_low.find(val)
                window = q_low[max(0, match_idx - 20):match_idx]
                is_neg = any(w in window for w in ['not', 'except', 'excluding', 'outside', 'neither'])
                
                op = "!=" if is_neg else "="
                # Store the filter to be injected into d_filters later
                found["filters"].append(('target', col_name, op, f"'{val.title() if col_name != 'cloud_region' else val}'"))

    # 3. Digest Extraction
    found["digests"] = [d for d in GLOBAL_REGISTRY["digests"] if d.lower() in q_low]
    if not found["digests"]:
        found["digests"] = re.findall(r'\b([a-zA-Z_]+--[a-zA-Z0-9]+)\b', question.upper())

    return found

# ---------------------------------------------------------# ---------------------------------------------------------
anchor_texts = []
anchor_intents = []

# Load standard FAQ intents
for intent_id, data in faq_intents.items():
    for anchor in data["anchors"]:
        anchor_texts.append(anchor)
        anchor_intents.append(intent_id)

# NEW: Load the DBA target intents
for intent_id, data in dba_intents.items():
    for anchor in data["anchors"]:
        anchor_texts.append(anchor)
        anchor_intents.append(intent_id)

anchor_embeddings = router_model.encode(anchor_texts, convert_to_tensor=True)




def semantic_metadata_router(user_question, embedding_model, anchor_embeddings, threshold_high=0.85, threshold_mid=0.55):
    q_emb = embedding_model.encode(user_question, convert_to_tensor=True)
    cosine_scores = util.cos_sim(q_emb, anchor_embeddings)[0]
    best_score_val, best_idx = torch.max(cosine_scores, dim=0)
    best_score = best_score_val.item()
    winning_intent_id = anchor_intents[best_idx]

    print(f"  [MASTER ROUTER] Intent={winning_intent_id} Score={best_score:.2f}")

    # Combine both dictionaries for a unified lookup
    all_intents = {**faq_intents, **dba_intents}
    
    if winning_intent_id not in all_intents or best_score < threshold_mid:
        return {"status": "PASS", "intent": None, "score": best_score}

    intent_data = all_intents[winning_intent_id]

    # --- 1. Handle Out of Scope ---
    if winning_intent_id == "OUT_OF_SCOPE":
        return {"status": "REJECT", "message": intent_data["clarification"]}

    # --- 2. Handle Remote DBA Diagnostic Queries ---
    if winning_intent_id in dba_intents:
        return {
            "status": "DBA_QUERY",
            "intent": winning_intent_id,
            "action_key": intent_data["intent_key"]
        }

    # --- 3. Handle Derby/Influx Metadata FAQs ---
    if best_score >= threshold_high:
        return {"status": "MATCH", "intent": winning_intent_id, "sql": intent_data["sql"], "target_db": intent_data["target_db"]}
    else:
        return {"status": "CLARIFY", "intent": winning_intent_id, "message": intent_data["clarification"], "sql": intent_data["sql"], "target_db": intent_data["target_db"]}


        



ROUTING_ANCHORS = {
    'influx_system': [
        "worst performing queries", "top anomalies", "performance bottlenecks",
        "average anomaly score", "spike in metrics", "time series data",
        "slowest queries", "highest anomaly", "performance issues",
        "bottleneck targets", "metric alerts", "query performance score",
        "execution anomalies", "worst sql performance", "top issues by score",
    ],
    'derby_system': [
        "list my targets", "show collectors", "cloud pricing", "report schedule",
        "dashboard templates", "target configuration", "database type",
        "integration settings", "collector status", "credit cost",
        "execution plan for digest", "sql digest text", "target details",
        "which targets are stopped", "show me all databases",
        "targets on linux", "cloud hosted targets", "targets not in region",
        "targets running oracle", "targets that are stopped", "show monitored databases"
    ]
}

routing_anchor_texts = []
routing_anchor_labels = []
for db, anchors in ROUTING_ANCHORS.items():
    for anchor in anchors:
        routing_anchor_texts.append(anchor)
        routing_anchor_labels.append(db)

routing_anchor_embeddings = router_model.encode(
    routing_anchor_texts, convert_to_tensor=True
)






def route_query(question, influx_threshold=0.40, derby_threshold=0.40):
    # Cache result on the question string to avoid redundant encode calls
    if not hasattr(route_query, '_cache'):
        route_query._cache = {}
    if question in route_query._cache:
        return route_query._cache[question]

    q_lower = question.lower()

    # 🛡️ NEW FIX: Derby-Exclusive Keyword Lock
    # If they explicitly ask for Derby-only concepts, bypass the router!
    derby_exclusive = ['report', 'reports', 'template', 'templates', 'dashboard', 'integration', 'collector', 'pricing']
    if any(re.search(rf'\b{w}\b', q_lower) for w in derby_exclusive):
        route_query._cache[question] = 'derby_system'
        return 'derby_system'

    # 🛡️ EXISTING FIX: Digest ID Router
    if re.search(r'\b[A-Za-z_]+--[A-Za-z0-9]+\b', question):
        if any(w in q_lower for w in ['score', 'anomaly', 'value', 'metric']):
            route_query._cache[question] = 'influx_system'
            return 'influx_system'
        if any(w in q_lower for w in ['sql', 'plan', 'text', 'query', 'digest']):
            route_query._cache[question] = 'derby_system'
            return 'derby_system'
    
    q_emb = router_model.encode(question, convert_to_tensor=True)
    scores = util.cos_sim(q_emb, routing_anchor_embeddings)[0]
    
    db_scores = {}
    for i, label in enumerate(routing_anchor_labels):
        db_scores[label] = max(db_scores.get(label, -1), float(scores[i]))
    
    influx_score = db_scores.get('influx_system', 0)
    derby_score = db_scores.get('derby_system', 0)
    
    print(f"[ROUTE SCORES] influx={influx_score:.3f} derby={derby_score:.3f}")
    
    influx_confident = influx_score >= influx_threshold
    derby_confident = derby_score >= derby_threshold
    
    if influx_confident and derby_confident:
        result = 'multi_step'
    elif influx_confident:
        result = 'influx_system'
    else:
        result = 'derby_system'
    
    route_query._cache[question] = result
    return result


def word_to_num(word):
    word_map = {
        'single': 1, 'one': 1, 'two': 2, 'three': 3, 'four': 4,
        'five': 5, 'six': 6, 'seven': 7, 'eight': 8, 'nine': 9,
        'ten': 10, 'dozen': 12
    }
    return word_map.get(word.lower(), None)



# =========================================================
# DIALOGUE STATE TRACKER (DST)
# =========================================================


class DialogueStateTracker:
    def __init__(self):
        self.state = {
            "active_target": None,     
            "active_db_type": None,    
            "last_intent": None,
            "neural_tables": [],       
            "neural_filters": [],
            "target_schema_cache": {},
            # --- NEW FIELDS ---
            "last_table_columns": [],
            "last_table_results": []
        }
        self.state["pending_clarification"] = None  

    # --- NEW METHOD ---
    def update_last_results(self, columns, results):
        """Saves the last executed query results as a list of dictionaries."""
        self.state["last_table_columns"] = columns
        # Zip columns and rows together so we can easily look up by column name later
        self.state["last_table_results"] = [dict(zip(columns, row)) for row in results]


    # In DialogueStateTracker
    def set_pending_clarification(self, intent_id, sql, target_db, original_question=None):
        self.state["pending_clarification"] = {
            "intent_id": intent_id, "sql": sql, "target_db": target_db,
            "original_question": original_question  # ← store it
        }

    def set_skip_faq(self):
        self.state["skip_faq"] = True

    def clear_skip_faq(self):
        self.state["skip_faq"] = False

    def should_skip_faq(self):
        return self.state.get("skip_faq", False)

    def clear_clarification(self):
        self.state["pending_clarification"] = None

    def update_state(self, target_name=None, db_type=None, intent=None):
        if target_name is not None: self.state["active_target"] = target_name
        if db_type is not None: self.state["active_db_type"] = db_type
        if intent is not None: self.state["last_intent"] = intent

    def clear_target_context(self):
        # The dedicated kill switch
        self.state["active_target"] = None
        self.state["active_db_type"] = None
        # THE FIX: Wipe neural memory so old filters don't haunt new targets
        self.state["neural_tables"] = []
        self.state["neural_filters"] = []

    def update_neural_context(self, tables, filters):
        self.state["neural_tables"] = tables
        self.state["neural_filters"] = filters

    def cache_schema(self, target_name, schema_graph):
        self.state["target_schema_cache"][target_name] = schema_graph

    def get_state(self):
        return self.state




# Initialize this in your main app execution block (e.g., Gradio/Streamlit state)
chat_memory = DialogueStateTracker()






import json
import sqlite3
import re


def format_results_as_md_table(cursor, results):
    if not results:
        return "_No data found for this query._"
    
    # Get all column names from the cursor
    all_col_names = [desc[0] for desc in cursor.description]
    
    # Check if this was a "wide" result (like from a SELECT *)
    # We truncate to the first 4 columns to prevent information overload
    display_limit = 4
    col_names = all_col_names[:display_limit]
    
    # Add a note if we truncated columns
    truncation_note = ""
    if len(all_col_names) > display_limit:
        truncation_note = f"\n\n*> Note: Showing first {display_limit} of {len(all_col_names)} columns.*"

    # Handle the Single Record Detail View
    if len(results) == 1 and len(all_col_names) > 5:
        detail_view = "### 📋 Record Detail View\n\n| Property | Value |\n| :--- | :--- |\n"
        row = results[0]
        # For detail view, we still show all columns as it's vertical and readable
        for col, val in zip(all_col_names, row):
            detail_view += f"| **{col}** | {str(val).replace('|', '&#124;')} |\n"
        return detail_view

    header = "| # | " + " | ".join(col_names) + " |"
    separator = "|---|" + "|".join(["---" for _ in col_names]) + "|"
    
    rows = []
    for index, r in enumerate(results):
        display_row = r[:display_limit]
        # Inject the 1-based index at the start of the row
        rows.append(f"| {index + 1} | " + " | ".join([str(i).replace('|', '&#124;') for i in display_row]) + " |")
    
    table_md = f'<div style="overflow-x: auto;">\n\n' + "\n".join([header, separator] + rows) + "\n\n</div>"
    
    return table_md + truncation_note




def parse_composite_digest(digest_string):
    """
    Implements the 'schema_name--original_digest' rule from the training data.
    Returns (schema_name, original_digest)
    """
    if not digest_string or "--" not in digest_string:
        return "public", digest_string # Default fallback
    
    parts = digest_string.split("--", 1)
    schema = parts[0] if parts[0] != "NULL" else "public"
    original_digest = parts[1]
    
    return schema, original_digest




def _run_neural_only(question, engine_derby, engine_influx, chat_memory, debug=False):
    """Runs the neural DB query path directly, skipping all FAQ/router logic."""
    q_lower = question.lower()
    current_state = chat_memory.get_state()
    target_db = route_query(question)
    engine = engine_derby if target_db == 'derby_system' else engine_influx

    # 🛡️ THE FIX: Define ctx in this scope!
    ctx = extract_universal_context(question)
    
    INFLUX_VALID_COLUMNS = {'target', 'metric', 'data_value_double', 'data_value_long', 
                            'time', 'capture', 'is_demo_alert', 'confidencescore'}

    active_tables, col_hits, raw_filters = engine.extract_intent(question, debug=debug)
    active_tables, col_hits, new_filters = enrich_and_inject_nuance(
        question, active_tables, col_hits, raw_filters, engine
    )

    # 🛡️ SECOND FIX: Safely unpack the filters (handling both 4 and 5 tuple formats)
    for f in ctx["filters"]:
        if len(f) == 4:
            f_t, f_c, f_o, f_v = f
            f_i = 'EXCLUDE' if '!' in f_o else 'INCLUDE'
        else:
            f_t, f_c, f_o, f_v, f_i = f
            
        # If we are in Influx, only inject if the column is valid for Influx
        if target_db == 'influx_system' and f_c not in INFLUX_VALID_COLUMNS:
            continue
        if not any(exist[1] == f_c for exist in new_filters):
            new_filters.append((f_t, f_c, f_o, f_v, f_i))

    chat_memory.update_neural_context(active_tables, new_filters)

    hrm_result = generate_sql_hrm(
        question, engine, hrm_model, router_model, device, debug=debug,
        injected_tables=active_tables,
        injected_col_hits=col_hits,
        injected_filters=new_filters
    )
    sql_query = hrm_result["sql"]

    exec_plan = f"### 🛠️ Generated Execution Plan\n* **Step 1** (Target: `{target_db}`):\n  ```sql\n  {sql_query}\n  ```\n"
    db_path = f"./{target_db}/{target_db}.sqlite"
    try:
        conn = sqlite3.connect(db_path)
        cursor = conn.cursor()
        cursor.execute(sql_query)
        results = cursor.fetchall()
        
        col_names = [desc[0] for desc in cursor.description]
        chat_memory.update_last_results(col_names, results)
        
        formatted_results = format_results_as_md_table(cursor, results)
        conn.close()
    except Exception as e:
        formatted_results = f"🚨 **SQL EXECUTION ERROR:**\n```\n{e}\n```"

    return f"{exec_plan}\n### 📊 Execution Results:\n{formatted_results}"




import os
import re
import json
import sqlite3
import torch
from sentence_transformers import util


def resolve_shorthand(question, chat_memory):
    state = chat_memory.get_state()
    last_results = state.get("last_table_results", [])
    last_tables = state.get("neural_tables", [])
    last_cols = state.get("last_table_columns", [])
    
    if not last_results:
        return question 
        
    match = re.search(r'(?:^|\s)(#|row\s+|entry\s+|number\s+|target\s+|digest\s+)(\d+)\b', question, re.IGNORECASE)
    if not match:
        return question
        
    target_idx = int(match.group(2)) - 1 
    
    if 0 <= target_idx < len(last_results):
        row_data = last_results[target_idx]
        id_columns = ['digest', 'metric', 'name', 'target', 'id', 'integration_name']
        resolved_value = None
        resolved_col = None
        
        for col in id_columns:
            if col in row_data:
                resolved_value = str(row_data[col])
                resolved_col = col
                break
                
        if not resolved_value and row_data:
            resolved_col = list(row_data.keys())[0]
            resolved_value = str(row_data[resolved_col])
            
        if resolved_value:
            # Fallback context inference if FAQ router bypassed neural state
            if not last_tables:
                if 'db_target_status' in last_cols: last_tables = ['target']
                elif 'platform' in last_cols and 'status' in last_cols: last_tables = ['dbactdc_instance']
            
            context_table = last_tables[0] if last_tables else "record"
            if context_table == 'sys_target_alerts': context_table = 'target'
            elif context_table == 'dbactdc_instance': context_table = 'collector'
            
            replacement = f"{context_table} {resolved_col} '{resolved_value}'"
            new_question = question[:match.start(1)] + replacement + question[match.end():]
            print(f"[SHORTHAND RESOLVER] Replaced '{match.group(0).strip()}' -> '{replacement}'")
            return new_question
            
    return question



BRIDGE_TRIGGER_ANCHORS = [
    "slowest queries", "worst performing targets", "top anomalies",
    "performance bottlenecks", "spike in metrics", "most severe issues",
    "targets with high anomaly scores", "slowest databases",
    "queries causing problems", "top issues by score",
    "worst sql performance", "targets not on linux",
    "cloud targets with anomalies", "non-demo targets with spikes"
]
bridge_anchor_embeddings = router_model.encode(BRIDGE_TRIGGER_ANCHORS, convert_to_tensor=True)

def needs_bridge_semantic(question, threshold=0.45):
    q_emb = router_model.encode(question, convert_to_tensor=True)
    scores = util.cos_sim(q_emb, bridge_anchor_embeddings)[0]
    best = float(scores.max())
    print(f"[BRIDGE SEMANTIC] score={best:.3f}")
    return best >= threshold




def run_multi_step_workflow(question, engine_derby, engine_influx, chat_memory, debug=True):
    question = resolve_shorthand(question, chat_memory)
    ctx = extract_universal_context(question)
    q_lower = question.lower()
    current_state = chat_memory.get_state()

    # =========================================================
    # ── EXPLICIT RAW EXECUTE COMMAND INTERCEPTOR ──
    # =========================================================
    # Catches: "execute <query> in <target>" OR "execute <query> in target name '<target>'"
    exec_match = re.search(r"(?i)\bexecute\s+([\s\S]+?)\s+in\s+(?:[a-zA-Z_]+\s+[a-zA-Z_]+\s+)?['\"]?([a-zA-Z0-9_-]+)['\"]?\s*$", question)
    
    if exec_match:
        sql_query = exec_match.group(1).strip().strip('\'"`') # Clean off any accidental quotes
        target_db_name = exec_match.group(2).strip().upper()
        
        if debug: print(f"[ORCHESTRATOR] ⚡ Action: EXPLICIT RAW EXECUTE intercepted for {target_db_name}")
        
        # Return exact requested output format
        return (f"**target db :** `{target_db_name}`\n\n"
                f"**Query :**\n```sql\n{sql_query}\n```\n\n"
                f"*(Note: Payload prepared for direct remote execution)*")

    
    INFLUX_VALID_COLUMNS = {'target', 'metric', 'data_value_double', 'data_value_long',
                            'time', 'capture', 'is_demo_alert', 'confidencescore'}
    BRIDGE_INVALID_COLUMNS = {
        'hlc_password', 'hlc_password2', 'hlc_pass_phrase', 'hlc_private_key_location',
        'hlc_is_private_key_location', 'hlc_database_auth_secret_name', 'hlc_username',
        'hlc_database_jmx', 'hlc_database_jmx_user', 'hlc_database_jmx_password',
        'hlc_database_jmx_password2', 'hlc_database_home', 'hlc_data_collector_home',
        'hlc_collector_configuration', 'hlc_environment_configuration', 'hlc_local_home',
        'hlc_local_configuration', 'hlc_logs_location', 'hlc_translation_source',
        'hlc_translation_destination', 'appdynamics_url', 'appdynamics_user_name',
        'appdynamics_account_name', 'appdynamics_account_password', 'newrelic_api_key',
        'datadog_api_key', 'datadog_application_key', 'cloud_pricing_lookup_code',
        'license', 'dtype', 'owner_id',
    }

    # =========================================================
    # 0. PENDING CLARIFICATION
    # =========================================================
    pending = current_state.get("pending_clarification")
    if pending:
        rejection_signals = ['no', 'nope', 'not that', "that's not", 'wrong', 'never mind']
        affirmation_signals = ['yes', 'yeah', 'correct', 'right', 'that', 'sure', 'go ahead']
        # if any(s in q_lower.split() or q_lower == s for s in rejection_signals):
        #     original_question = pending.get("original_question")
        #     chat_memory.clear_clarification()
        #     if original_question:
        #         # return _run_neural_only(original_question, engine_derby, engine_influx, chat_memory, debug)
        #         return run_multi_step_workflow(original_question, engine_derby, engine_influx, chat_memory, debug)
        if any(s in q_lower.split() or q_lower == s for s in rejection_signals):
            original_question = pending.get("original_question")
            chat_memory.clear_clarification()
            chat_memory.set_skip_faq()
            if original_question:
                return run_multi_step_workflow(original_question, engine_derby, engine_influx, chat_memory, debug)

        
        elif any(s in q_lower.split() or q_lower == s for s in affirmation_signals):
            chat_memory.clear_clarification()
            sql_query, target_db = pending["sql"], pending["target_db"]
            db_path = f"./{target_db}/{target_db}.sqlite"
            exec_plan = f"### 🛠️ Generated Execution Plan\n* **Step 1** (Target: `{target_db}`):\n ```sql\n {sql_query}\n ```\n"
            try:
                conn = sqlite3.connect(db_path)
                cursor = conn.cursor()
                cursor.execute(sql_query)
                results = cursor.fetchall()
                chat_memory.update_last_results([d[0] for d in cursor.description], results)
                formatted_results = format_results_as_md_table(cursor, results)
                conn.close()
            except Exception as e:
                formatted_results = f"🚨 **SQL EXECUTION ERROR:**\n```\n{e}\n```"
            return f"{exec_plan}\n### 📊 Execution Results:\n{formatted_results}"

    # =========================================================
    # 1. CONTEXT EXTRACTION
    # =========================================================
    extracted_name = ctx["target"]
    extracted_type = ctx["type"]
    extracted_digests = ctx["digests"]

    is_global_query = any(w in q_lower for w in [
        "all target", "every target", "all the target", "list all", "across all",
        "all of them", "system wide", "overall"
    ])

    if extracted_name:
        if extracted_name != current_state["active_target"]:
            if debug: print(f"[MEMORY] New target detected: {extracted_name}. Switching context.")
            chat_memory.clear_target_context()
            chat_memory.update_state(target_name=extracted_name, db_type=extracted_type)
        target_name, db_type = extracted_name, extracted_type
    elif is_global_query:
        chat_memory.clear_target_context()
        target_name, db_type = None, None
    else:
        target_name = current_state["active_target"]
        db_type = current_state["active_db_type"]
        if debug and target_name: print(f"[MEMORY] Recalled target: {target_name}")

    # =========================================================
    # 2. MASTER GATEKEEPER (FAQ / DBA / REJECT)
    # =========================================================
    # route = semantic_metadata_router(question, router_model, anchor_embeddings)

    if chat_memory.should_skip_faq():
        chat_memory.clear_skip_faq()
        route = {"status": "PASS"}
    else:
        route = semantic_metadata_router(question, router_model, anchor_embeddings)


    if route["status"] == "REJECT":
        return f"🤖 {route['message']}"

    if route["status"] == "DBA_QUERY":

        anomaly_keywords = ['anomal', 'spike', 'bottleneck', 'worst', 'top', 'slow', 'performance', 'severe']
        if any(kw in question.lower() for kw in anomaly_keywords):
        # This is a performance analysis question, not a DBA dialect query — pass through
            pass  # fall through to routing/multi-step logic
        else:
            if not target_name or not db_type:
                formatted_intent = route['intent'].replace('_', ' ').lower()
                return f"🤖 You asked to `{formatted_intent}`, but I don't know which database you want to check. Could you specify the target name?"
        dialect_sql = TARGET_DIALECTS.get(route["action_key"], {}).get(db_type)
        if not dialect_sql:
            return f"🤖 No `{route['action_key']}` dialect for `{db_type}` yet."
        if route["action_key"] in ["SHOW_COLUMNS", "SHOW_INDEXES"]:
            table_match = re.search(r"\btable\s+['\"]?([a-zA-Z0-9_]+)['\"]?", question, re.IGNORECASE)
            if table_match:
                tbl = table_match.group(1)
                t_col = {"postgres": "relname" if route["action_key"] == "SHOW_COLUMNS" else "tablename",
                         "db2": "tabname", "redshift": "tablename",
                         "sqlserver": "object_name(object_id)",
                         "sybase": "object_name(id)", "clickhouse": "table"}.get(db_type, "table_name")
                fuzzy = f" LOWER({t_col}) LIKE LOWER('%{tbl}%')"
                dialect_sql = dialect_sql.replace(";", f" AND{fuzzy};" if "WHERE" in dialect_sql.upper() else f" WHERE{fuzzy};")
        return (f"### 🛠️ Generated Target Payload\n* **Target:** `{target_name}`\n"
                f"* **Engine:** `{db_type.upper()}`\n* **Action:** `{route['action_key']}`\n"
                f" ```sql\n {dialect_sql}\n ```\n*Note: Remote query payload ready for dispatch.*")

    if route["status"] == "CLARIFY":
        chat_memory.set_pending_clarification(route["intent"], route["sql"], route["target_db"], original_question=question)
        return f"🤖 {route['message']}"

    if route["status"] == "MATCH":
        sql_query = route["sql"]
        faq_db = route["target_db"]
        if route["intent"] == "TARGET_DETAILS" and target_name:
            sql_query = f"SELECT * FROM target WHERE name = '{target_name}';"
        exec_plan = f"### 🛠️ Generated Execution Plan\n* **Step 1** (Target: `{faq_db}`):\n ```sql\n {sql_query}\n ```\n"
        try:
            conn = sqlite3.connect(f"./{faq_db}/{faq_db}.sqlite")
            cursor = conn.cursor()
            cursor.execute(sql_query)
            results = cursor.fetchall()
            chat_memory.update_last_results([d[0] for d in cursor.description], results)
            formatted_results = "\n".join([f"- **{r[0]}**" for r in results if r[0] != 'sqlite_sequence']) \
                if "sqlite_master" in sql_query else format_results_as_md_table(cursor, results)
            conn.close()
        except Exception as e:
            formatted_results = f"🚨 **SQL EXECUTION ERROR:**\n```\n{e}\n```"
        return f"{exec_plan}\n### 📊 Execution Results:\n{formatted_results}"

    # =========================================================
    # 3. ROUTING DECISION
    # =========================================================
    routing_decision = route_query(question)
    aggs = engine_influx.detect_aggregation(question)
    needs_sql_lookup = bool(re.search(r'\b(sql|quer(?:y|ies)|execution plans?|digest|what query|what sql)\b', q_lower))
    needs_anomalies = any(kw in q_lower for kw in [
        'anomal', 'spike', 'bottleneck', 'worst', 'top', 'performance', 'severe', 'issue', 'slow'
    ])

    trigger_multistep = (
        (routing_decision == 'multi_step' and needs_anomalies and not aggs)
        or (needs_sql_lookup and needs_anomalies)
        or (needs_sql_lookup and bool(target_name) and not is_global_query)
    )

    if trigger_multistep:
        target_db = 'influx_system'
    elif routing_decision == 'influx_system' or (routing_decision == 'multi_step' and needs_anomalies):
        target_db = 'influx_system'
    else:
        target_db = 'derby_system'

    engine = engine_derby if target_db == 'derby_system' else engine_influx

    # =========================================================
    # 4. BRIDGE — Scan Derby for metadata filters
    # =========================================================
    bridge_execution_step = ""
    injected_bridge_filters = []

    if trigger_multistep or needs_bridge_semantic(question):
        d_tables, d_cols, raw_d_filters = engine_derby.extract_intent(question, debug=False)
        _, _, d_filters = enrich_and_inject_nuance(question, d_tables, d_cols, raw_d_filters, engine_derby)

        # Inject universal deterministic filters
        for f_t, f_c, f_o, f_v in ctx["filters"]:
            if not any(exist[1] == f_c for exist in d_filters):
                d_filters.append((f_t, f_c, f_o, f_v, 'EXCLUDE' if '!' in f_o else 'INCLUDE'))

        # First operator on each column is authoritative
        col_first_op = {}
        for ft, fc, fop, fval, fint in d_filters:
            if fc not in col_first_op:
                col_first_op[fc] = (fop, fint)
        d_filters = [(ft, fc, col_first_op[fc][0], fval, col_first_op[fc][1])
                     for ft, fc, fop, fval, fint in d_filters]
        d_filters = [f for f in d_filters if f[1] not in BRIDGE_INVALID_COLUMNS]

        # Build bridge SQL from target-table filters only
        from collections import defaultdict
        col_groups = defaultdict(list)
        for ft, fc, fop, fval, fint in d_filters:
            if ft == 'target':
                col_groups[fc].append((fop, fval))

        derby_conditions = []
        for col, entries in col_groups.items():
            excludes = [v for op, v in entries if op == '!=']
            includes = [v for op, v in entries if op == '=']
            if len(excludes) > 1: derby_conditions.append(f"{col} NOT IN ({', '.join(excludes)})")
            elif excludes:        derby_conditions.append(f"{col} != {excludes[0]}")
            if len(includes) > 1: derby_conditions.append(f"{col} IN ({', '.join(includes)})")
            elif includes:        derby_conditions.append(f"{col} = {includes[0]}")

        # # Dedup by column+operator key
        # seen_keys, derby_conditions = set(), [
        #     cond for cond in derby_conditions
        #     if (key := ' '.join(cond.split()[:2])) not in seen_keys and not seen_keys.add(key)
        # ]

        seen_keys = set()
        final_conditions = []
        for cond in derby_conditions:
            key = ' '.join(cond.split()[:2])
            if key not in seen_keys:
                seen_keys.add(key)
                final_conditions.append(cond)
        derby_conditions = final_conditions

        if derby_conditions:
            bridge_sql = f"SELECT name FROM target WHERE {' AND '.join(derby_conditions)}"
            if debug: print(f"[BRIDGE SQL] {bridge_sql}")
            try:
                conn_d = sqlite3.connect("./derby_system/derby_system.sqlite")
                cross_targets = [r[0] for r in conn_d.execute(bridge_sql).fetchall()]
                conn_d.close()
                if not cross_targets:
                    return f"🤖 **Zero Results:** No targets matched your filters.\n\n*(Bridge Query: `{bridge_sql}`)*"
                target_list_str = ", ".join([f"'{t}'" for t in cross_targets])
                injected_bridge_filters.append(('sys_target_alerts', 'target', 'IN', f"({target_list_str})", 'INCLUDE'))
                bridge_execution_step = (f"* **Step 1 (Cross-DB Bridge):** `derby_system`\n"
                                         f" *(Resolving metadata filters)*\n ```sql\n {bridge_sql}\n ```\n"
                                         f"→ Found {len(cross_targets)} qualifying targets\n")
            except Exception as e:
                if debug: print(f"[BRIDGE ERROR] {e}")

    # =========================================================
    # 5. MULTI-STEP PIPELINE
    # =========================================================
    if trigger_multistep:
        if debug: print(f"[ORCHESTRATOR] ⚡ MULTI-STEP")
        chat_memory.update_state(intent="MULTI_STEP_PERFORMANCE")
        execution_steps = []

        inf_tables, inf_cols, inf_filters = engine_influx.extract_intent(question, debug=debug)

        if target_name and not is_global_query:
            inf_filters = [f for f in inf_filters if f[1] != 'target']
            inf_filters.append(('sys_target_alerts', 'target', '=', f"'{target_name}'", 'INCLUDE'))
            if 'sys_target_alerts' not in inf_tables: inf_tables.append('sys_target_alerts')
        elif injected_bridge_filters:
            inf_filters.extend(injected_bridge_filters)
            if 'sys_target_alerts' not in inf_tables: inf_tables.append('sys_target_alerts')

        if extracted_digests:
            inf_filters = [f for f in inf_filters if not any(d.lower() in str(f[3]).lower() for d in extracted_digests)]
            inf_filters += [('sys_target_alerts', 'metric', '=', f"'{d}'", 'INCLUDE') for d in extracted_digests]

        inf_filters = [f for f in inf_filters if f[1] in INFLUX_VALID_COLUMNS]

        influx_sql = generate_sql_hrm(
            question=question, engine=engine_influx, hrm_model=hrm_model,
            embedder=router_model, device=device, debug=debug,
            injected_tables=inf_tables, injected_col_hits=inf_cols,
            injected_filters=inf_filters, target_db='influx_system'
        )["sql"]

        step_num = 1
        if bridge_execution_step:
            execution_steps.append(bridge_execution_step.strip())
            step_num += 1
        execution_steps.append(f"* **Step {step_num}** (Target: `influx_system`):\n ```sql\n {influx_sql}\n ```")
        step_num += 1

        try:
            conn_in = sqlite3.connect("./influx_system/influx_system.sqlite")
            conn_in.row_factory = sqlite3.Row
            influx_results = conn_in.execute(influx_sql).fetchall()
            conn_in.close()
        except Exception as e:
            return f"🚨 InfluxDB Execution Error: {e}"

        digests, influx_data_map = [], []
        for row in influx_results:
            r = dict(row)
            d = r.get('metric') or r.get('METRIC')
            if d:
                digests.append(d)
                influx_data_map.append(r)

        if not digests:
            return "\n".join(execution_steps) + "\n\n**Result:** Found metrics but none linked to SQL digests in Derby."

        derby_text_sql = f"SELECT digest, digest_text, sql_plan_id FROM performance_schema WHERE digest IN ({', '.join(f'{chr(39)}{d}{chr(39)}' for d in digests)})"
        execution_steps.append(f"* **Step {step_num}** (Target: `derby_system`):\n *(Extracts digests from Step {step_num-1})*\n ```sql\n {derby_text_sql}\n ```")
        step_num += 1

        try:
            conn_derby = sqlite3.connect("./derby_system/derby_system.sqlite")
            conn_derby.row_factory = sqlite3.Row
            derby_text_results = conn_derby.execute(derby_text_sql).fetchall()
        except Exception as e:
            return f"🚨 Derby Query Error: {e}"

        sql_metadata, plan_ids_to_fetch = {}, set()
        for row in derby_text_results:
            r = dict(row)
            sql_metadata[r['digest']] = {"text": r['digest_text'], "plan_id": r['sql_plan_id']}
            if r['sql_plan_id']: plan_ids_to_fetch.add(str(r['sql_plan_id']))

        plan_metadata = {}
        if plan_ids_to_fetch:
            derby_plan_sql = f"SELECT id, sql_plain_text_plan FROM sql_plan WHERE id IN ({', '.join(plan_ids_to_fetch)})"
            execution_steps.append(f"* **Step {step_num}** (Target: `derby_system`):\n *(Extracts plan IDs from Step {step_num-1})*\n ```sql\n {derby_plan_sql}\n ```")
            try:
                for pr in [dict(r) for r in conn_derby.execute(derby_plan_sql).fetchall()]:
                    plan_metadata[str(pr['id'])] = pr['sql_plain_text_plan']
            except Exception:
                pass
        conn_derby.close()

        llm_payload = [{
            "target_server": r.get('target', 'unknown'),
            "digest": (d_id := r.get('metric') or r.get('METRIC')),
            "anomaly_score": r.get('data_value_double', 0),
            "timestamp": r.get('time'),
            "sql_query": sql_metadata.get(d_id, {}).get("text", "-- Metadata Not Found --"),
            "execution_plan": plan_metadata.get(str(sql_metadata.get(d_id, {}).get("plan_id", "")), "-- No Plan Linked --")
        } for r in influx_data_map]

        return (f"### 🛠️ Generated Execution Plan\n{chr(10).join(execution_steps)}\n\n"
                f"### 📦 Final JSON Payload (For LLM)\n```json\n{json.dumps(llm_payload, indent=2)}\n```")

    # =========================================================
    # 6. NEURAL FALLBACK
    # =========================================================
    if debug: print(f"[ORCHESTRATOR] ➡️ NEURAL FALLBACK.")

    active_tables, col_hits, raw_filters = engine.extract_intent(question, debug=debug)
    active_tables, col_hits, new_filters = enrich_and_inject_nuance(question, active_tables, col_hits, raw_filters, engine)

    if injected_bridge_filters:
        new_filters.extend(injected_bridge_filters)
        if target_db == 'influx_system' and 'sys_target_alerts' not in active_tables:
            active_tables.append('sys_target_alerts')

    # Scrubber
    table_names = {tbl.lower() for tbl in engine.schema.tables}
    all_col_names = {col.lower() for cols in engine.schema.tables.values() for col in cols}
    new_filters = [
        (ft, fc, fop, fval, fint) for ft, fc, fop, fval, fint in new_filters
        if not (
            (fval.replace("'","").lower() in table_names and fc not in {'name','target','hlc_server_name'}) or
            fval.replace("'","").lower() in all_col_names or
            (target_name and not is_global_query and fop != 'IN' and (
                (fc == ('target' if target_db == 'influx_system' else 'name') and fval.replace("'","").lower() != target_name.lower()) or
                (target_name.lower() in fval.replace("'","").lower() and fc != ('target' if target_db == 'influx_system' else 'name'))
            ))
        )
    ]

    # Force target filter
    if target_name and not is_global_query:
        target_col = 'target' if target_db == 'influx_system' else 'name'
        target_tbl = 'sys_target_alerts' if target_db == 'influx_system' else 'target'
        if not any(fc == target_col and target_name.lower() in fval.lower() for _, fc, _, fval, _ in new_filters):
            new_filters.append((target_tbl, target_col, '=', f"'{target_name}'", 'INCLUDE'))
        if target_tbl not in active_tables:
            active_tables.append(target_tbl)

    # Universal context filters
    for f_t, f_c, f_o, f_v in ctx["filters"]:
        if target_db == 'influx_system' and f_c not in INFLUX_VALID_COLUMNS: continue
        if not any(exist[1] == f_c for exist in new_filters):
            new_filters.append((f_t, f_c, f_o, f_v, 'INCLUDE' if f_o == '=' else 'EXCLUDE'))

    # Digest injection
    if extracted_digests:
        new_filters = [f for f in new_filters if not any(d.lower() in str(f[3]).lower() for d in extracted_digests)]
        for d in extracted_digests:
            tbl = 'performance_schema' if target_db == 'derby_system' else 'sys_target_alerts'
            col = 'digest' if target_db == 'derby_system' else 'metric'
            new_filters.append((tbl, col, '=', f"'{d}'", 'INCLUDE'))

    # Table pruning
    explicit_tables = {tbl for term, tbl in DOMAIN_LEXICON["table_mappings"].items()
                       if re.search(rf'\b{term}\b', q_lower)}

    if 'cloud_database_pricing' in active_tables:
        if not any(w in q_lower for w in ['cost', 'price', 'pricing', 'credit', 'paying']):
            active_tables.remove('cloud_database_pricing')
            new_filters = [f for f in new_filters if f[0] != 'cloud_database_pricing']
        else:
            if not any(f[0] == 'target' and f[1] not in ('hlc_server_type', 'is_on_cloud') for f in new_filters):
                active_tables = [t for t in active_tables if t != 'target']
                new_filters = [f for f in new_filters if not (f[0] == 'target' and f[1] in ('hlc_server_type', 'is_on_cloud'))]

    if 'target' in active_tables:
        for p in ['dbactdc_instance', 'cloud_database_pricing', 'integration']:
            if p in active_tables and p not in explicit_tables \
                    and not any(f[0] == p for f in new_filters) \
                    and not any(ch[0] == p and ch[2] > 1.5 for ch in col_hits):
                active_tables.remove(p)

    if 'target' in active_tables and 'target' not in explicit_tables:
        if not any(f[0] == 'target' and f[1] not in ('is_on_cloud', 'is_demo_target') for f in new_filters) \
                and len(active_tables) > 1:
            active_tables.remove('target')
            new_filters = [f for f in new_filters if f[0] != 'target']

    # Memory injection
    if re.search(r'\b(those|them|these|what about|how about|which of|out of)\b', q_lower) \
            and current_state["neural_filters"]:
        merged = {(f[0], f[1]): f for f in current_state["neural_filters"]}
        merged.update({(f[0], f[1]): f for f in new_filters})
        new_filters = list(merged.values())
        for t in current_state["neural_tables"]:
            if t not in active_tables: active_tables.append(t)

    chat_memory.update_neural_context(active_tables, new_filters)

    if target_db == 'influx_system':
        new_filters = [f for f in new_filters if f[1] in INFLUX_VALID_COLUMNS]

    sql_query = generate_sql_hrm(
        question, engine, hrm_model, router_model, device, debug=debug,
        injected_tables=active_tables, injected_col_hits=col_hits,
        injected_filters=new_filters, target_db=target_db
    )["sql"]

    exec_plan = "### 🛠️ Generated Execution Plan\n"
    exec_plan += bridge_execution_step + f"* **Step 2** (Target: `{target_db}`):\n *(Extracting metrics)*\n ```sql\n {sql_query}\n ```\n" \
        if bridge_execution_step else f"* **Step 1** (Target: `{target_db}`):\n ```sql\n {sql_query}\n ```\n"

    try:
        conn = sqlite3.connect(f"./{target_db}/{target_db}.sqlite")
        cursor = conn.cursor()
        cursor.execute(sql_query)
        results = cursor.fetchall()
        chat_memory.update_last_results([d[0] for d in cursor.description], results)
        formatted_results = format_results_as_md_table(cursor, results)
        conn.close()
    except Exception as e:
        formatted_results = f"🚨 **SQL EXECUTION ERROR:**\n```\n{e}\n```"

    return f"{exec_plan}\n### 📊 Execution Results:\n{formatted_results}"




import sqlite3
# =============================================================================
# 0. INITIALIZE ENGINES & SCHEMAS
# =============================================================================
def build_schema_from_sqlite(db_path):
    """Helper to dynamically generate the SchemaGraph from the SQLite files"""
    conn = sqlite3.connect(db_path)
    schema_text = ""
    for row in conn.execute("SELECT sql FROM sqlite_master WHERE type='table' AND sql IS NOT NULL"):
        schema_text += row[0] + ";\n"
    conn.close()
    return SchemaGraph(schema_text=schema_text)

print("🏗️ Building Schemas from SQLite...")
# derby_schema = build_schema_from_sqlite("/kaggle/working/spider_data/database/derby_system/derby_system.sqlite")
# influx_schema = build_schema_from_sqlite("/kaggle/working/spider_data/database/influx_system/influx_system.sqlite")
derby_schema = build_derby_schema()
influx_schema = build_influx_schema()

print("🔥 Initializing Linguistic Engines (Loading CrossEncoders)...")
# Note: If your CrossEncoder models are saved in a different path, update these!
TABLE_MODEL_PATH = './model_tables'
COLUMN_MODEL_PATH = './model_columns'
VALUE_MODEL_PATH = './model_values'

# Initialize the Derby Engine
engine_derby = LinguisticEngine(
    schema_graph=derby_schema,
    table_model_path=TABLE_MODEL_PATH,
    column_model_path=COLUMN_MODEL_PATH,
    value_model_path=VALUE_MODEL_PATH,
    skeleton_model_path=None, # We deleted the Skeleton model, so pass None
    db_id='derby_system'
)

# Initialize the Influx Engine
engine_influx = LinguisticEngine(
    schema_graph=influx_schema,
    table_model_path=TABLE_MODEL_PATH,
    column_model_path=COLUMN_MODEL_PATH,
    value_model_path=VALUE_MODEL_PATH,
    skeleton_model_path=None, 
    db_id='influx_system'
)

# ==============================================================
# 🚨 ADD THIS BLOCK: INITIALIZE THE HRM MODEL & DEVICE
# ==============================================================
print("🧠 Loading HRM Neural SQL Composer...")
device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')

# Instantiate the model architecture
hrm_model = HRMSQLComposer(embed_dim=384, hidden_dim=256, num_actions=9).to(device)

# Load the weights (⚠️ UPDATE THIS PATH to point to your actual HRM weights file)
HRM_WEIGHTS_PATH = 'hrm_model_v8_noise.pt' 

hrm_model.load_state_dict(torch.load(HRM_WEIGHTS_PATH, map_location=device))
hrm_model.eval()
# ==============================================================

print("✅ Engines initialized successfully! Ready for evaluation.")

print("✅ Engines initialized successfully! Ready for evaluation.")


#HRM