import sys
import os

import smtplib
from email.mime.text import MIMEText
from email.mime.multipart import MIMEMultipart
from email.mime.base import MIMEBase
from email import encoders

# Add project root to path for imports
project_root = os.path.dirname(os.path.dirname(os.path.dirname(os.path.abspath(__file__))))
if project_root not in sys.path:
    sys.path.insert(0, project_root)

import oracledb
from tavily import TavilyClient

# Core imports
import ast
import json
import re
import hashlib
import time
from typing import List, Optional, Dict, Any, Tuple, Union, Callable
from datetime import datetime, timedelta

import pandas as pd
from app.core.config import settings
from sqlalchemy.exc import OperationalError
from sqlalchemy.orm import Session

from app.core.database import SessionLocal, get_db_session
from app.models.question import TalkbackSessionState
from app.services.document_nodes import timer
from app.services.llm_service import LLMService, _qwen_configured
from askme_open_agent.exceptions import QueryTimeoutError


def _talkback_is_retryable_lock_error(exc: BaseException) -> bool:
    """SQLite busy/locked (matches logs.py intent; covers common driver messages)."""
    msg = str(exc).lower()
    return (
        "database is locked" in msg
        or "database table is locked" in msg
        or "sqlite_busy" in msg
        or "sqlite_locked" in msg
    )


def _talkback_db_execute(fn: Callable[[Session], Any]) -> Any:
    """
    Same SQLite handling as ``app/routes/logs.py`` (update-error-log): use ``SessionLocal()``,
    ``commit`` after work, ``rollback`` on failure, retry on busy/locked with exponential backoff.
    """
    max_retries = 5
    retry_delay = 0.1
    for attempt in range(max_retries):
        db = SessionLocal()
        try:
            result = fn(db)
            db.commit()
            return result
        except OperationalError as e:
            db.rollback()
            if _talkback_is_retryable_lock_error(e) and attempt < max_retries - 1:
                time.sleep(retry_delay * (2 ** attempt))
            else:
                raise
        except Exception:
            db.rollback()
            raise
        finally:
            db.close()


# Optional: Oracle database
try:
    import oracledb
    ORACLE_AVAILABLE = True
except ImportError:
    ORACLE_AVAILABLE = False
    oracledb = None

# Optional: Tavily client
try:
    from tavily import TavilyClient
    TAVILY_AVAILABLE = True
except ImportError:
    TAVILY_AVAILABLE = False
    TavilyClient = None

# Optional: Updated folder imports (these may have Streamlit dependencies)
UPDATED_MODULES_AVAILABLE = False
OpenAI_Chat = None
ChromaDB_VectorStore = None
APIChatService = None

try:
    from askme_open_agent.askme_openai.openai_chat import OpenAI_Chat
    from askme_open_agent.chromadb.chromadb_vector import ChromaDB_VectorStore
    from app.services.api_call import APIChatService
    UPDATED_MODULES_AVAILABLE = True
except ImportError as e:
    print(f"Warning: Could not import updated modules: {e}")
except Exception as e:
    print(f"Warning: Error importing updated modules: {e}")



# Define AskmeAgent class conditionally based on available modules
if UPDATED_MODULES_AVAILABLE and ChromaDB_VectorStore is not None and OpenAI_Chat is not None:
    class AskmeAgent(ChromaDB_VectorStore, OpenAI_Chat):
        def __init__(self, config=None):
            ChromaDB_VectorStore.__init__(self, config=config)
            OpenAI_Chat.__init__(self, config=config)
else:
    # Fallback stub class when updated modules are not available
    class AskmeAgent:
        def __init__(self, config=None):
            self.config = config or {}
            print("Warning: AskmeAgent running in stub mode - updated modules not available")
        
        def generate_sql(self, question, **kwargs):
            return None
        
        def run_sql(self, sql, **kwargs):
            return None
        
        def get_training_data(self):
            return None
        
        def connect_to_oracle(self, **kwargs):
            return None
        
        def connect_to_mysql(self, **kwargs):
            return None

        def connect_to_mssql(self, **kwargs):
            return None

        def generate_json_code_for_chart(self, question, columns, model=None):
            return ""

        def suggest_intent_and_agents(self, question, agent_profiles, **kwargs):
            return {
                "intent": "other",
                "suggested_agent_ids": [],
                "needs_clarification": True,
                "clarification_message": "AskmeAgent stub mode: routing unavailable.",
            }

        def merge_sql_query_summaries_into_details(self, response_details, user_question):
            """Stub: no LLM layer."""
            return None


class AskmeService:
    def __init__(self, model="gpt-4.1", chroma_path="./data/askme",traning_data = "",lib_path = "",open_ai ="",claude_api_key="",connection_string="",db_user="",db_password="", web_key ="", db_type="", host="", port="", dbname="", dsn="", user="", password="", max_row_count="1000", agent_ids=None):
        self._agent = None
        self._model = model
        self._chroma_path = chroma_path
        self._oracle_ready = False
        self._training_data = traning_data
        self._lib_path = lib_path
        self._open_ai = open_ai
        self._claude_api_key = claude_api_key or ""
        self._connection_string = connection_string or dsn
        self._db_user = db_user or user
        self._db_password = db_password or password
        self._web_key = web_key
        self._db_type = db_type
        self._host = host
        self._port = port
        self._dbname = dbname
        self._db_path = (dbname or "").strip() if dbname else ""
        self._api_chat_service = None
        self._api_chat_db_path = "api_chat.db"
        self._max_row_count = max_row_count
        self._agent_ids = agent_ids if isinstance(agent_ids, list) and len(agent_ids) > 1 else None

    def _parse_df_input(self, df_input):
        """
        Parse dataframe input which can be a list of dicts or a JSON string.
        Handles multiple formats:
        - List of row dicts (one dict per row)
        - List of multi-agent/API blocks: ``[{ "agent_id", "headers", "data": [rows...] }, ...]``
          — row lists from each block are concatenated (single-agent prior-``df`` follow-up)
        - Dict ``{"headers": [...], "data": [rows...]}`` (and optional ``agent_id``)
        - JSON string that parses to any of the above
        Returns a list of dicts or None if parsing fails.
        """
        if df_input is None:
            return None

        if isinstance(df_input, dict):
            data = df_input.get("data")
            if isinstance(data, list):
                return data
            return []

        if isinstance(df_input, list):
            if len(df_input) == 0:
                return []
            first = df_input[0]
            if (
                isinstance(first, dict)
                and "data" in first
                and isinstance(first.get("data"), list)
                and ("headers" in first or "agent_id" in first)
            ):
                rows_out: List[Any] = []
                for block in df_input:
                    if not isinstance(block, dict):
                        continue
                    part = block.get("data")
                    if isinstance(part, list):
                        rows_out.extend(part)
                return rows_out
            return df_input

        if isinstance(df_input, str):
            if not df_input.strip():
                return []
            try:
                parsed = json.loads(df_input)
                return self._parse_df_input(parsed)
            except json.JSONDecodeError as e:
                print(f"Error parsing DF JSON: {e}")
                return []
                
        print(f"Warning: Unexpected DF input type: {type(df_input)}")
        return []

    def _parse_multi_agent_df_blocks(self, df_input) -> List[Dict[str, Any]]:
        """
        Parse multi-agent/API dataframe input preserving per-agent blocks.

        Accepted shapes:
        - List[{"agent_id": "...", "headers": [...], "data": [rows...]}]
        - Dict {"agent_id": "...", "headers": [...], "data": [rows...]}
        - JSON string for either of the above

        Returns:
            List of normalized blocks:
            [{"agent_id": str, "headers": List[str], "data": List[dict]}]
        """
        if df_input is None:
            return []

        parsed = df_input
        if isinstance(df_input, str):
            s = df_input.strip()
            if not s:
                return []
            try:
                parsed = json.loads(s)
            except json.JSONDecodeError:
                return []

        if isinstance(parsed, dict):
            parsed = [parsed]
        if not isinstance(parsed, list):
            return []

        blocks: List[Dict[str, Any]] = []
        for idx, item in enumerate(parsed):
            if not isinstance(item, dict):
                continue
            data = item.get("data")
            if not isinstance(data, list):
                continue

            # Keep only row-like items to avoid noisy prompt input.
            rows = [r for r in data if isinstance(r, dict)]
            if not rows:
                continue

            raw_headers = item.get("headers")
            headers: List[str] = []
            if isinstance(raw_headers, list):
                headers = [str(h) for h in raw_headers if h is not None and str(h).strip()]
            if not headers:
                # Derive headers from rows when omitted.
                seen = set()
                derived: List[str] = []
                for row in rows:
                    for k in row.keys():
                        ks = str(k)
                        if ks not in seen:
                            seen.add(ks)
                            derived.append(ks)
                headers = derived

            blocks.append(
                {
                    "agent_id": str(item.get("agent_id") or item.get("source_agent") or f"block_{idx+1}"),
                    "headers": headers,
                    "data": rows,
                }
            )
        return blocks

    def _ensure_agent(self):
        if self._agent is None:
            cfg = {
                "api_key": self._open_ai,
                "claude_api_key": self._claude_api_key,
                "model": self._model,
                "path": self._chroma_path,
                "max_rows": self._max_row_count,
            }
            self._agent = AskmeAgent(config=cfg)
            if getattr(self._agent.get_training_data(), "empty", False):
                self.add_ddl(self.get_train_data())
        return self._agent

    def _ensure_sql(self):
        if self._db_type and self._db_type.lower() == "oracle":
            if not self._oracle_ready:
                print(self._lib_path)
                lib_dir = self._lib_path
                oracledb.init_oracle_client(lib_dir=lib_dir)
                self._oracle_ready = True
            a = self._ensure_agent()
            self._conn = a.connect_to_oracle(dsn=self._connection_string, user=self._db_user, password=self._db_password)
        elif self._db_type and self._db_type.lower() == "sql":
            a = self._ensure_agent()
            self._conn = a.connect_to_mysql(host=self._host, port=int(self._port), user=self._db_user, password=self._db_password, dbname=self._dbname)
        elif self._db_type and self._db_type.lower() == "mssql":
            a = self._ensure_agent()
            try:
                port_int = int(self._port) if self._port and str(self._port).strip().upper() != "NA" else 1433
            except (TypeError, ValueError):
                port_int = 1433
            self._conn = a.connect_to_mssql(
                server=self._host,
                port=port_int,
                user=self._db_user,
                password=self._db_password,
                dbname=self._dbname,
            )
        elif self._db_type and self._db_type.lower() == "sqlite":
            a = self._ensure_agent()
            db_path = (self._db_path or self._connection_string or "").strip()
            if not db_path:
                raise ValueError(
                    "SQLite agent requires db_name or connection (absolute path or sqlite:/// URL) in agent_management"
                )
            self._conn = a.connect_to_sqlite(db_path=db_path)

        


    # --- seed ---

    def get_train_data(self):
        return (
           self._training_data
        )

    def add_ddl(self, ddl):
        self._ensure_agent().train(ddl=ddl)

    # --- parity with your original helpers ---
    def generate_questions_cached(self, count):
        return self._ensure_agent().generate_init_question(self.get_train_data(), count=count)

    def generate_sql_cached(self, question, **kwargs):
        """Generate SQL - supports both single and multi-agent (multi-database) queries."""
        a = self._ensure_agent()
        
        # Multi-agent mode: per-agent training data, one SQL per agent, result format for LLM merge
        if self._agent_ids:
            from app.services.agent_config import get_agent_config

            agent_results = []
            for agent_id in self._agent_ids:
                config = get_agent_config(agent_id=agent_id)
                if not config:
                    agent_results.append({
                        "agent_id": agent_id,
                        "agent_name": str(agent_id),
                        "db_type": "unknown",
                        "config_snapshot": {},
                        "sql": "",
                        "execution_success": False,
                        "row_count": None,
                        "data_preview": None,
                        "error": "No config found for agent",
                    })
                    continue

                agent_name = config.get("agent_name", agent_id)
                db_type = config.get("db_type", "unknown")
                training_data = config.get("traning_data", "")

                # Config snapshot for LLM merge (full config keys; traning_data as length to avoid huge payload)
                config_snapshot = {
                    "agent_id": config.get("agent_id"),
                    "agent_name": config.get("agent_name"),
                    "db_type": config.get("db_type"),
                    "dbname": config.get("dbname", ""),
                    "host": config.get("host", ""),
                    "max_row_count": config.get("max_row_count", "1000"),
                    "traning_data_length": len(training_data or ""),
                }

                # Use this agent's dialect for generation
                dt = (db_type or "").lower()
                if dt == "oracle":
                    a.dialect = "ORACLE SQL"
                elif dt == "sqlite":
                    a.dialect = "SQLite"
                elif dt == "mssql":
                    a.dialect = "MSSQL"
                else:
                    a.dialect = "MYSQL"

                # Run pipeline with this agent's training_data only - no cross-agent context.
                # Explicitly pass empty question_sql_list and doc_list so ChromaDB (shared across agents)
                # does not inject examples/docs from other agents, which would cause join queries.
                ddl_list = [training_data] if training_data else []
                state = a.run_sql_generation_pipeline(
                    question,
                    max_retries=2,
                    ddl_list=ddl_list,
                    question_sql_list=[],
                    doc_list=[],
                )

                gr = state.generation_result if state else None
                sql = gr.sql if gr and gr.sql else ""
                execution_success = gr.execution_success if gr else False
                err_msg = (gr.error_message if gr else None) or (state.errors[-1] if state and state.errors else None)

                agent_results.append({
                    "agent_id": agent_id,
                    "agent_name": agent_name,
                    "db_type": db_type,
                    "config_snapshot": config_snapshot,
                    "sql": sql,
                    "execution_success": execution_success,
                    "row_count": None,
                    "data_preview": None,
                    "error": err_msg,
                })

            return {
                "question": question,
                "agent_results": agent_results,
                "merge_hint": "Use these per-agent results to produce a unified answer for the question.",
            }
        
        # Single agent mode (existing logic)
        db_type = (self._db_type or "sql").lower()
        if db_type == "oracle":
            a.dialect = "ORACLE SQL"
        elif db_type == "sqlite":
            a.dialect = "SQLite"
        elif db_type == "mssql":
            a.dialect = "MSSQL"
        else:
            a.dialect = "MySQL"
        return a.generate_sql(question=question, allow_llm_to_see_data=True, **kwargs)

    def generate_init_question(self, sql=None, count=5, model=None):
        """Generate initial business questions"""
        # If sql is provided, we might want to use it, but for now using training data logic as per original
        # Or if available in agent:
        return self._ensure_agent().generate_init_question(self.get_train_data(), count=count)

    def generate_question(self, sql, model=None):
        """Generate a question from SQL"""
        # Assuming agent has this method, or fallback
        if hasattr(self._ensure_agent(), 'generate_question'):
            return self._ensure_agent().generate_question(sql=sql)
        return "Question generation from SQL not supported"

    def is_sql_valid_cached(self, sql):
        return self._ensure_agent().is_sql_valid(sql=sql)

    def _resolve_step_placeholders(self, query_sql: str, step_results: list) -> str:
        """Replace {{STEP_N:column_name}} in query_sql with values from step N's DataFrame as SQL IN list."""
        import re
        import math

        def repl(match):
            step_idx = int(match.group(1))
            col_name = match.group(2)
            if step_idx < 0 or step_idx >= len(step_results) or step_results[step_idx] is None:
                return "(NULL)"
            df_step = step_results[step_idx]
            if col_name not in df_step.columns:
                return "(NULL)"
            values = df_step[col_name].dropna().unique().tolist()
            if not values:
                return "(NULL)"
            parts = []
            for v in values:
                if v is None or (isinstance(v, float) and math.isnan(v)):
                    continue
                if isinstance(v, (int, float)) and not isinstance(v, bool):
                    parts.append(str(v))
                else:
                    s = str(v).replace("'", "''")
                    parts.append(f"'{s}'")
            return "(" + ",".join(parts) + ")" if parts else "(NULL)"

        return re.sub(r"\{\{STEP_(\d+):(\w+)\}\}", repl, query_sql)

    @staticmethod
    def check_merge_data_loss(
        results: Dict[str, pd.DataFrame],
        merged: pd.DataFrame,
        merge_strategy: str,
        merge_key: Optional[str] = None,
    ) -> Dict[str, Any]:
        """
        Check whether the merged DataFrame preserves data from per-agent results.
        Returns a report dict with pre/post row counts, key coverage, and a data_loss flag.

        - concatenate: expected merged rows = sum(agent rows); loss if merged rows != that sum.
        - outer: every row from every agent should appear (by merge_key); loss if any key set is not fully present.
        - inner/left/right: inner/right can drop rows by design; left keeps first agent's rows.
        """
        report = {
            "merge_strategy": merge_strategy,
            "merge_key": merge_key,
            "pre_merge_rows_per_agent": {},
            "pre_merge_total_rows": 0,
            "merged_rows": len(merged) if merged is not None and hasattr(merged, "__len__") else 0,
            "merged_columns": list(merged.columns) if merged is not None and hasattr(merged, "columns") else [],
            "data_loss": False,
            "message": "",
        }
        if not results or merged is None:
            report["message"] = "No results or no merged DataFrame"
            return report

        for aid, df in results.items():
            n = len(df) if df is not None and hasattr(df, "__len__") else 0
            report["pre_merge_rows_per_agent"][aid] = n
            report["pre_merge_total_rows"] += n

        if merge_strategy == "concatenate":
            expected = report["pre_merge_total_rows"]
            actual = report["merged_rows"]
            report["expected_rows_after_merge"] = expected
            if actual != expected:
                report["data_loss"] = True
                report["message"] = f"Concatenate: expected {expected} rows, got {actual}"
            else:
                report["message"] = f"Concatenate: row count preserved ({actual} rows)"
            return report

        if merge_key and merge_key in merged.columns:
            # Key-based merge: check that keys from each agent appear in merged
            keys_in_merged = set(merged[merge_key].dropna().astype(str))
            for aid, df in results.items():
                if df is None or merge_key not in df.columns:
                    continue
                keys_in_agent = set(df[merge_key].dropna().astype(str))
                missing = keys_in_agent - keys_in_merged
                if missing and merge_strategy == "outer":
                    report["data_loss"] = True
                    report["message"] = f"Agent {aid}: {len(missing)} keys from source not in merged (outer merge should keep all)"
                    report[f"missing_keys_count_agent_{aid}"] = len(missing)
                    return report
                if missing and merge_strategy == "left" and aid == list(results.keys())[0]:
                    report["data_loss"] = True
                    report["message"] = f"Left merge: first agent {aid} has {len(missing)} keys not in merged"
                    return report
            report["message"] = "Key-based merge: key coverage checked; no unexpected loss"
        else:
            report["message"] = "No merge_key or key not in merged; cannot verify key coverage"
        return report

    def _agent_id_to_name_map(self, agent_ids: List[str]) -> Dict[str, str]:
        """Resolve agent_id to agent_name for use in merged column names. Returns dict agent_id -> safe label (agent_name, never raw agent_id)."""
        from app.services.agent_config import get_agent_config
        out = {}
        for aid in agent_ids:
            try:
                cfg = get_agent_config(agent_id=aid)
                name = None
                if cfg:
                    name = cfg.get("agent_name") or cfg.get("name") or None
                if not name or str(name).strip() == "" or str(name) == str(aid):
                    name = f"Agent_{aid}"
                # Sanitize for column names: replace spaces/special with underscore
                out[aid] = str(name).replace(" ", "_").replace("/", "_").replace("\\", "_").strip() or f"Agent_{aid}"
            except Exception:
                out[aid] = f"Agent_{aid}"
        return out

    def _apply_merge_info_to_results(
        self,
        results: Dict[str, pd.DataFrame],
        merge_info: dict,
    ) -> Union[pd.DataFrame, Dict[str, pd.DataFrame]]:
        """Apply merge strategy to per-agent results. Returns merged DataFrame or dict of DataFrames for independent."""
        merge_key = merge_info.get("merge_key", "id")
        merge_strategy = merge_info.get("merge_strategy", "outer")
        merge_steps = merge_info.get("merge_steps") or []
        agent_name_map = self._agent_id_to_name_map(list(results.keys())) if results else {}


        # Log per-agent results before merging (for debugging)
        print("[run_sql_cached] Results before merging:", {aid: (df.shape if df is not None and hasattr(df, "shape") else None) for aid, df in results.items()})
        for aid, df in results.items():
            if df is not None and hasattr(df, "head"):
                print(f"  agent {aid} preview:\n{df.head().to_string()}\n")

        # Handle independent queries: return multiple DataFrames (one per agent)
        if merge_strategy == "independent":
            if not results:
                return pd.DataFrame()
            return results

        # Handle concatenate strategy: stack rows vertically
        if merge_strategy == "concatenate":
            parts = []
            for agent_id, df in results.items():
                if df is None or (hasattr(df, 'empty') and df.empty):
                    continue
                df_copy = df.copy()
                if "source_agent" not in df_copy.columns:
                    df_copy["source_agent"] = agent_name_map.get(agent_id, str(agent_id))
                parts.append(df_copy)
            if not parts:
                return pd.DataFrame()
            merged = pd.concat(parts, axis=0, ignore_index=True)
            report = self.check_merge_data_loss(results, merged, "concatenate", merge_key=None)
            print("[merge_data_check]", report)
            if report.get("data_loss"):
                print("[merge_data_check] WARNING: possible data loss after concatenate merge")
            return merged

        # ADVANCED MERGE: Multi-step merging with different strategies
        if merge_steps and isinstance(merge_steps, list) and len(merge_steps) > 0:
            intermediate_results = {}
            for step_idx, step in enumerate(merge_steps):
                agent_list = step.get("agents", [])
                step_merge_key = step.get("merge_key")
                step_merge_strategy = step.get("merge_strategy", "outer")
                if len(agent_list) < 2:
                    print(f"Warning: merge_step {step_idx} has fewer than 2 agents, skipping")
                    continue
                dfs_to_merge = []
                agent_ids_in_step = []
                agent_labels_in_step = [] 

                for agent_ref in agent_list:
                    if str(agent_ref).startswith("result_"):
                        if agent_ref in intermediate_results:
                            dfs_to_merge.append(intermediate_results[agent_ref])
                            agent_ids_in_step.append(agent_ref)
                            agent_labels_in_step.append(agent_ref)
                        else:
                            print(f"Warning: intermediate result {agent_ref} not found")
                    else:
                        actual_agent_id = agent_ref[5:] if str(agent_ref).startswith("agent") else str(agent_ref)
                        if actual_agent_id in results:
                            dfs_to_merge.append(results[actual_agent_id])
                            agent_ids_in_step.append(actual_agent_id)
                            agent_labels_in_step.append(agent_name_map.get(actual_agent_id, actual_agent_id))
                        else:
                            print(f"Warning: agent {actual_agent_id} not found in results")
                if len(dfs_to_merge) < 2:
                    print(f"Warning: Not enough dataframes to merge in step {step_idx}")
                    continue
                step_merged = dfs_to_merge[0].copy()
                for i, df in enumerate(dfs_to_merge[1:], 1):
                    if df is None or (hasattr(df, 'empty') and df.empty):
                        continue
                    step_label = agent_labels_in_step[i] if i < len(agent_labels_in_step) else agent_ids_in_step[i]
                    try:
                        if step_merge_key and step_merge_key in step_merged.columns and step_merge_key in df.columns:
                            step_merged = pd.merge(
                                step_merged,
                                df,
                                on=step_merge_key,
                                how=step_merge_strategy,
                                suffixes=("", f"_{step_label}")
                            )
                        else:
                            df_renamed = df.copy()
                            for col in df.columns:
                                if col != step_merge_key:
                                    df_renamed.rename(columns={col: f"{step_label}_{col}"}, inplace=True)
                            step_merged = pd.concat([step_merged, df_renamed], axis=1)
                    except Exception as e:
                        print(f"Error in merge_step {step_idx}, merging with {step_label}: {e}. Returning data as separate DataFrames.")
                        return results
                result_name = f"result_{'_'.join(str(a).replace('agent', '') for a in agent_ids_in_step)}"
                intermediate_results[result_name] = step_merged
                print(f"[merge_step {step_idx}] Created {result_name} with shape {step_merged.shape}")
            if intermediate_results:
                final_result_key = list(intermediate_results.keys())[-1]
                merged = intermediate_results[final_result_key]
                report = self.check_merge_data_loss(
                    results, merged, "merge_steps", merge_key=merge_info.get("merge_key")
                )
                print("[merge_data_check]", report)
                if report.get("data_loss"):
                    print("[merge_data_check] WARNING: possible data loss after multi-step merge")
                return merged
            return pd.DataFrame()

        # SIMPLE MERGE: Single strategy for all agents
        merged_df = None
        for agent_id, df in results.items():
            if df is None or (hasattr(df, 'empty') and df.empty):
                continue
            label = agent_name_map.get(agent_id, str(agent_id))
            if merged_df is None:
                merged_df = df.copy()
                continue
            try:
                if merge_key and merge_key in merged_df.columns and merge_key in df.columns:
                    merged_df = pd.merge(
                        merged_df,
                        df,
                        on=merge_key,
                        how=merge_strategy,
                        suffixes=("", f"_{label}")
                    )
                else:
                    df_renamed = df.copy()
                    for col in df.columns:
                        if col != merge_key:
                            df_renamed.rename(columns={col: f"{label}_{col}"}, inplace=True)
                    merged_df = pd.concat([merged_df, df_renamed], axis=1)
            except Exception as e:
                print(f"Error merging agent {agent_id}: {e}. Returning data as separate DataFrames.")
                return results
        if merged_df is not None:
            report = self.check_merge_data_loss(results, merged_df, merge_strategy, merge_key=merge_key)
            print("[merge_data_check]", report)
            if report.get("data_loss"):
                print("[merge_data_check] WARNING: possible data loss after simple merge")
        return merged_df if merged_df is not None else pd.DataFrame()

    def run_sql_cached(self, sql, merge_info=None, question=None):
        """Execute SQL - supports single and multi-agent (independent and dependent) queries. merge_info is per-request to avoid races. If question is provided and merge_strategy is independent, LLM may be used to decide a better merge strategy from result shapes."""
        print("sql", sql)
        print("merge_info", merge_info)
        print("question", question)
        if self._agent_ids:
            if merge_info is None:
                parsed_merge = _merge_info_from_sql_blob(sql)
                if parsed_merge and parsed_merge.get("queries"):
                    merge_info = parsed_merge

            if merge_info is not None:
                queries = merge_info.get("queries", {})
                execution_order = merge_info.get("execution_order") or []
                merge_key = merge_info.get("merge_key", "id")
                merge_strategy = merge_info.get("merge_strategy", "outer")
                merge_steps = merge_info.get("merge_steps") or []

                if not isinstance(execution_order, list):
                    execution_order = []

                if not queries:
                    self._ensure_sql()
                    return self._ensure_agent().run_sql(sql=sql)

                from app.services.askme_service import get_askme_service
                results = {}

                if execution_order:
                    step_results = []
                    for step_idx, agent_key in enumerate(execution_order):
                        query = (queries.get(agent_key) or "").strip()
                        if not query:
                            step_results.append(None)
                            continue
                        actual_agent_id = agent_key[5:] if str(agent_key).startswith("agent") else str(agent_key)
                        resolved_sql = self._resolve_step_placeholders(query, step_results)
                        try:
                            svc = get_askme_service(agent_id=str(actual_agent_id))
                            if svc and svc != self:
                                df = svc.run_sql_cached(resolved_sql)
                                step_results.append(df)
                                if df is not None:
                                    results[actual_agent_id] = df
                            else:
                                step_results.append(None)
                        except Exception as e:
                            if isinstance(e, QueryTimeoutError):
                                raise
                            print(f"Error at step {step_idx} ({agent_key}): {e}")
                            step_results.append(None)
                else:
                    for agent_key, query in queries.items():
                        if not query or not str(query).strip():
                            continue
                        actual_agent_id = agent_key[5:] if str(agent_key).startswith("agent") else str(agent_key)
                        try:
                            svc = get_askme_service(agent_id=str(actual_agent_id))
                            if svc and svc != self:
                                df = svc.run_sql_cached(query)
                                if df is not None:
                                    results[actual_agent_id] = df
                        except Exception as e:
                            if isinstance(e, QueryTimeoutError):
                                raise
                            print(f"Error executing query for {agent_key}: {e}")

                if not results:
                    return pd.DataFrame()

                # Optionally ask LLM for merge strategy from result shapes when question provided and currently independent
                if question and merge_strategy == "independent":
                    result_shapes = {}
                    for aid, df in results.items():
                        if df is not None and hasattr(df, "columns"):
                            result_shapes[aid] = {
                                "columns": list(df.columns),
                                "row_count": len(df),
                                "dtypes": df.dtypes.astype(str).to_dict(),
                            }
                    if result_shapes and queries:
                        try:
                            agent = self._ensure_agent()
                            if hasattr(agent, "get_merge_strategy_from_result_shapes"):
                                llm_merge_info = agent.get_merge_strategy_from_result_shapes(
                                    question=question,
                                    queries=queries,
                                    result_shapes=result_shapes,
                                )
                                if llm_merge_info and llm_merge_info.get("merge_strategy") != "independent":
                                    merge_info = llm_merge_info
                                    merge_key = merge_info.get("merge_key", "id")
                                    merge_strategy = merge_info.get("merge_strategy", "outer")
                                    merge_steps = merge_info.get("merge_steps") or []
                        except Exception as e:
                            print(f"[run_sql_cached] LLM merge strategy failed, keeping independent: {e}")

                return self._apply_merge_info_to_results(results, merge_info)
        # Single-agent mode
        self._ensure_sql()
        return self._ensure_agent().run_sql(sql=sql)
    def should_generate_chart_cached(self,question, sql, df):
        return self._ensure_agent().should_generate_chart(df=df)

    def generate_plotly_code_cached(self, question, sql, df):
        return self._ensure_agent().generate_plotly_code(question=question, sql=sql, df=df)

    def generate_plot_cached(self, code, df):
        return self._ensure_agent().get_plotly_figure(plotly_code=code, df=df)

    def generate_json_code_cached(self, question, columns, model=None, use_qwen=None):
        """
        Generate Python code that creates a Plotly-compatible JSON structure from a DataFrame.
        Delegates to the agent's generate_json_code_for_chart (LLM call in AskmeBase).

        When *use_qwen* is omitted, uses OpenAI only (automatic charts / reconcile / API flows).
        Pass ``use_qwen=settings.default_use_qwen_for_summary_and_web()`` from explicit @askme
        analysis routes when scoped local LLM is desired.
        """
        if use_qwen is None:
            use_qwen = False
        return self._ensure_agent().generate_json_code_for_chart(
            question=question, columns=columns, model=model, use_qwen=use_qwen
        )

    def reconcile_and_merge(
        self,
        df: pd.DataFrame,
        excel_df: pd.DataFrame,
        question: str = "",
        do_df: bool = True,
        model: Optional[str] = None,
    ) -> Tuple[pd.DataFrame, Dict[str, Any]]:
        """Delegate to agent's reconcile_and_merge; pass instance max_row_count for full-data cap."""
        try:
            max_rows = int(self._max_row_count) if self._max_row_count is not None else None
        except (TypeError, ValueError):
            max_rows = None
        return self._ensure_agent().reconcile_and_merge(
            df, excel_df, question=question, do_df=do_df, model=model, max_rows=max_rows
        )

    def generate_followup_cached(self, question, sql, columns, use_qwen=None):
        if use_qwen is None:
            use_qwen = settings.default_use_qwen_for_summary_and_web()
        return self._ensure_agent().generate_followup_questions(
            question=question, sql=sql, columns=columns, use_qwen=use_qwen
        )

    def _norm_continuation_meta(self, val: Any, max_len: int = 8000) -> Optional[str]:
        """Stringify SQL/API metadata for continuation talkback; drop empty and literal 'N'."""
        if val is None:
            return None
        if isinstance(val, (dict, list)):
            s = json.dumps(val, default=str)
        else:
            s = str(val).strip()
        if not s or s.upper() == "N":
            return None
        if len(s) > max_len:
            return s[:max_len] + "...(truncated)"
        return s

    def generate_continuation_talkback_cached(
        self,
        question,
        answer_context,
        columns,
        n_items=3,
        use_qwen=None,
        df=None,
        sql=None,
        api_endpoint=None,
        api_params=None,
        api_response_validator=None,
    ):
        """
        Generate a structured continuation talkback with follow-up questions, next steps, and episodic context.
        When ``use_qwen`` is true and ``df`` is provided, suggestions are grounded in those rows.
        When ``use_qwen`` is omitted, defaults to ``settings.default_use_qwen_for_summary_and_web()`` so that
        local Qwen handles suggested questions when configured (per project routing rule).
        For OpenAI (``use_qwen`` false), suggestions use ``columns`` (and SQL/API metadata) only — no row payload.
        When ``sql`` or API fields are provided, suggestions align with how data was fetched.
        """
        if use_qwen is None:
            use_qwen = settings.default_use_qwen_for_summary_and_web()
        parsed_df = None
        if df is not None:
            if isinstance(df, list) and df and isinstance(df[0], pd.DataFrame):
                try:
                    parsed_df = pd.concat(df, ignore_index=True)
                except Exception:
                    parsed_df = df[0] if df else None
            elif hasattr(df, "to_dict") and hasattr(df, "columns"):
                try:
                    if not getattr(df, "empty", True):
                        parsed_df = df
                except Exception:
                    parsed_df = None
            else:
                data = self._parse_df_input(df)
                if data:
                    try:
                        parsed_df = pd.DataFrame(data)
                    except Exception:
                        parsed_df = None

        sql_s = self._norm_continuation_meta(sql, 8000)
        api_ep_s = self._norm_continuation_meta(api_endpoint, 6000)
        api_pr_s = self._norm_continuation_meta(api_params, 6000)
        api_val_s = self._norm_continuation_meta(api_response_validator, 6000)

        # Qwen/local: pass rows into the prompt when available; OpenAI/OIC: columns (+ sql/api) only.
        # Even when use_qwen=True, only inject rows when local Qwen is actually configured — otherwise
        # the call falls back to OIC/OpenAI and per project rule those LLMs must see columns only.
        df_for_continuation = parsed_df if (use_qwen and _qwen_configured()) else None

        with timer("talkback.generate_continuation_talkback_cached"):
            return self._ensure_agent().generate_continuation_talkback(
                question=question,
                answer_context=answer_context,
                columns=columns,
                n_items=n_items,
                df=df_for_continuation,
                sql=sql_s,
                api_endpoint=api_ep_s,
                api_params=api_pr_s,
                api_response_validator=api_val_s,
                use_qwen=use_qwen,
            )

    def merge_sql_query_summaries_into_details_cached(
        self, response_details: Dict[str, Any], user_question: str
    ) -> None:
        """
        After successful SQL execution, attach ``details.sql_query_summaries`` using the
        same agent LLM stack as ``generate_continuation_talkback_cached`` (AskmeBase.submit_prompt).
        """
        try:
            self._ensure_agent().merge_sql_query_summaries_into_details(
                response_details, user_question
            )
        except Exception as e:
            print(f"merge_sql_query_summaries_into_details_cached: {e}")

    def generate_summary_cached(self, question, columns, df=None, use_qwen=None):
        # Parse df if provided
        parsed_df = None
        if df is not None:
            data = self._parse_df_input(df)
            if data:
                try:
                    parsed_df = pd.DataFrame(data)
                except Exception as e:
                    print(f"Error creating DataFrame from parsed data: {e}")
                    parsed_df = None

        if use_qwen is None:
            use_qwen = settings.default_use_qwen_for_summary_and_web()
        # When the client omits or mis-sends `columns` but `df` parsed correctly, fill field names
        # so the LLM and downstream prompts are not "empty schema" with real rows.
        if parsed_df is not None and (not columns or (isinstance(columns, list) and len(columns) == 0)):
            try:
                columns = list(parsed_df.columns)
            except Exception:
                pass
        # Pass parsed rows for all backends; askme_base injects a bounded table into the system message
        # (not only the local Qwen path) so OIC/OpenAI and Qwen-fallback still summarize real data.
        df_for_summary = parsed_df
        return self._ensure_agent().generate_summary(
            question=question, columns=columns, df=df_for_summary, use_qwen=use_qwen
        )

    def generate_summary_multi_agent_cached(self, question, columns, df=None, use_qwen=None, model=None):
        """
        Generate a summary for multi-agent responses.
        """
        agent_data_blocks = self._parse_multi_agent_df_blocks(df)

        # Parse df if provided
        parsed_df = None
        if df is not None:
            data = self._parse_df_input(df)
            if data:
                try:
                    parsed_df = pd.DataFrame(data)
                except Exception as e:
                    print(f"Error creating DataFrame from parsed data: {e}")
                    parsed_df = None

        if use_qwen is None:
            use_qwen = settings.default_use_qwen_for_summary_and_web()
        # Per-agent and merged DataFrame samples: include in the prompt for every LLM so summaries
        # are not "no data" when the client only posts rows. Local Qwen still uses ``qwen_data`` in addition.
        df_for_summary = parsed_df
        blocks_for_summary = agent_data_blocks
        return self._ensure_agent().generate_summary_multi_agent(
            question=question,
            columns=columns,
            df=df_for_summary,
            use_qwen=use_qwen,
            model=model,
            agent_data_blocks=blocks_for_summary,
        )

    def generate_contextual_question(self, previous_questions, current_question, **kwargs):
        return self._ensure_agent().generate_contextual_question(previous_questions, current_question, **kwargs)

    def derive_talkback_session_key(
        self,
        user_id: str,
        agent_ids: Optional[List[str]] = None,
        client_session_key: Optional[str] = None,
    ) -> str:
        """
        Stable hashed key for ``talkback_session_state.session_key``.

        ``agent_ids`` is accepted for call-site compatibility but is not part of the hash
        (scope is stored on the row via ``agent_scope`` / upserts).

        Optional ``client_session_key`` (from talkback ``session_key`` form field) namespaces
        multiple concurrent threads per user; it is hashed together with ``user_id`` so the
        raw client value is never stored as the DB primary lookup key.
        """
        uid = str(user_id or "").strip() or "anonymous"
        extra = str(client_session_key or "").strip()
        if extra:
            raw = f"talkback_session::{uid}::{extra}"
        else:
            raw = f"talkback_session::{uid}"
        return hashlib.sha256(raw.encode("utf-8")).hexdigest()

    def get_talkback_session_state(self, session_key: str) -> Optional[Dict[str, Any]]:
        def _read(db: Session) -> Optional[Dict[str, Any]]:
            row = db.query(TalkbackSessionState).filter(TalkbackSessionState.session_key == session_key).first()
            if not row:
                return None
            return {
                "session_key": row.session_key,
                "user_id": row.user_id,
                "agent_scope": row.agent_scope,
                "history_json": row.history_json or "[]",
                "openai_response_id": row.openai_response_id,
                "pending_clarification": bool(row.pending_clarification),
                "clarification_prompt": row.clarification_prompt,
                "clarification_context_json": row.clarification_context_json,
            }

        return _talkback_db_execute(_read)

    def upsert_talkback_session_state(
        self,
        session_key: str,
        user_id: str,
        agent_ids: List[str],
        history: Optional[List[Dict[str, Any]]] = None,
        openai_response_id: Optional[str] = None,
        pending_clarification: Optional[bool] = None,
        clarification_prompt: Optional[str] = None,
        clarification_context: Optional[Dict[str, Any]] = None,
        agent_scope_ids: Optional[List[str]] = None,
    ) -> Dict[str, Any]:
        """
        agent_scope_ids: when set, persisted to ``agent_scope`` (typically latest ``dispatch_agent_ids``).
        When None, ``agent_ids`` is used for ``agent_scope`` (session / resolved scope).
        """

        def _upsert(db: Session) -> Dict[str, Any]:
            row = db.query(TalkbackSessionState).filter(TalkbackSessionState.session_key == session_key).first()
            scope_for_row = agent_scope_ids if agent_scope_ids is not None else agent_ids
            norm_scope = sorted(
                {str(a).strip() for a in (scope_for_row or []) if str(a).strip()}
            )
            scope_json = json.dumps(norm_scope)
            if not row:
                row = TalkbackSessionState(
                    session_key=session_key,
                    user_id=str(user_id),
                    agent_scope=scope_json,
                    history_json="[]",
                )
                db.add(row)
                db.flush()
            if history is not None:
                row.history_json = json.dumps(history, ensure_ascii=False, default=str)
            if openai_response_id is not None:
                row.openai_response_id = openai_response_id
            if pending_clarification is not None:
                row.pending_clarification = 1 if pending_clarification else 0
            if clarification_prompt is not None:
                row.clarification_prompt = clarification_prompt
            if clarification_context is not None:
                row.clarification_context_json = json.dumps(clarification_context, ensure_ascii=False, default=str)
            row.agent_scope = scope_json
            row.updated_at = datetime.utcnow()
            db.flush()
            return {
                "session_key": row.session_key,
                "history_json": row.history_json or "[]",
                "openai_response_id": row.openai_response_id,
                "pending_clarification": bool(row.pending_clarification),
                "clarification_prompt": row.clarification_prompt,
                "clarification_context_json": row.clarification_context_json,
                "agent_scope": row.agent_scope,
            }

        return _talkback_db_execute(_upsert)

    def clear_talkback_pending_clarification_only(self, session_key: str) -> None:
        """Clear pending_clarification and clarification fields without wiping history."""

        def _clear(db: Session) -> None:
            row = db.query(TalkbackSessionState).filter(TalkbackSessionState.session_key == session_key).first()
            if not row:
                return
            row.pending_clarification = 0
            row.clarification_prompt = None
            row.clarification_context_json = None
            row.updated_at = datetime.utcnow()
            db.flush()

        _talkback_db_execute(_clear)

    def clear_talkback_session_state(
        self, session_key: str, reset_agent_scope_ids: Optional[List[str]] = None
    ) -> None:
        """
        Clear talkback history and pending clarification.

        If ``reset_agent_scope_ids`` is provided, persist that list as ``agent_scope`` (typically the
        request's full authorized agent ids) so the column is not left blank and auto-routing can use it.

        If omitted, ``agent_scope`` is set to ``[]`` (legacy callers).
        """
        def _clear(db: Session) -> None:
            row = db.query(TalkbackSessionState).filter(TalkbackSessionState.session_key == session_key).first()
            if not row:
                return
            row.history_json = "[]"
            row.openai_response_id = None
            row.pending_clarification = 0
            row.clarification_prompt = None
            row.clarification_context_json = None
            if reset_agent_scope_ids is not None:
                norm = sorted(
                    {str(a).strip() for a in reset_agent_scope_ids if str(a).strip()}
                )
                row.agent_scope = json.dumps(norm)
            else:
                row.agent_scope = "[]"
            row.updated_at = datetime.utcnow()
            db.flush()

        _talkback_db_execute(_clear)

    def purge_expired_talkback_sessions(self, ttl_hours: int = 24) -> int:
        """
        Delete talkback session rows older than ttl_hours based on updated_at.
        Returns number of deleted rows.

        Deletes in batches with a separate commit per batch (via ``_talkback_db_execute``) to
        avoid holding a single long write lock on SQLite, same retry/backoff as other talkback ops.
        """
        cutoff = datetime.utcnow() - timedelta(hours=max(int(ttl_hours or 24), 1))
        batch_size = 200
        total = 0

        while True:
            def _purge_batch(db: Session) -> int:
                ids = [
                    row[0]
                    for row in (
                        db.query(TalkbackSessionState.id)
                        .filter(TalkbackSessionState.updated_at < cutoff)
                        .limit(batch_size)
                        .all()
                    )
                ]
                if not ids:
                    return 0
                deleted = (
                    db.query(TalkbackSessionState)
                    .filter(TalkbackSessionState.id.in_(ids))
                    .delete(synchronize_session=False)
                )
                db.flush()
                return int(deleted or 0)

            n = _talkback_db_execute(_purge_batch)
            total += n
            if n < batch_size:
                break

        return total

    def append_turn_to_talkback_session(
        self,
        session_key: str,
        user_id: str,
        agent_ids: List[str],
        question: str,
        answer: str,
        limit: int = 10,
        request_id: Optional[str] = None,
        answer_meta: Optional[Dict[str, Any]] = None,
        agent_scope_ids: Optional[List[str]] = None,
    ) -> Dict[str, Any]:
        """Single transaction: load row, idempotency check, append — reduces SQLite lock contention vs read+upsert."""

        def _append(db: Session) -> Dict[str, Any]:
            row = db.query(TalkbackSessionState).filter(TalkbackSessionState.session_key == session_key).first()
            history: List[Dict[str, Any]] = []
            if row and row.history_json:
                try:
                    parsed = json.loads(row.history_json or "[]")
                    if isinstance(parsed, list):
                        history = parsed
                except Exception:
                    history = []

            if request_id:
                rid = str(request_id).strip()
                if rid:
                    for t in history:
                        if isinstance(t, dict) and str(t.get("request_id", "")).strip() == rid:
                            # Same shape as ``get_talkback_session_state`` (idempotent replay).
                            if row:
                                return {
                                    "session_key": row.session_key,
                                    "user_id": row.user_id,
                                    "agent_scope": row.agent_scope,
                                    "history_json": row.history_json or "[]",
                                    "openai_response_id": row.openai_response_id,
                                    "pending_clarification": bool(row.pending_clarification),
                                    "clarification_prompt": row.clarification_prompt,
                                    "clarification_context_json": row.clarification_context_json,
                                }
                            break

            turn: Dict[str, Any] = {
                "question": question,
                "answer": answer,
                "ts": datetime.utcnow().isoformat(),
                "request_id": str(request_id).strip() if request_id else None,
            }
            if answer_meta:
                turn["answer_meta"] = answer_meta
            history.append(turn)
            history = history[-limit:]

            scope_for_row = agent_scope_ids if agent_scope_ids is not None else agent_ids
            norm_scope = sorted(
                {str(a).strip() for a in (scope_for_row or []) if str(a).strip()}
            )
            scope_json = json.dumps(norm_scope)
            if not row:
                row = TalkbackSessionState(
                    session_key=session_key,
                    user_id=str(user_id),
                    agent_scope=scope_json,
                    history_json=json.dumps(history, ensure_ascii=False, default=str),
                )
                db.add(row)
            else:
                row.history_json = json.dumps(history, ensure_ascii=False, default=str)
                row.agent_scope = scope_json
                row.updated_at = datetime.utcnow()
            db.flush()
            return {
                "session_key": row.session_key,
                "history_json": row.history_json or "[]",
                "openai_response_id": row.openai_response_id,
                "pending_clarification": bool(row.pending_clarification),
                "clarification_prompt": row.clarification_prompt,
                "clarification_context_json": row.clarification_context_json,
                "agent_scope": row.agent_scope,
            }

        return _talkback_db_execute(_append)

    def generate_forecast_cached(self, question, df):
        """Generate forecast dataframe based on user question and historical data"""
        return self._ensure_agent().generate_forecast(question=question, df=df)
    
    def ask_data_analysis_cached(self, question, df, use_qwen=None):
        """Generate flexible data analysis based on user question - can return text, dataframe, chart code, or combination"""
        if use_qwen is None:
            use_qwen = settings.default_use_qwen_for_summary_and_web()
        return self._ensure_agent().ask_data_analysis(question=question, df=df, use_qwen=use_qwen)
    
    def generate_error_recovery_questions(self, failed_question: str, max_questions: int = 3):
        """Generate validated questions when user's question fails"""
        self._ensure_sql()
        return self._ensure_agent().generate_error_recovery_questions(failed_question, max_questions)
        

    def generate_validated_followup_questions(self, question, sql, df, max_questions):
        """Generate validated followup questions - requires database connection for SQL testing"""
        self._ensure_sql()  # Ensure database connection is established
        return self._ensure_agent().generate_validated_followup_questions(question, sql, df, max_questions)
    
    def generate_multi_validated_followup_questions(self, question, sql_list, df_list, max_questions):
        """Generate validated followup questions for multiple SQL/DF pairs"""
        self._ensure_sql()
        return self._ensure_agent().generate_multi_validated_followup_questions(question, sql_list, df_list, max_questions)

    
    def get_traning_data(self):
        return self._ensure_agent().get_training_data()

    def add_traning_data(self, type, query, question=None):
        a = self._ensure_agent()
        if type == "ddl":
            a.train(ddl=query)
        elif type == "documentation":
            a.train(documentation=query)
        else:
            a.train(question=question, sql=query)

    def delete_traning_data(self, id):
        return self._ensure_agent().remove_training_data(id=id)

    def open_search(self, query):
        resp = self._ensure_agent().generate_normal_question(query=query)
        return resp

    def generate_normal_question(self, query, model=None):
        """Generate response for a general question"""
        agent = self._ensure_agent()
        kwargs = {}
        if model:
            kwargs['model'] = model
        return agent.generate_normal_question(query=query, **kwargs)

    def generate_rewritten_question(self, last_question, new_question, model=None):
        """Generate a rewritten question by combining the last question and the new question if they are related"""
        agent = self._ensure_agent()
        kwargs = {}
        if model:
            kwargs['model'] = model
        return agent.generate_rewritten_question(last_question=last_question, new_question=new_question, **kwargs)

    def rephrase_text_for_agent_classification_cached(self, text: str, model=None):
        """Rephrase a new-agent description (no training data) for unambiguous multi-agent routing."""
        agent = self._ensure_agent()
        kwargs: Dict[str, Any] = {}
        if model:
            kwargs["model"] = model
        return agent.rephrase_text_for_agent_classification(text, **kwargs)

    def generate_compare_data(self, query, df, web_data, model=None):
        """Compare data from database with web data and generate response"""
        agent = self._ensure_agent()
        kwargs = {}
        if model:
            kwargs['model'] = model
        # Convert df (list of dicts) to string representation for the agent
        if isinstance(df, list):
            df_str = pd.DataFrame(df).to_string()
        else:
            df_str = str(df)
        return agent.generate_compare_data(query=query, df=df_str, web_data=web_data, **kwargs)

    def _create_default_plotly_json(self, df, columns, question):
        """
        Create a default Plotly JSON structure from DataFrame.
        Helper function to ensure consistent Plotly format.
        """
        # Ensure df is a DataFrame
        if not isinstance(df, pd.DataFrame):
            df = pd.DataFrame(df)
        
        # Determine chart type based on number of columns
        if len(columns) == 1:
            # Single column - create bar chart with index as x
            chart_type = "bar"
            data = [{
                "x": list(range(len(df))),
                "y": df[columns[0]].tolist(),
                "type": chart_type,
                "name": columns[0]
            }]
            layout = {
                "title": question[:50] if len(question) > 50 else question,
                "xaxis": {"title": "Index"},
                "yaxis": {"title": columns[0]}
            }
        elif len(columns) >= 2:
            # Multiple columns - use first as x, second as y
            chart_type = "bar"
            data = [{
                "x": df[columns[0]].tolist(),
                "y": df[columns[1]].tolist(),
                "type": chart_type,
                "name": columns[1]
            }]
            layout = {
                "title": question[:50] if len(question) > 50 else question,
                "xaxis": {"title": columns[0]},
                "yaxis": {"title": columns[1]}
            }
        else:
            # No columns - empty chart
            data = [{"type": "bar", "x": [], "y": []}]
            layout = {"title": question[:50] if len(question) > 50 else question}
        
        return {"data": data, "layout": layout}

    def search_and_format_response(self, user_query, last_response, use_qwen=None):
        try:
            if use_qwen is None:
                use_qwen = settings.default_use_qwen_for_summary_and_web()

            # Gemini+search (if opted in): only user_query is sent — prior df is not sent to Gemini.
            use_gsearch = bool(getattr(settings, "USE_GEMINI_GOOGLE_SEARCH_WITH_QWEN", True))
            gkey = (getattr(settings, "GEMINI_API_KEY", None) or "").strip()
            resp = None
            if use_qwen and use_gsearch and gkey:
                g = self._ensure_agent().fetch_google_search_context_for_user_query(
                    user_query, gemini_api_key=gkey, gemini_model="gemini-3-flash-preview"
                )
                if (g or "").strip():
                    resp = (
                        "Web context (Google Search; only your question was sent to Gemini, not your table):\n\n"
                        + g.strip()
                    )

            if resp is None:
                tv = TavilyClient(api_key=self._web_key)
                data = tv.search(user_query, max_results=2, search_depth="advanced", language="en")
                resp = f"I found several results for **{data['query']}**. Here are the top options:\n\n"
                for i, r in enumerate(data.get("results", []), 1):
                    resp += f"**{i}. {r['title']}**\n{r['content']}\n🔗 {r['url']}\n\n"

            return self._ensure_agent().generate_compare_data(
                query=user_query, df=last_response, web_data=resp, use_qwen=use_qwen
            )
        except Exception as e:
            return {"formatted_response": f"Error: {e}", "raw_data": None, "success": False}

    def get_users_list(self, sql=None):
        if sql is None:
            sql = (
                "Select fu.user_id, fu.user_name, fu.email_address ,ppx.full_name "
                "from fnd_user fu , per_people_x ppx "
                "where nvl(fu.end_date , sysdate+1) > sysdate "
                "and fu.user_id > 1000 "
                "and ppx.person_id = fu.employee_id"
            )
        self._ensure_sql()
        return self._ensure_agent().run_sql(sql=sql)
    

    def _ensure_api_chat_service(self):
        """Ensure API Chat Service is initialized"""
        if self._api_chat_service is None:
            self._api_chat_service = APIChatService(
                openai_api_key=self._open_ai,
                claude_api_key=self._claude_api_key,
                db_path=self._api_chat_db_path
            )
        return self._api_chat_service
    
    def add_api(self, name, endpoint, base_params, response_schema, description=""):
        """Add a new API configuration"""
        service = self._ensure_api_chat_service()
        return service.add_api(name, endpoint, base_params, response_schema, description)
    
    def get_all_apis(self,agent_id=""):
        """Get all configured APIs"""
        service = self._ensure_api_chat_service()
        return service.get_all_apis(agent_id)
    
    def get_api_by_name(self, name, agent_id="" ):
        """Get API configuration by name"""
        service = self._ensure_api_chat_service()
        return service.get_api_by_name(name, agent_id)
    
    def update_api(self, api_id, name, endpoint, base_params, response_schema, description):
        """Update an existing API"""
        service = self._ensure_api_chat_service()
        return service.update_api(api_id, name, endpoint, base_params, response_schema, description)
    
    def delete_api(self, api_id):
        """Delete an API"""
        service = self._ensure_api_chat_service()
        return service.delete_api(api_id)
    
    def save_api_chat_message(self, session_id, role, content, api_used=None, 
                             params_used=None, raw_data=None, transformation_code=None):
        """Save a chat message to database"""
        service = self._ensure_api_chat_service()
        return service.save_chat_message(session_id, role, content, api_used, 
                                        params_used, raw_data, transformation_code)
    
    def get_api_chat_history(self, session_id):
        """Get chat history for a session"""
        service = self._ensure_api_chat_service()
        return service.get_chat_history(session_id)
    
    def clear_api_chat_history(self, session_id):
        """Clear chat history for a session"""
        service = self._ensure_api_chat_service()
        return service.clear_chat_history(session_id)
    
    def select_appropriate_api(self, user_query, available_apis=None,agent_id=""):
        """Use ChatGPT to select the most appropriate API for the user query"""
        service = self._ensure_api_chat_service()
        return service.select_appropriate_api(user_query, available_apis, agent_id)
    
    def extract_api_params(self, user_query, api_config):
        """Use ChatGPT to extract parameters from user query"""
        service = self._ensure_api_chat_service()
        return service.extract_api_params(user_query, api_config)
    
    def generate_transformation_code(self, user_query, response_schema):
        """Generate Python code to transform API response"""
        service = self._ensure_api_chat_service()
        return service.generate_transformation_code(user_query, response_schema)
    
    def execute_transformation(self, api_response, transformation_code):
        """Execute the generated transformation code"""
        service = self._ensure_api_chat_service()
        return service.execute_transformation(api_response, transformation_code)
    
    def call_api_and_transform(self, params, endpoint, response_schema,transformation_code, user_query):
        """Call the API and transform response using generated code"""
        service = self._ensure_api_chat_service()
        return service.call_api_and_transform(params, endpoint, response_schema, transformation_code,user_query)
    
    def process_api_user_query(self, user_query, session_id,agent_id=""):
        """
        Process a user query end-to-end: select API, extract params, call API, transform data
        
        Args:
            user_query: User's natural language query
            session_id: Session identifier for chat history
            
        Returns:
            Dict with keys: success, message, api_used, params, raw_data, df, code
        """
        service = self._ensure_api_chat_service()
        return service.process_user_query(user_query, session_id, agent_id)
    
    def process_sap_query(self, service_root, entity_set, user_query, response_schema, transformation_code,
                          username=None, password=None, verify_ssl=True, paginate=True, api_query_params=None):
        """
        Process SAP OData query end-to-end
        
        Args:
            service_root: SAP OData service root URL
            entity_set: Entity set name
            user_query: User's natural language query
            response_schema: Response schema definition
            transformation_code: Code to transform the response
            username: SAP username
            password: SAP password
            verify_ssl: Whether to verify SSL certificate
            paginate: Whether to follow pagination links
            api_query_params: Pre-extracted query parameters (optional)
            
        Returns:
            Dict with keys: success, message, df, query_params, url, code
        """
        service = self._ensure_api_chat_service()
        return service.process_sap_query(
            service_root=service_root,
            entity_set=entity_set,
            user_query=user_query,
            response_schema=response_schema,
            transformation_code=transformation_code,
            username=username,
            password=password,
            verify_ssl=verify_ssl,
            paginate=paginate,
            api_query_params=api_query_params
        )
    
    def generate_api_init_questions(self, api_name, count=5, agent_id=""):
        """
        Generate initial business questions based on API configuration
        
        Args:
            api_name: Name of the API to generate questions for
            count: Number of questions to generate (default: 5)
            
        Returns:
            List of business questions as strings
        """
        service = self._ensure_api_chat_service()
        return service.generate_init_questions_from_api(api_name, count, agent_id)
    
    def generate_all_api_init_questions(self, count_per_api=2,agent_id=""):
        """
        Generate initial questions for all configured APIs
        
        Args:
            count_per_api: Number of questions to generate per API (default: 3)
            
        Returns:
            Dictionary mapping API names to lists of questions
        """
        service = self._ensure_api_chat_service()
        a = self._ensure_agent()
        return service.generate_init_questions_from_all_apis(count_per_api, agent_id)
    
    def generate_questions_from_schema(self, response_schema, count=5):
        """
        Generate business questions based on API response schema/structure
        
        Args:
            response_schema: The structure/schema of API response
            count: Number of questions to generate
            
        Returns:
            List of business questions
        """
        service = self._ensure_api_chat_service()
        return service.generate_questions_from_response_schema(response_schema, count)
    
    # ==================== SAP OData Methods ====================
    
    def extract_sap_query_params(self, user_query, entity_set):
        """
        Use ChatGPT to extract SAP OData query parameters from user query
        
        Args:
            user_query: User's natural language query
            entity_set: SAP entity set name (e.g., SalesOrderSet)
            
        Returns:
            Dictionary of OData query parameters ($select, $filter, $orderby, etc.)
        
        Example:
            params = service.extract_sap_query_params(
                "Show me top 10 sales orders", 
                "A_SalesOrder"
            )
            # Returns: {"$top": "10", "$orderby": "CreatedOn desc"}
        """
        service = self._ensure_api_chat_service()
        return service.extract_sap_query_params(user_query, entity_set)
    
    def call_sap_odata_api(self, service_root, entity_set, user_query,
                           username=None, password=None,
                           verify_ssl=True, timeout=60,
                           paginate=True, max_pages=None):
        """
        Call SAP OData API based on user query, extract parameters, and return DataFrame
        
        Args:
            service_root: SAP OData service root URL 
                         (e.g., https://host/sap/opu/odata/sap/API_SALES_ORDER_SRV/)
            entity_set: Entity set name (e.g., A_SalesOrder, A_BusinessPartner)
            user_query: User's natural language query
            username: SAP username (if None, uses st.session_state.sap_username)
            password: SAP password (if None, uses st.session_state.sap_password)
            verify_ssl: Whether to verify SSL certificate
            timeout: Request timeout in seconds
            paginate: Whether to follow pagination links
            max_pages: Maximum number of pages to fetch (None = unlimited)
            
        Returns:
            Tuple of (DataFrame or None, status message)
        
        Example:
            # Set credentials in session
            st.session_state['sap_username'] = 'user'
            st.session_state['sap_password'] = 'pass'
            
            df, msg = service.call_sap_odata_api(
                service_root="https://host/sap/opu/odata/sap/API_SALES_ORDER_SRV/",
                entity_set="A_SalesOrder",
                user_query="Get top 20 sales orders with value over 5000"
            )
            
            if df is not None:
                st.dataframe(df)
            else:
                st.error(msg)
        """
        service = self._ensure_api_chat_service()
        return service.call_sap_odata_api(
            service_root=service_root,
            entity_set=entity_set,
            user_query=user_query,
            username=username,
            password=password,
            verify_ssl=verify_ssl,
            timeout=timeout,
            paginate=paginate,
            max_pages=max_pages
        )
    
    def process_sap_query(self, service_root, entity_set, user_query,response_schema,transformation_code,
                          username=None, password=None,
                          verify_ssl=True, paginate=True,api_query_params=None):
        """
        High-level method to process SAP OData query end-to-end
        Returns complete query information including URL, parameters, and DataFrame
        
        Args:
            service_root: SAP OData service root URL
            entity_set: Entity set name
            user_query: User's natural language query
            username: SAP username (optional, uses session if not provided)
            password: SAP password (optional, uses session if not provided)
            verify_ssl: Whether to verify SSL certificate
            paginate: Whether to follow pagination links
            
        Returns:
            Dict with keys: success, message, df, query_params, url
        
        Example:
            result = service.process_sap_query(
                service_root="https://host/sap/opu/odata/sap/API_BUSINESS_PARTNER/",
                entity_set="A_BusinessPartner",
                user_query="Show me customers in Germany with credit over 50000"
            )
            
            if result['success']:
                st.write(f"Generated URL: {result['url']}")
                st.write(f"Query Parameters: {result['query_params']}")
                st.dataframe(result['df'])
            else:
                st.error(result['message'])
        """
        service = self._ensure_api_chat_service()
        
        result = service.process_sap_query(
            service_root=service_root,
            entity_set=entity_set,
            user_query=user_query,
            response_schema=response_schema,
            transformation_code=transformation_code,
            username=username,
            password=password,
            verify_ssl=verify_ssl,
            paginate=paginate,
            max_row_count ="1000",
            api_query_params=api_query_params
        )
       
        return result
    
    def close_api_chat_service(self):
        """Close API chat service database connection"""
        if self._api_chat_service:
            self._api_chat_service.close()


    def generate_initial_questions(self, initial_training, num_questions: int, **kwargs) -> list:
        """
        Generate initial questions from training data
 
        Args:
            initial_training: Initial training data
            num_questions: Number of questions to generate
 
        Returns:
            List of generated questions
        """
        self._ensure_sql()        
        a = self._ensure_agent()
        return a.generate_initial_question_with_run(initial_training= initial_training,num_questions=num_questions, **kwargs)
    
    def generate_sql_for_question(self, question: str, max_retries: int = 1, **kwargs) -> dict:
        """
        Generate and validate SQL for a single user-provided question.
        Uses LLM to generate SQL, validates by executing with LIMIT 1,
        retries up to max_retries times. If still fails, rephrases the
        question via LLM and retries again with max_retries.
        
        Args:
            question: The question text to generate SQL for
            max_retries: Number of retries per phase if SQL validation fails (default: 1)
        
        Returns:
            dict with 'question', 'rephrased_question', 'sql', 'success', and optionally 'error'
        """
        self._ensure_sql()
        a = self._ensure_agent()
        return a.generate_sql_for_question_with_rephrase(
            question=question, max_retries=max_retries, **kwargs
        )


    @staticmethod
    def _row_matches_suggestion_datasource(
        question_data_source: Any, datasource_filter: Optional[str]
    ) -> bool:
        if not datasource_filter:
            return True
        t = re.sub(
            r"\s+",
            " ",
            (str(question_data_source) if question_data_source is not None else "")
            .replace("\ufeff", "")
            .strip()
            .lower(),
        )
        d = str(datasource_filter).lower()
        if d == "api":
            if t in ("api", "oic"):
                return True
            if re.search(r"(^|[^a-z0-9_])api([^a-z0-9_]|$)", t):
                return True
            return t == "a_p_i"
        if d == "database":
            if t in ("database", "db", "sql"):
                return True
            return re.search(
                r"(^|[^a-z0-9_])(database|db|sql)([^a-z0-9_]|$)", t
            ) is not None
        return True

    def get_agent_questions_from_db(
        self, agent_id: str, limit: int = 5, datasource_filter: Optional[str] = None
    ) -> Tuple:
        """
        Get initial suggested questions for an agent from `initial_suggestion_question`.

        datasource_filter: when ``"api"`` or ``"database"``, only matching
        ``question_data_source`` rows (BOM/whitespace tolerant, ``oic``, word-boundary ``api``).
        None preserves legacy behavior (no filter).

        Intentionally does *not* call ``_ensure_agent()`` so Chroma/embeddings cannot
        make this read-only path fail and return an empty list.
        """
        try:
            from app.core.database import SessionLocal
            from app.models.question import InitialSuggestionQuestion

            db = SessionLocal()
            try:
                # Pull recent rows then filter in Python so we match stored labels reliably.
                cap = 500
                rows = (
                    db.query(InitialSuggestionQuestion)
                    .filter(InitialSuggestionQuestion.agent_id == agent_id)
                    .order_by(InitialSuggestionQuestion.created_at.asc())
                    .limit(cap)
                    .all()
                )
                picked = []
                for r in rows:
                    if not self._row_matches_suggestion_datasource(
                        r.question_data_source, datasource_filter
                    ):
                        continue
                    picked.append(r)
                    if len(picked) >= max(limit, 0):
                        break
                questions_list: List[str] = []
                question_details_map: Dict[str, Any] = {}
                for r in picked:
                    questions_list.append(r.question)
                    question_details_map[r.question] = {
                        "id": r.id,
                        "sql": r.sql,
                        "api_endpoint": r.api_endpoint,
                        "api_params": r.api_params,
                        "api_response_validator": r.api_response_validator,
                        "question_data_source": r.question_data_source,
                        "question_type": r.question_type,
                    }
                return (questions_list, question_details_map)
            finally:
                db.close()
        except Exception as e:
            print(f"Error getting agent questions from DB: {str(e)}")
            return ([], {})

    def send_email(
        self,
        sender_email,
        sender_name,
        recipient_email,
        subject,
        body_text,
        body_html=None,
        attachment_path=None,
        attachment_paths=None,
        agent_conf=None,
    ):
        """
        Send email with optional HTML body and attachment(s).
        
        Args:
            sender_email (str): Sender's email address
            sender_name (str): Sender's name
            recipient_email (list): List of recipient email addresses
            subject (str): Email subject
            body_text (str): Plain text email body
            body_html (str, optional): HTML email body
            attachment_path (str, optional): Path to a single attachment file
            attachment_paths (list, optional): Paths to multiple attachment files
            agent_conf (dict): Agent configuration containing SMTP settings
        
        Returns:
            bool: True if email sent successfully, False otherwise
        """
        try:
            # Create message
            msg = MIMEMultipart('alternative')
            msg['From'] = f"{sender_name} <{sender_email}>"
            msg['To'] = ", ".join(recipient_email)
            msg['Subject'] = subject
            
            # Add text version
            part1 = MIMEText(body_text, 'plain')
            msg.attach(part1)
            
            # Add HTML version if provided
            if body_html:
                part2 = MIMEText(body_html, 'html')
                msg.attach(part2)
            
            paths_to_attach = []
            if attachment_paths:
                paths_to_attach.extend(
                    attachment_paths if isinstance(attachment_paths, list) else [attachment_paths]
                )
            elif attachment_path:
                paths_to_attach.append(attachment_path)

            for path in paths_to_attach:
                if not path or not os.path.isfile(path):
                    continue
                with open(path, "rb") as attachment:
                    part = MIMEBase('application', 'octet-stream')
                    part.set_payload(attachment.read())

                encoders.encode_base64(part)
                part.add_header(
                    'Content-Disposition',
                    f'attachment; filename= {os.path.basename(path)}',
                )
                msg.attach(part)
            
            # Create SMTP session
            server = smtplib.SMTP(agent_conf.get("smtp_server"), int(agent_conf.get("smtp_port")))
            server.starttls()  # Enable TLS encryption
            server.login(agent_conf.get("smtp_username"), agent_conf.get("smtp_password"))
            
            # Send email
            text = msg.as_string()
            server.sendmail(sender_email, recipient_email, text)
            server.quit()
            
            print(f"Email sent successfully to {recipient_email}")
            return True
            
        except Exception as e:
            print(f"Error sending email: {str(e)}")
            return False

    def process_api_user_query_multi_agent(self, user_query, session_id, agent_ids):
        """
        Process a user query for multiple agents (Non-SAP APIs).
        
        Args:
            user_query: User's natural language query
            session_id: Session identifier
            agent_ids: List of agent IDs
            
        Returns:
            Dict mapping agent_id to result dict
        """
        results = {}
        # Ensure agent_ids is a list
        if isinstance(agent_ids, str):
            agent_ids = [agent_ids]
            
        for agent_id in agent_ids:
            # Since process_api_user_query relies on the service instance's config (like API key or DB path),
            # we must use the specific service instance for each agent.
            
            # Local import to avoid circular dependency if any (though already inside class)
            from app.services.askme_service import get_askme_service 
            
            service = get_askme_service(agent_id=str(agent_id))
            if service:
                # We call the single-agent version
                # Note: process_api_user_query signature is (user_query, session_id, agent_id="")
                # We pass agent_id to it as well, though it primarily uses self config.
                res = service.process_api_user_query(user_query, session_id, agent_id=str(agent_id))
                results[agent_id] = res
            else:
                results[agent_id] = {"success": False, "message": f"Service not found for agent {agent_id}"}
                
        return results

    def process_sap_query_multi_agent(self, service_root, entity_set, user_query, response_schema, transformation_code, agent_ids, username=None, password=None, verify_ssl=True, paginate=True, api_query_params=None):
        """
        Process SAP OData query for multiple agents.
        
        Args:
            agent_ids: List of agent IDs
            [Other args same as process_sap_query]
        """
        results = {}
        if isinstance(agent_ids, str):
            agent_ids = [agent_ids]
            
        for agent_id in agent_ids:
            from app.services.askme_service import get_askme_service
            service = get_askme_service(agent_id=str(agent_id))
            if service:
                res = service.process_sap_query(
                    service_root, entity_set, user_query, response_schema, transformation_code,
                    username, password, verify_ssl, paginate, api_query_params
                )
                results[agent_id] = res
            else:
                results[agent_id] = {"success": False, "message": f"Service not found for agent {agent_id}"}
        return results


# --- Multi-agent intent routing, mixed execution, heterogeneous merge (no extra modules) ---

def normalize_request_datasource_override(raw: Optional[str]) -> Optional[str]:
    """Map client datasource string to DATABASE | API | FILE, or None if not specified."""
    if raw is None or (isinstance(raw, str) and not str(raw).strip()):
        return None
    u = str(raw).strip().upper()
    if u in ("DATABASE", "DB", "SQL"):
        return "DATABASE"
    if u == "API":
        return "API"
    if u in ("FILE", "DOCUMENT"):
        return "FILE"
    return None


def build_multi_agent_datasource_map(
    agent_ids: List[str],
    question: str,
    request_datasource: Optional[str],
) -> Tuple[Optional[Dict[str, Dict[str, Any]]], Optional[str]]:
    """
    Per-agent datasource selection: override > single configured > LLM intent.
    Returns (map, error_message). map[agent_id] = {type, confidence, reasoning}.
    """
    from app.services.agent_config import get_agent_by_id, ordered_available_datasources

    override = normalize_request_datasource_override(request_datasource)
    out: Dict[str, Dict[str, Any]] = {}
    llm_agent = None

    for aid in agent_ids:
        row = get_agent_by_id(aid)
        if not row:
            return None, f"Agent not found: {aid}"
        avail = ordered_available_datasources(row)
        if not avail:
            return None, f"No datasources configured for agent {aid}"

        if override is not None:
            if override not in avail:
                return None, (
                    f"Datasource override '{request_datasource}' is not enabled for agent {aid}. "
                    f"Allowed: {avail}"
                )
            out[aid] = {"type": override, "confidence": 1.0, "reasoning": "request_override"}
            continue

        if len(avail) == 1:
            out[aid] = {"type": avail[0], "confidence": 1.0, "reasoning": "single_configured_source"}
            continue

        if llm_agent is None:
            svc = get_askme_service(aid)
            if not svc:
                return None, f"No AskmeService for agent {aid}"
            llm_agent = svc._ensure_agent()
        pick = llm_agent.select_datasource_from_intent(question, avail)
        out[aid] = {
            "type": pick["datasource"],
            "confidence": float(pick.get("confidence") or 0.8),
            "reasoning": str(pick.get("reasoning") or ""),
        }
    return out, None


def classify_multi_agent_strategy(
    datasource_map: Dict[str, Dict[str, Any]], agent_ids: List[str]
) -> str:
    """SINGLE | DB_MULTI | API_MULTI | MIXED"""
    if len(agent_ids) <= 1:
        return "SINGLE"
    types = {datasource_map[a]["type"] for a in agent_ids if a in datasource_map}
    if len(types) == 1:
        t = next(iter(types))
        if t == "DATABASE":
            return "DB_MULTI"
        if t == "API":
            return "API_MULTI"
        return "MIXED"
    return "MIXED"


def _meta_exclude_cols(cols: List[str]) -> List[str]:
    return [c for c in cols if c not in ("_datasource_type", "_agents")]


def _schema_similarity(cols_a: List[str], cols_b: List[str]) -> float:
    a, b = set(_meta_exclude_cols(cols_a)), set(_meta_exclude_cols(cols_b))
    if not a and not b:
        return 1.0
    if not a or not b:
        return 0.0
    return len(a & b) / float(len(a | b))


def _pick_merge_key(common: List[str]) -> Optional[str]:
    if not common:
        return None
    cl = [str(c).lower() for c in common]
    for i, c in enumerate(common):
        if any(x in cl[i] for x in ("date", "time", "timestamp", "period")):
            return c
    for i, c in enumerate(common):
        if any(x in cl[i] for x in ("id", "code", "key")):
            return c
    return common[0]


def heterogeneous_merge_dataframes(
    db_df: Optional[pd.DataFrame],
    api_df: Optional[pd.DataFrame],
    max_rows: int = 5000,
) -> Tuple[pd.DataFrame, Dict[str, Any]]:
    """Merge DATABASE and/or API group DataFrames only. File/document results are never merged here."""

    def _ok(d: Optional[pd.DataFrame]) -> bool:
        return d is not None and not getattr(d, "empty", True)

    parts: List[Tuple[str, pd.DataFrame]] = []
    for label, d in (("DATABASE", db_df), ("API", api_df)):
        if _ok(d):
            parts.append((label, d.copy()))

    meta: Dict[str, Any] = {"strategies": []}
    if not parts:
        return pd.DataFrame(), {**meta, "merge_strategy": "empty"}
    if len(parts) == 1:
        mdf = parts[0][1].head(max_rows)
        return mdf, {**meta, "merge_strategy": "single"}

    def _merge_two(left: pd.DataFrame, right: pd.DataFrame) -> pd.DataFrame:
        c1 = _meta_exclude_cols(list(left.columns))
        c2 = _meta_exclude_cols(list(right.columns))
        common = [c for c in c1 if c in c2]
        key = _pick_merge_key(common)
        if key:
            try:
                meta["strategies"].append({"step": "merge_on_key", "key": key})
                return pd.merge(left, right, on=key, how="outer", suffixes=("", "_r"))
            except Exception as e:
                print(f"[heterogeneous_merge] key merge failed: {e}")
        sim = _schema_similarity(list(left.columns), list(right.columns))
        if sim > 0.7 and len(c1) == len(c2) and len(c1) > 0:
            meta["strategies"].append({"step": "concat_rows", "similarity": sim})
            try:
                return pd.concat([left, right], axis=0, ignore_index=True)
            except Exception:
                pass
        meta["strategies"].append({"step": "concat_columns"})
        try:
            return pd.concat(
                [left.reset_index(drop=True), right.reset_index(drop=True)], axis=1
            )
        except Exception as e:
            print(f"[heterogeneous_merge] concat_columns failed: {e}")
            return left

    merged = parts[0][1]
    for _, nxt in parts[1:]:
        merged = _merge_two(merged, nxt)

    truncated = len(merged) > max_rows
    if truncated:
        merged = merged.head(max_rows)
    meta["truncated"] = truncated
    meta["merge_strategy"] = "heterogeneous_chain"
    return merged, meta


def _convert_page_images_to_urls(page_images: Dict[str, Any]) -> Dict[str, Any]:
    """Convert page_images path fields to static API URLs (frontend @find contract)."""
    if not isinstance(page_images, dict):
        return page_images
    converted: Dict[str, Any] = {}
    for page_key, imgs in page_images.items():
        converted[page_key] = []
        for img in imgs or []:
            if isinstance(img, dict) and img.get("path"):
                c = dict(img)
                c["path"] = _mixed_file_image_path_to_url(img["path"])
                converted[page_key].append(c)
            else:
                converted[page_key].append(img)
    return converted


def _mixed_file_image_path_to_url(image_path: str) -> str:
    """Same contract as askme.convert_image_path_to_url (no route import)."""
    if not image_path:
        return ""
    s = str(image_path).strip().strip('"').strip("'")
    if s.startswith("http://") or s.startswith("https://"):
        return s
    try:
        from app.core.config import get_base_url

        base_url = get_base_url()
    except Exception:
        base_url = "http://localhost:8000"
    normalized = s.replace("\\", "/")
    return f"{base_url}/api/images/{normalized}"


def _decorate_file_handle_result_for_response(
    agent_id: str, user_question: str, query_str: str, doc_results: Dict[str, Any]
) -> Dict[str, Any]:
    """
    Build one agent's file payload aligned with /ask-file details (document_search / analysis / find / error).
    Image paths are converted to backend URLs like the ask-file route.
    """
    result_type = doc_results.get("type", "unknown")
    base: Dict[str, Any] = {
        "agent_id": agent_id,
        "datasource": "FILE",
        "question": user_question,
        "document_query": query_str,
        "sql": "N",
        "api_endpoint": "",
        "api_params": "",
        "api_response_validator": "",
        "document_result_type": result_type,
    }

    if result_type == "document_search":
        try:
            from app.services.document_service import enrich_document_search_file_paths_from_db

            enrich_document_search_file_paths_from_db(doc_results, agent_id=agent_id)
        except Exception:
            pass
        folder_groups = doc_results.get("folder_groups", {}) or {}
        for _folder_key, items in folder_groups.items():
            if not isinstance(items, list):
                continue
            for item in items:
                if not isinstance(item, dict) or "relevant_images" not in item:
                    continue
                converted = []
                for img in item.get("relevant_images") or []:
                    if isinstance(img, dict) and "path" in img:
                        c = dict(img)
                        c["path"] = _mixed_file_image_path_to_url(img["path"])
                        converted.append(c)
                    else:
                        converted.append(img)
                item["relevant_images"] = converted
        chunk_hits = doc_results.get("chunk_hits") or []
        for hit in chunk_hits:
            if isinstance(hit, dict) and hit.get("relevant_images"):
                hit["relevant_images"] = [
                    {**img, "path": _mixed_file_image_path_to_url(img["path"])}
                    if isinstance(img, dict) and img.get("path")
                    else img
                    for img in hit["relevant_images"]
                ]
        search_results = doc_results.get("results", []) or []
        for hit in search_results:
            if isinstance(hit, dict) and hit.get("relevant_images"):
                hit["relevant_images"] = [
                    {**img, "path": _mixed_file_image_path_to_url(img["path"])}
                    if isinstance(img, dict) and img.get("path")
                    else img
                    for img in hit["relevant_images"]
                ]
        base["document_search"] = {
            "query": doc_results.get("query", ""),
            "folder_detected": doc_results.get("folder_detected"),
            "chunk_hits": chunk_hits,
            "folder_groups": folder_groups,
            "results": search_results,
            "overall_summary": doc_results.get("overall_summary", ""),
            "matched_agent_ids": doc_results.get("matched_agent_ids", []),
        }
    elif result_type == "document_analysis":
        converted_images = []
        for img in doc_results.get("all_images") or []:
            if isinstance(img, dict):
                c = dict(img)
                if "path" in c:
                    c["path"] = _mixed_file_image_path_to_url(c["path"])
                converted_images.append(c)
            elif isinstance(img, str):
                converted_images.append(_mixed_file_image_path_to_url(img))
            else:
                converted_images.append(img)
        analysis_block: Dict[str, Any] = {
            "file_id": doc_results.get("file_id"),
            "filename": doc_results.get("filename", ""),
            "file_type": doc_results.get("file_type", ""),
            "folder_path": doc_results.get("folder_path", ""),
            "relative_path": doc_results.get("relative_path", ""),
            "folder_detected": doc_results.get("folder_detected"),
            "analysis": doc_results.get("analysis", ""),
            "all_images": converted_images,
            "file_path": doc_results.get("file_path", ""),
            "file_size": doc_results.get("file_size"),
            "upload_date": doc_results.get("upload_date"),
        }
        fid = doc_results.get("file_id")
        if fid and agent_id:
            from app.utils.document_preview import apply_preview_contract_to_mapping

            apply_preview_contract_to_mapping(
                analysis_block,
                file_id=str(fid),
                agent_id=str(agent_id),
                filename=str(analysis_block.get("filename") or ""),
                file_type=str(analysis_block.get("file_type") or ""),
                relative_url=True,
            )
        base["document_analysis"] = analysis_block
    elif result_type == "document_find":
        find_results = doc_results.get("results", []) or []
        for result in find_results:
            if not isinstance(result, dict):
                continue
            if result.get("page_images"):
                result["page_images"] = _convert_page_images_to_urls(result["page_images"])
            result.pop("relevant_images", None)
        base["document_find"] = {
            "query": doc_results.get("query", ""),
            "folder_detected": doc_results.get("folder_detected"),
            "results": find_results,
            "total_files_searched": doc_results.get("total_files_searched", 0),
            "message": doc_results.get("message"),
            "matched_agent_ids": doc_results.get("matched_agent_ids", []),
        }
    else:
        base["document_error"] = {
            "error": doc_results.get("error", doc_results.get("message", str(doc_results))),
            "folder_detected": doc_results.get("folder_detected"),
        }

    try:
        from app.utils.document_preview import collect_talkback_documents_from_file_results

        docs = collect_talkback_documents_from_file_results([base])
        if docs:
            base["documents"] = docs
    except Exception:
        pass

    return base


def run_file_agent_document_query_sync(
    agent_id: Union[str, List[str]], question: str
) -> Optional[Dict[str, Any]]:
    """
    Run DocumentService.handle_query for one agent (same prefixes as /ask-file).
    Returns a dict shaped like ask-file details for that agent — not a DataFrame (never merged with DB/API).
    """
    try:
        from app.services.document_service import (
            DocumentService,
            FILE_SYS_METADATA_DB,
            FILE_VECTORS_DB_DIR,
            FILE_UPLOAD_DISK_ROOT,
            FILE_IMAGE_DISK_ROOT,
            FILE_EXTRACT_DISK_ROOT,
        )
    except ImportError:
        print("[mixed] DocumentService not available")
        return None

    from app.services.agent_config import get_agent_config, agent_has_file_datasource

    # Normalize agent_id to a list for capability checking, but we need one primary agent for config
    agent_ids = [agent_id] if isinstance(agent_id, str) else agent_id
    if not agent_ids:
        return None
        
    primary_agent_id = agent_ids[0]
    
    # Check if at least one agent has file datasource (though usually the planner only picks ones that do)
    any_file_capable = any(agent_has_file_datasource(str(aid)) for aid in agent_ids)
    if not any_file_capable:
        return None
        
    cfg = get_agent_config(agent_id=primary_agent_id)
    if not cfg:
        return None
    api_key = cfg.get("api_key") or cfg.get("open_ai")
    if not api_key:
        return None

    doc = DocumentService(
        upload_folder=FILE_UPLOAD_DISK_ROOT,
        image_folder=FILE_IMAGE_DISK_ROOT,
        extracted_folder=FILE_EXTRACT_DISK_ROOT,
        db_file=FILE_SYS_METADATA_DB,
        chroma_path=FILE_VECTORS_DB_DIR,
        openai_api_key=str(api_key),
        claude_api_key=str(cfg.get("claude_api_key") or ""),
        agent_id=agent_id, # Can be string or list
    )

    query_str = str(question).strip() if question else ""
    ql = query_str.lower()
    if query_str and not any(
        ql.startswith(p) for p in ("@file", "@find", "@analyze", "@analyse")
    ):
        query_str = f"@file {query_str}"
    query_str = re.sub(
        r"@analyse\b", "@analyze", query_str, flags=re.IGNORECASE
    )

    doc_results = doc.handle_query(query_str)
    return _decorate_file_handle_result_for_response(
        primary_agent_id, question or "", query_str, doc_results
    )


def _sql_payload_is_reusable(sql_val: Any) -> bool:
    """
    True when sql_to_use should be executed as-is (matches multi-agent / saved-prompt
    semantics: N, NONE, [] are not reusable).
    """
    if sql_val is None:
        return False
    if isinstance(sql_val, str):
        s = sql_val.strip().upper()
        if not s or s in ("N", "NONE", "[]", "{}"):
            return False
        return True
    if isinstance(sql_val, (dict, list)):
        return bool(sql_val)
    return True


def _merge_info_from_sql_blob(sql_blob: Any) -> Optional[Dict[str, Any]]:
    """
    Normalize structured SQL payloads to run_sql_cached merge_info.

    Accepted shapes:
      - {"queries": {...}, "execution_order": [...]}
      - {"agent_results": [{"agent_id","sql",...}]}
      - [{"agent_id","sub_question","sql"}, ...]
      - JSON string of any of the above
    """
    parsed: Any = sql_blob
    if isinstance(parsed, str):
        s = parsed.strip()
        if not s.startswith(("{", "[")):
            return None
        try:
            parsed = json.loads(s)
        except Exception:
            return None

    if isinstance(parsed, dict):
        if "queries" in parsed and isinstance(parsed.get("queries"), dict):
            queries: Dict[str, str] = {}
            for k, v in (parsed.get("queries") or {}).items():
                key = str(k).strip()
                if not key:
                    continue
                sql_text = v if isinstance(v, str) else json.dumps(v, default=str)
                if not str(sql_text).strip():
                    continue
                queries[key] = str(sql_text)
            if not queries:
                return None
            order_raw = parsed.get("execution_order")
            if isinstance(order_raw, list):
                order = [str(x).strip() for x in order_raw if str(x).strip() in queries]
            else:
                order = list(queries.keys())
            if not order:
                order = list(queries.keys())
            return {
                "queries": queries,
                "execution_order": order,
                "merge_key": parsed.get("merge_key", "id"),
                "merge_strategy": parsed.get("merge_strategy", "independent"),
                "merge_steps": parsed.get("merge_steps") or [],
            }
        if "agent_results" in parsed:
            parsed = parsed.get("agent_results") or []
        else:
            return None

    if isinstance(parsed, list):
        queries: Dict[str, str] = {}
        order: List[str] = []
        for idx, item in enumerate(parsed):
            if not isinstance(item, dict):
                continue
            aid = str(
                item.get("agent_id")
                or item.get("id")
                or item.get("agent")
                or f"agent{idx + 1}"
            ).strip()
            if not aid:
                continue
            sql_part = item.get("sql")
            if sql_part is None:
                continue
            sql_text = sql_part if isinstance(sql_part, str) else json.dumps(sql_part, default=str)
            if not str(sql_text).strip():
                continue
            queries[aid] = str(sql_text)
            if aid not in order:
                order.append(aid)
        if not queries:
            return None
        return {
            "queries": queries,
            "execution_order": order or list(queries.keys()),
            "merge_key": "id",
            "merge_strategy": "independent",
            "merge_steps": [],
        }
    return None


def mixed_execute_db_group(
    db_agent_ids: List[str],
    question: str,
    sql_to_use: Any,
    *,
    force_sql_reuse_only: bool = False,
) -> Tuple[Optional[pd.DataFrame], Any, None]:
    """Returns (combined_df_with_meta, generated_sql, None)."""
    if not db_agent_ids:
        return None, None, None

    svc = get_askme_service(db_agent_ids if len(db_agent_ids) > 1 else db_agent_ids[0])
    if not svc:
        return None, None, None

    if force_sql_reuse_only and not _sql_payload_is_reusable(sql_to_use):
        print("[mixed] DB group: force_sql_reuse_only but sql_to_use is not reusable")
        return None, None, None

    generated_sql = sql_to_use
    if not _sql_payload_is_reusable(generated_sql):
        generated_sql = svc.generate_sql_cached(question=question)
    if not generated_sql:
        return None, None, None

    parsed = _merge_info_from_sql_blob(generated_sql)
    is_multi_agent_json = bool(parsed and parsed.get("queries"))
    sql_for_execution = json.dumps(parsed) if is_multi_agent_json else generated_sql

    # Skip execution entirely when multi-agent SQL generation produced no queries
    # (e.g. all DB agents failed SQL generation due to DB timeout / connection error).
    # Attempting run_sql_cached with an empty-queries JSON wastes another TCP round-trip.
    if is_multi_agent_json and parsed and not parsed.get("queries"):
        print(f"[mixed] DB group: skipping run_sql_cached — no SQL queries were generated (empty queries dict)")
        return None, generated_sql, None

    merge_info_arg = parsed if is_multi_agent_json else None
    if is_multi_agent_json and parsed and len(db_agent_ids) == 1:
        # Single DB agent inside mixed flow: execute plain SQL for that agent.
        # Passing JSON/merge_info to a single-agent AskmeService can fall through
        # to raw SQL execution and trigger ORA-00900.
        qmap = parsed.get("queries") or {}
        order = parsed.get("execution_order") or list(qmap.keys())
        picked_key = str(order[0]) if order else (str(next(iter(qmap.keys()))) if qmap else "")
        picked_sql = qmap.get(picked_key) if picked_key else None
        if picked_sql is None and qmap:
            picked_sql = next(iter(qmap.values()))
        if picked_sql is not None and str(picked_sql).strip():
            sql_for_execution = str(picked_sql)
            merge_info_arg = None
    try:
        df = svc.run_sql_cached(
            sql_for_execution, merge_info=merge_info_arg, question=question
        )
    except Exception as e:
        print(f"[mixed] DB group run_sql_cached failed: {e}")
        return None, generated_sql, None

    agents_suffix = ",".join(db_agent_ids)
    if isinstance(df, dict):
        pieces = []
        for _aid, sub in df.items():
            if sub is None or getattr(sub, "empty", True):
                continue
            x = sub.copy()
            x["_datasource_type"] = "DATABASE"
            x["_agents"] = agents_suffix
            pieces.append(x)
        if not pieces:
            return None, generated_sql, None
        db_df = pd.concat(pieces, axis=0, ignore_index=True)
        return db_df, generated_sql, None

    if df is None or getattr(df, "empty", True):
        return None, generated_sql, None
    db_df = df.copy()
    db_df["_datasource_type"] = "DATABASE"
    db_df["_agents"] = agents_suffix
    return db_df, generated_sql, None


def _parse_saved_api_list_field(field_val: Any) -> List[Any]:
    """Same shape as parse_list_field in multi-agent-ask (saved API columns)."""
    if not field_val or field_val == "N":
        return []
    if isinstance(field_val, list):
        return field_val
    if isinstance(field_val, str):
        s = field_val.strip()
        if s.startswith("["):
            try:
                return json.loads(s)
            except Exception:
                pass
        return [field_val]
    return [field_val]


def _normalize_saved_api_params(params: Any) -> Dict[str, Any]:
    if isinstance(params, dict):
        return params
    if isinstance(params, str) and params.strip():
        try:
            if params.strip().startswith("{"):
                return ast.literal_eval(params)
            if params.strip().startswith("["):
                temp_list = json.loads(params)
                return temp_list[0] if temp_list else {}
            return json.loads(params)
        except Exception:
            try:
                return json.loads(params)
            except Exception:
                pass
    return {}


def mixed_execute_api_group(
    api_agent_ids: List[str],
    question: str,
    *,
    api_endpoint_to_use: Any = None,
    api_params_to_use: Any = None,
    api_response_validator_to_use: Any = None,
    force_api_replay_only: bool = False,
) -> Tuple[Optional[pd.DataFrame], Dict[str, Any]]:
    if not api_agent_ids:
        return None, {}

    endpoints = _parse_saved_api_list_field(api_endpoint_to_use)
    params_list = _parse_saved_api_list_field(api_params_to_use)
    validators_list = _parse_saved_api_list_field(api_response_validator_to_use)

    use_replay = bool(
        endpoints
        and params_list
        and api_endpoint_to_use not in (None, "", "N", "[]")
        and api_params_to_use not in (None, "", "N", "[]")
    )

    if force_api_replay_only and not use_replay:
        print(
            "[mixed] API group: force_api_replay_only but api_endpoint/api_params "
            "are missing or not parseable — skipping LLM API selection"
        )
        return None, {}

    svc0 = get_askme_service(api_agent_ids[0])
    if not svc0:
        return None, {}

    merge_result: Dict[str, Any] = {}
    if use_replay:
        api_chat = svc0._ensure_api_chat_service()
        dataframes: List[pd.DataFrame] = []
        api_info_list: List[Dict[str, Any]] = []
        result_ep: List[str] = []
        result_pr: List[str] = []
        result_val: List[str] = []

        for i, endpoint in enumerate(endpoints):
            if not endpoint:
                continue
            params_raw = params_list[i] if i < len(params_list) else {}
            validator = validators_list[i] if i < len(validators_list) else ""
            if not validator:
                validator = ""
            aid = str(
                api_agent_ids[i] if i < len(api_agent_ids) else api_agent_ids[0]
            )

            svc = get_askme_service(aid)
            if not svc:
                continue
            apis_list = svc.get_all_apis(aid) or []
            api_info = next(
                (a for a in apis_list if a.get("endpoint") == endpoint), {}
            )
            if not api_info:
                print(
                    f"[mixed] replay: endpoint not registered for agent {aid}: {endpoint}"
                )
                continue

            params = _normalize_saved_api_params(params_raw)

            is_sap = api_info.get("extra1", "N")
            sub_df = None
            if is_sap == "Y":
                sap_result = svc.process_sap_query(
                    service_root=endpoint,
                    entity_set=api_info.get("base_params"),
                    user_query=question,
                    response_schema=api_info.get("response_schema"),
                    transformation_code=validator,
                    username="S4h_aa",
                    password="Welcome1",
                    verify_ssl=False,
                    api_query_params=params,
                )
                sub_df = sap_result.get("df")
            else:
                _, _, sub_df, _ = svc.call_api_and_transform(
                    params,
                    endpoint,
                    api_info.get("response_schema", ""),
                    validator,
                    question,
                )

            row_detail = {
                "name": api_info.get("name", ""),
                "endpoint": endpoint,
                "params": str(params) if params else "",
                "validator": validator,
                "type": "SAP" if is_sap == "Y" else "REST",
                "agent_id": aid,
            }
            result_ep.append(endpoint)
            result_pr.append(row_detail["params"])
            result_val.append(validator)

            if isinstance(sub_df, dict):
                for key, sub in sub_df.items():
                    if sub is not None and isinstance(sub, pd.DataFrame) and not sub.empty:
                        dataframes.append(sub)
                        sub_d = row_detail.copy()
                        sub_d["name"] = f"{row_detail['name']} - {key}"
                        api_info_list.append(sub_d)
            elif sub_df is not None and hasattr(sub_df, "empty") and not sub_df.empty:
                dataframes.append(sub_df)
                api_info_list.append(row_detail)
            else:
                api_info_list.append(
                    {**row_detail, "status": "No Data/Error"}
                )

        if dataframes:
            if len(dataframes) == 1:
                merged_inner = dataframes[0]
            else:
                try:
                    merged_list = api_chat._merge_dataframes(dataframes, question)
                    merged_inner = (
                        merged_list[0]
                        if len(merged_list) == 1
                        else merged_list
                    )
                except Exception as e:
                    print(f"[mixed] replay merge failed: {e}")
                    merged_inner = dataframes

            merge_result = {
                "success": True,
                "message": "Replayed saved API parameters",
                "df": merged_inner,
                "api_details": api_info_list,
                "api_endpoints": result_ep,
                "api_params": result_pr,
                "api_validators": result_val,
            }
        else:
            merge_result = {
                "success": False,
                "message": "No data from replayed API calls",
                "df": None,
                "api_details": api_info_list,
                "api_endpoints": result_ep,
                "api_params": result_pr,
                "api_validators": result_val,
            }

        if use_replay and not merge_result.get("success"):
            if force_api_replay_only:
                print(
                    "[mixed] saved API replay produced no data; "
                    "force_api_replay_only — not using LLM API selection"
                )
                return None, merge_result or {}
            print(
                "[mixed] saved API replay produced no data; falling back to API selection"
            )
            merge_result = {}

    if not merge_result or not merge_result.get("success"):
        all_apis: List[Dict[str, Any]] = []
        for aid in api_agent_ids:
            s = get_askme_service(aid)
            if not s:
                continue
            try:
                all_apis.extend(s.get_all_apis(aid) or [])
            except Exception as e:
                print(f"[mixed] get_all_apis {aid}: {e}")

        if not all_apis:
            return None, merge_result or {}

        if force_api_replay_only:
            print(
                "[mixed] force_api_replay_only — skipping LLM API selection "
                "(replay path did not succeed)"
            )
            return None, merge_result or {}

        api_chat = svc0._ensure_api_chat_service()
        selected = api_chat._select_api_from_candidates(question, all_apis)
        if not selected:
            return None, merge_result or {}

        merge_result = api_chat.call_multiple_apis_and_merge(
            selected_apis=selected, user_query=question
        )
        if not merge_result or not merge_result.get("success"):
            return None, merge_result or {}

    df = merge_result.get("df")
    agents_suffix = ",".join(api_agent_ids)

    if isinstance(df, list) and df and isinstance(df[0], pd.DataFrame):
        pieces = []
        for sub in df:
            if sub is None or getattr(sub, "empty", True):
                continue
            x = sub.copy()
            x["_datasource_type"] = "API"
            x["_agents"] = agents_suffix
            pieces.append(x)
        if not pieces:
            return None, merge_result
        api_df = pd.concat(pieces, axis=0, ignore_index=True)
        return api_df, merge_result

    if df is not None and hasattr(df, "copy"):
        api_df = df.copy()
        api_df["_datasource_type"] = "API"
        api_df["_agents"] = agents_suffix
        return api_df, merge_result

    return None, merge_result


def execute_mixed_multi_agent_sync(
    agent_ids: List[str],
    datasource_map: Dict[str, Dict[str, Any]],
    question: str,
    sql_to_use: Any,
    max_rows: int = 5000,
    *,
    api_endpoint_to_use: Any = None,
    api_params_to_use: Any = None,
    api_response_validator_to_use: Any = None,
    mixed_reuse_both_payload: bool = False,
) -> Dict[str, Any]:
    """
    File-first: run FILE agents (document_search / document_analysis / document_find — not DataFrames).
    Then run DB and API groups in parallel; merge only DB + API with heterogeneous_merge_dataframes.

    When ``mixed_reuse_both_payload`` is True (request supplied both meaningful SQL and API fields),
    SQL is not regenerated and API calls do not fall back to LLM selection.
    """
    from concurrent.futures import ThreadPoolExecutor

    file_agents = [a for a in agent_ids if datasource_map.get(a, {}).get("type") == "FILE"]
    db_agents = [a for a in agent_ids if datasource_map.get(a, {}).get("type") == "DATABASE"]
    api_agents = [a for a in agent_ids if datasource_map.get(a, {}).get("type") == "API"]

    file_results: List[Dict[str, Any]] = []
    if file_agents:
        for aid in file_agents:
            try:
                fd = run_file_agent_document_query_sync(aid, question)
                if fd is not None:
                    file_results.append(fd)
            except Exception as e:
                print(f"[mixed] file agent {aid} failed: {e}")
                file_results.append(
                    {
                        "agent_id": aid,
                        "datasource": "FILE",
                        "document_result_type": "error",
                        "document_error": {"error": str(e)},
                        "question": question,
                        "document_query": "",
                        "sql": "N",
                        "api_endpoint": "",
                        "api_params": "",
                        "api_response_validator": "",
                    }
                )

    db_df = None
    gen_sql = None
    api_df = None
    api_meta: Dict[str, Any] = {}

    def _run_db():
        return mixed_execute_db_group(
            db_agents,
            question,
            sql_to_use,
            force_sql_reuse_only=mixed_reuse_both_payload,
        )

    def _run_api():
        return mixed_execute_api_group(
            api_agents,
            question,
            api_endpoint_to_use=api_endpoint_to_use,
            api_params_to_use=api_params_to_use,
            api_response_validator_to_use=api_response_validator_to_use,
            force_api_replay_only=mixed_reuse_both_payload,
        )

    with ThreadPoolExecutor(max_workers=2) as ex:
        f_db = ex.submit(_run_db) if db_agents else None
        f_api = ex.submit(_run_api) if api_agents else None
        if f_db is not None:
            try:
                db_df, gen_sql, _ = f_db.result()
            except Exception as e:
                print(f"[mixed] DB group failed: {e}")
        if f_api is not None:
            try:
                api_df, api_meta = f_api.result()
            except Exception as e:
                print(f"[mixed] API group failed: {e}")

    merged, merge_meta = heterogeneous_merge_dataframes(
        db_df, api_df, max_rows=max_rows
    )

    has_tabular = merged is not None and not getattr(merged, "empty", True)
    _ok_file_types = ("document_search", "document_analysis", "document_find")
    has_file_ok = any(
        fr.get("document_result_type") in _ok_file_types for fr in file_results
    )
    success = has_tabular or has_file_ok

    return {
        "success": success,
        "df": merged,
        "file_results": file_results,
        "generated_sql": gen_sql,
        "datasource_map": datasource_map,
        "merge_meta": merge_meta,
        "api_details": api_meta,
        "has_tabular": has_tabular,
    }


# Per-agent service cache
_askme_services: dict = {}

_SELF_DB_AGENT_ID = "0690"


def _sqlite_path_from_settings() -> str:
    url = (getattr(settings, "DATABASE_URL", "") or "").strip()
    if "sqlite" not in url.lower():
        return ""
    path = url.replace("sqlite:///", "").lstrip("./")
    return os.path.abspath(path)


def _apply_self_db_overrides(agent_id: Union[str, int], cfg: Optional[Dict[str, Any]]) -> Optional[Dict[str, Any]]:
    """
    Hardcode agent 0690 to run against AskMe360's own SQLite DB.
    Keep other details (training data, model, etc.) from the agent config row.
    """
    if cfg is None:
        return None
    if str(agent_id).strip() != _SELF_DB_AGENT_ID:
        return cfg
    db_path = _sqlite_path_from_settings()
    if not db_path:
        print(
            f"Warning: {_SELF_DB_AGENT_ID} expects sqlite DATABASE_URL, "
            f"got {getattr(settings, 'DATABASE_URL', '')!r}"
        )
        return cfg
    out = dict(cfg)
    out["db_type"] = "sqlite"
    out["host"] = ""
    out["port"] = ""
    out["db_user"] = ""
    out["user"] = ""
    out["db_password"] = ""
    out["password"] = ""
    out["dbname"] = db_path
    out["db_name"] = db_path
    out["connection_string"] = db_path
    out["dsn"] = db_path
    return out


def get_askme_service(agent_id: Union[str, List[str], None] = None, config: Optional[dict] = None):
    """Get or create AskmeService instance for single or multiple agents"""
    global _askme_services

    if config:
        return AskmeService(
            model=config.get("model", "gpt-4.1"),
            chroma_path=config.get("chroma_path", config.get("path", "./data/askme")),
            traning_data=config.get("traning_data", ""),
            lib_path=config.get("lib_path", ""),
            open_ai=config.get("open_ai", config.get("api_key", "")),
            claude_api_key=config.get("claude_api_key", ""),
            connection_string=config.get("connection_string", config.get("dsn", "")),
            db_user=config.get("db_user", config.get("user", "")),
            db_password=config.get("db_password", config.get("password", "")),
            web_key=config.get("web_key", ""),
            db_type=config.get("db_type", ""),
            host=config.get("host", ""),
            port=config.get("port", ""),
            dbname=config.get("dbname", ""),
            max_row_count=config.get("max_row_count", "1000")

        )

    if agent_id is None:
        if "_default" not in _askme_services:
            _askme_services["_default"] = AskmeService()
        return _askme_services["_default"]

    agent_ids_list = agent_id if isinstance(agent_id, list) else [agent_id]
    if len(agent_ids_list) == 0:
        if "_default" not in _askme_services:
            _askme_services["_default"] = AskmeService()
        return _askme_services["_default"]

    if len(agent_ids_list) > 1:
        cache_key = "_".join(sorted(str(a) for a in agent_ids_list))
        if cache_key not in _askme_services:
            from app.services.agent_config import get_agent_config
            first_config = get_agent_config(agent_id=agent_ids_list[0])
            first_config = _apply_self_db_overrides(agent_ids_list[0], first_config)
            if not first_config:
                print(f"Warning: No config found for agent_id: {agent_ids_list[0]}")
                return None
            _askme_services[cache_key] = AskmeService(
                model=first_config.get("model", "gpt-4.1"),
                chroma_path=first_config.get("chroma_path", first_config.get("path", "./data/askme")),
                traning_data=first_config.get("traning_data", ""),
                lib_path=first_config.get("lib_path", ""),
                open_ai=first_config.get("api_key", ""),
                claude_api_key=first_config.get("claude_api_key", ""),
                connection_string=first_config.get("connection_string", first_config.get("dsn", "")),
                db_user=first_config.get("db_user", first_config.get("user", "")),
                db_password=first_config.get("db_password", first_config.get("password", "")),
                web_key=first_config.get("web_key", ""),
                db_type=first_config.get("db_type", ""),
                host=first_config.get("host", ""),
                port=first_config.get("port", ""),
                dbname=first_config.get("dbname", ""),
                agent_ids=agent_ids_list,
            )
        return _askme_services[cache_key]

    single_id = agent_ids_list[0]
    if single_id not in _askme_services:
        from app.services.agent_config import get_agent_config
        agent_config = get_agent_config(agent_id=single_id)
        agent_config = _apply_self_db_overrides(single_id, agent_config)
        if not agent_config:
            print(f"Warning: No config found for agent_id: {single_id}")
            return None
        _askme_services[single_id] = AskmeService(
            model=agent_config.get("model", "gpt-4.1"),
            chroma_path=agent_config.get("chroma_path", agent_config.get("path", "./data/askme")),
            traning_data=agent_config.get("traning_data", ""),
            lib_path=agent_config.get("lib_path", ""),
            open_ai=agent_config.get("api_key", ""),
            claude_api_key=agent_config.get("claude_api_key", ""),
            connection_string=agent_config.get("connection_string", agent_config.get("dsn", "")),
            db_user=agent_config.get("db_user", agent_config.get("user", "")),
            db_password=agent_config.get("db_password", agent_config.get("password", "")),
            web_key=agent_config.get("web_key", ""),
            db_type=agent_config.get("db_type", ""),
            host=agent_config.get("host", ""),
            port=agent_config.get("port", ""),
            dbname=agent_config.get("dbname", ""),
            max_row_count=agent_config.get("max_row_count", "1000"),
        )
    return _askme_services.get(single_id)


def clear_askme_service_cache(agent_id: str = None):
    """Clear cached AskmeService instances
    
    Args:
        agent_id: Specific agent to clear, or None to clear all
    """
    global _askme_services
    if agent_id:
        _askme_services.pop(agent_id, None)
    else:
        _askme_services.clear()
