# ==== CREATED AND OWNED BY askme360 (HIPL) ====
import json
import os
import re
import sqlite3
import traceback
from html import escape
from abc import ABC, abstractmethod
from typing import Any, Dict, List, Optional, Tuple, Union
from urllib.parse import urlparse
from tabulate import tabulate
import pandas as pd
import plotly
import plotly.express as px
import plotly.graph_objects as go
import requests
import sqlparse
import datetime

# Optional: Scikit-learn and Statsmodels for forecasting (may be absent in env / type checker)
try:
    from sklearn.linear_model import LinearRegression  # type: ignore[import-untyped]
    SKLEARN_AVAILABLE = True
except ImportError:
    SKLEARN_AVAILABLE = False
    LinearRegression = None

try:
    import statsmodels.api as sm  # type: ignore[import-untyped]
    STATSMODELS_AVAILABLE = True
except ImportError:
    STATSMODELS_AVAILABLE = False
    sm = None

# Import node-based SQL generation components
from askme_open_agent.askme_base.sql_nodes import (
    SQLGenerationPipeline,
    SQLGenerationState,
    NodeStatus,
)


# Import exceptions, types, and utils from updated module
try:
    from askme_open_agent.exceptions import (
        DependencyError,
        ImproperlyConfigured,
        QueryTimeoutError,
        ValidationError,
    )
    from askme_open_agent.types import TrainingPlan, TrainingPlanItem
    from askme_open_agent.utils import validate_config_path
except ImportError:
    # Fallback: define minimal exceptions if imports fail
    class DependencyError(Exception):
        """Raise for missing dependencies."""
        pass
    
    class ImproperlyConfigured(Exception):
        """Raise for incorrect configuration."""
        pass
    
    class ValidationError(Exception):
        """Raise for validations."""
        pass

    class QueryTimeoutError(Exception):
        """Raised when a DB statement exceeds TIMEOUT_DB."""
        pass

    # Minimal type stubs
    class TrainingPlanItem:
        ITEM_TYPE_SQL = "sql"
        ITEM_TYPE_DDL = "ddl"
        ITEM_TYPE_IS = "is"
        def __init__(self, item_type, item_group, item_name, item_value):
            self.item_type = item_type
            self.item_group = item_group
            self.item_name = item_name
            self.item_value = item_value
    
    class TrainingPlan:
        def __init__(self, plan):
            self._plan = plan
    
    def validate_config_path(path):
            import os
            if not os.path.exists(path):
                raise ImproperlyConfigured(f'No such configuration file: {path}')
            if not os.path.isfile(path):
                raise ImproperlyConfigured(f'Config should be a file: {path}')
            if not os.access(path, os.R_OK):
                raise ImproperlyConfigured(f'Cannot read the config file. Please grant read privileges: {path}')


# --- Plain-language SQL summaries (details.sql_query_summaries; used via AskmeService) ---


def sql_plain_summary_enabled() -> bool:
    try:
        from app.core.config import get_initial_config, settings

        if not getattr(settings, "ENABLE_SQL_PLAIN_SUMMARY", True):
            return False
        cfg = get_initial_config()
        raw = (cfg.get("ENABLE_SQL_PLAIN_SUMMARY") or "").strip().lower()
        if raw in ("false", "0", "no", "n", "off"):
            return False
        return True
    except Exception:
        return True


def _sql_summary_parse_json_array_from_llm_text(text: str) -> Optional[List[Any]]:
    if text is None:
        return None
    t = str(text).strip()
    if not t:
        return None
    unfenced = re.sub(r"^\s*```(?:json)?\s*", "", t, flags=re.IGNORECASE | re.DOTALL)
    unfenced = re.sub(r"\s*```\s*$", "", unfenced, flags=re.DOTALL).strip()
    for body in (t, unfenced):
        if not body:
            continue
        try:
            data = json.loads(body)
            if isinstance(data, list):
                return data
        except json.JSONDecodeError:
            pass
        start = body.find("[")
        if start < 0:
            continue
        depth = 0
        for i in range(start, len(body)):
            ch = body[i]
            if ch == "[":
                depth += 1
            elif ch == "]":
                depth -= 1
                if depth == 0:
                    try:
                        data = json.loads(body[start : i + 1])
                        if isinstance(data, list):
                            return data
                    except json.JSONDecodeError:
                        pass
                    break
    return None


def _sql_summary_coerce_details_sql_to_str(details_sql: Any) -> str:
    if details_sql is None:
        return ""
    if isinstance(details_sql, (dict, list)):
        try:
            return json.dumps(details_sql, default=str)
        except Exception:
            return str(details_sql)
    if isinstance(details_sql, bytes):
        return details_sql.decode("utf-8", errors="replace")
    return str(details_sql)


def _sql_summary_inline_markdown_to_html(text: str) -> str:
    if not text:
        return ""
    out: List[str] = []
    pos = 0
    for m in re.finditer(r"\*\*(.+?)\*\*", text):
        out.append(escape(text[pos : m.start()]))
        out.append(f"<strong>{escape(m.group(1))}</strong>")
        pos = m.end()
    out.append(escape(text[pos:]))
    return "".join(out)


def _sql_summary_markdown_table_to_html(lines: List[str]) -> str:
    rows: List[List[str]] = []
    for line in lines:
        row = line.strip()
        if not (row.startswith("|") and row.endswith("|")):
            continue
        cells = [c.strip() for c in row[1:-1].split("|")]
        if not cells:
            continue
        if all(re.fullmatch(r"[-:\s]+", c or "") for c in cells):
            continue
        rows.append(cells)
    if not rows:
        return ""
    header = rows[0]
    body = rows[1:]
    cols = len(header)
    head_html = "".join(
        (
            '<th style="text-align:left; padding:8px; border-bottom:1px solid #9ca3af;">'
            f"{_sql_summary_inline_markdown_to_html(c)}</th>"
        )
        for c in header
    )
    body_parts: List[str] = []
    for row in body:
        padded = row + ([""] * max(0, cols - len(row)))
        padded = padded[:cols]
        td = "".join(
            (
                '<td style="padding:8px; border-bottom:1px solid #e5e7eb; vertical-align:top;">'
                f"{_sql_summary_inline_markdown_to_html(c)}</td>"
            )
            for c in padded
        )
        body_parts.append(f"<tr>{td}</tr>")
    return (
        '<table style="width:100%; border-collapse:collapse; margin:8px 0 16px 0;">'
        f"<thead><tr>{head_html}</tr></thead>"
        f"<tbody>{''.join(body_parts)}</tbody>"
        "</table>"
    )


def _sql_summary_markdown_to_html(markdown_text: str) -> str:
    text = (markdown_text or "").strip()
    if not text:
        return ""
    lines = text.splitlines()
    out: List[str] = []
    i = 0
    while i < len(lines):
        raw = lines[i].rstrip()
        s = raw.strip()
        if not s:
            i += 1
            continue
        if s == "---":
            out.append('<hr style="border:none; border-top:1px solid #d1d5db; margin:16px 0;" />')
            i += 1
            continue
        if s.startswith("|") and s.endswith("|"):
            table_lines: List[str] = []
            while i < len(lines):
                ss = lines[i].strip()
                if ss.startswith("|") and ss.endswith("|"):
                    table_lines.append(lines[i])
                    i += 1
                    continue
                break
            out.append(_sql_summary_markdown_table_to_html(table_lines))
            continue
        m = re.match(r"^\*\*(.+?)\*\*$", s)
        if m:
            out.append(f"<h3>{_sql_summary_inline_markdown_to_html(m.group(1))}</h3>")
            i += 1
            continue
        if re.match(r"^Step\s+\d+\s*[–-]", s):
            out.append(f"<h4>{_sql_summary_inline_markdown_to_html(s)}</h4>")
            i += 1
            continue
        if s in (
            "What This Query Does (Plain English)",
            "The Final Output",
            "Quick Summary Table",
        ):
            out.append(f"<h3>{_sql_summary_inline_markdown_to_html(s)}</h3>")
            i += 1
            continue
        if s.startswith("- ") or s.startswith("* "):
            lis: List[str] = []
            while i < len(lines):
                item = lines[i].strip()
                if item.startswith("- ") or item.startswith("* "):
                    lis.append(
                        f"<li>{_sql_summary_inline_markdown_to_html(item[2:].strip())}</li>"
                    )
                    i += 1
                    continue
                break
            out.append(f"<ul>{''.join(lis)}</ul>")
            continue
        out.append(f"<p>{_sql_summary_inline_markdown_to_html(s)}</p>")
        i += 1
    body = "".join(out)
    return (
        '<div class="sql-summary-html" style="line-height:1.6; color:#111827;">'
        f"{body}"
        "</div>"
    )


def database_sql_items_from_details_sql(details_sql: Any) -> List[Dict[str, str]]:
    """
    Extract DATABASE legs from details['sql'] envelope JSON as
    [{agent_id, sub_question, sql_text}, ...].
    """
    raw = _sql_summary_coerce_details_sql_to_str(details_sql).strip()
    if not raw:
        return []
    try:
        d = json.loads(raw)
    except Exception:
        return []
    if not isinstance(d, dict):
        return []
    traces = d.get("sql_traces") or []
    if not isinstance(traces, list):
        return []

    db_agents: Optional[set] = None
    if bool(d.get("mixed_datasource_execution")):
        db_agents = set()
        for st in d.get("sub_tasks") or []:
            if not isinstance(st, dict):
                continue
            if str(st.get("datasource") or "DATABASE").upper() == "DATABASE":
                aid = str(st.get("agent_id") or "").strip()
                if aid:
                    db_agents.add(aid)

    def _rows_filtered(agent_set: Optional[set]) -> List[Dict[str, str]]:
        rows: List[Dict[str, str]] = []
        for t in traces:
            if not isinstance(t, dict):
                continue
            aid = str(t.get("agent_id") or "").strip()
            if agent_set is not None and agent_set and aid not in agent_set:
                continue
            sql = t.get("sql")
            if sql is None:
                continue
            st = sql if isinstance(sql, str) else json.dumps(sql, default=str)
            if not str(st).strip():
                continue
            rows.append(
                {
                    "agent_id": aid,
                    "sub_question": str(t.get("sub_question") or "").strip(),
                    "sql_text": str(st),
                }
            )
        return rows

    out = _rows_filtered(db_agents)
    if not out and db_agents and traces:
        out = _rows_filtered(None)
    return out


class AskmeBase(ABC):
    def __init__(self, config=None):
        if config is None:
            config = {}
        self.config = config
        self.run_sql_is_set = False
        self.static_documentation = ""
        self.dialect = self.config.get("dialect", "SQL")
        self.language = self.config.get("language", None)
        self.max_tokens = self.config.get("max_tokens", 14000)
        self.max_rows = self.config.get("max_rows", "1000")
        self.back_year = self.config.get("back_year", "2024")
        self.conn = None

    def _get_sql_statement_timeout_ms(self) -> int:
        """
        Maximum time for LLM-generated SQL to run (milliseconds for ``oracledb`` ``call_timeout``).

        ``app.core.config.settings.TIMEOUT`` is always **seconds** (see ``api_request_timeout_seconds``
        and ``initial_config.txt`` normalization). Older code wrongly treated it as ms, which made
        ``TIMEOUT=300`` become a **300 ms** Oracle call timeout (ORA-03156 / DPY-4024).
        """
        try:
            from app.core.config import db_statement_timeout_ms, settings

            sec = int(getattr(settings, "TIMEOUT", 300) or 300)
        except Exception:
            sec = 300
        sec = max(1, sec)
        return sec * 1000

    def log(self, message: str, title: str = "Info"):
        print(f"{title}: {message}")

    def _response_language(self) -> str:
        if self.language is None:
            return ""

        return f"Respond in the {self.language} language."

    def generate_sql(self, question: str, allow_llm_to_see_data=False, skip_node_pipeline=False, **kwargs) -> str:
        """
        Example:
        ```python
        self.generate_sql("What are the top 10 customers by sales?")
        ```

        Uses the LLM to generate a SQL query that answers a question. It runs the following methods:

        - [`get_similar_question_sql`][askme_openai_base.AskmeBase.get_similar_question_sql]

        - [`get_related_ddl`][askme_openai_base.AskmeBase.get_related_ddl]

        - [`get_related_documentation`][askme_openai_base.AskmeBase.get_related_documentation]

        - [`get_sql_prompt`][askme_openai_base.AskmeBase.get_sql_prompt]

        - [`submit_prompt`][askme_openai_base.AskmeBase.submit_prompt]
        
        If config['use_node_pipeline'] is True, this method will use the node-based pipeline
        (generate_sql_with_nodes) instead of the traditional approach.


        Args:
            question (str): The question to generate a SQL query for.
            allow_llm_to_see_data (bool): Whether to allow the LLM to see the data (for the purposes of introspecting the data to generate the final SQL).
            skip_node_pipeline (bool): Whether to skip the node pipeline and use the traditional approach.

        Returns:
            str: The SQL query that answers the question.
        """

        # Check if node-based pipeline should be used
        if not skip_node_pipeline:
            max_retries = self.config.get("node_pipeline_max_retries", 1)
            return self.generate_sql_with_nodes(question=question, max_retries=max_retries, **kwargs)
        

        if self.config is not None:
            initial_prompt = self.config.get("initial_prompt", None)
        else:
            initial_prompt = None
        question_sql_list = self.get_similar_question_sql(question, **kwargs)
        ddl_list = self.get_related_ddl(question, **kwargs)
        doc_list = self.get_related_documentation(question, **kwargs)
        prompt = self.get_sql_prompt(
            initial_prompt=initial_prompt,
            question=question,
            question_sql_list=question_sql_list,
            ddl_list=ddl_list,
            doc_list=doc_list,
            **kwargs,
        )
        self.log(title="SQL Prompt", message=prompt)
        llm_response = self.submit_prompt(prompt, **kwargs)
        self.log(title="LLM Response", message=llm_response)

        if 'intermediate_sql' in llm_response:
            if not allow_llm_to_see_data:
                return "The LLM is not allowed to see the data in your database. Your question requires database introspection to generate the necessary SQL. Please set allow_llm_to_see_data=True to enable this."

            if allow_llm_to_see_data:
                intermediate_sql = self.extract_sql(llm_response)

                try:
                    self.log(title="Running Intermediate SQL", message=intermediate_sql)
                    df = self.run_sql(intermediate_sql)

                    prompt = self.get_sql_prompt(
                        initial_prompt=initial_prompt,
                        question=question,
                        question_sql_list=question_sql_list,
                        ddl_list=ddl_list,
                        doc_list=doc_list+[f"The following is a pandas DataFrame with the results of the intermediate SQL query {intermediate_sql}: \n" + df.to_markdown()],
                        **kwargs,
                    )
                    self.log(title="Final SQL Prompt", message=prompt)
                    llm_response = self.submit_prompt(prompt, **kwargs)
                    self.log(title="LLM Response", message=llm_response)
                except Exception as e:
                    return f"Error running intermediate SQL: {self.extract_sql(llm_response)}"


        return self.extract_sql(llm_response)

    def run_sql_generation_pipeline(self, question: str, max_retries: int = 2, **kwargs) -> SQLGenerationState:
        """
        Run the SQL generation pipeline and return the full state (for callers that need execution_success, errors, etc.).

        Args:
            question: The natural language question to convert to SQL.
            max_retries: Maximum number of retry attempts for SQL generation.
            **kwargs: Passed to context retrieval; can include ddl_list, question_sql_list, doc_list for per-agent context.

        Returns:
            SQLGenerationState: The pipeline state with context, plan, generation_result, errors.
        """
        pipeline = SQLGenerationPipeline(askme_instance=self, max_retries=max_retries)
        return pipeline.run(question, **kwargs)

    def generate_sql_with_nodes(self, question: str, max_retries: int = 2, **kwargs) -> str:
        """
        Generate SQL using a node-based pipeline architecture.

        This method uses a two-node pipeline:
        - PlannerNode: Creates a structured SQL plan from context
        - GeneratorNode: Generates SQL based on plan, tests with LIMIT 1, retries on failure

        Example:
        ```python
        sql = self.generate_sql_with_nodes("What are the top 10 customers by sales?")
        ```

        Args:
            question (str): The question to generate a SQL query for.
            max_retries (int): Maximum number of retry attempts for SQL generation (default: 2).
            **kwargs: Additional keyword arguments passed to context retrieval.

        Returns:
            str: The SQL query that answers the question, or an error message if generation fails.
        """
        self.log(f"Starting node-based SQL generation for: {question}", title="NodePipeline")

        state = self.run_sql_generation_pipeline(question, max_retries=max_retries, **kwargs)

        # Check if generation was successful
        if state.generation_result and state.generation_result.execution_success:
            self.log(f"SQL generated successfully after {state.retry_count + 1} attempt(s)", title="NodePipeline")
            return state.generation_result.sql

        # If validation passed but execution failed, still return the SQL
        if state.generation_result and state.generation_result.is_valid and state.generation_result.sql:
            self.log(
                f"SQL generated but execution test failed. Returning SQL anyway. Error: {state.generation_result.error_message}",
                title="NodePipeline"
            )
            return state.generation_result.sql

        # If generation failed completely, log errors and return message
        if state.errors:
            error_summary = "; ".join(state.errors[-3:])  # Last 3 errors
            self.log(f"SQL generation failed: {error_summary}", title="NodePipeline")
            return f"SQL generation failed after {state.retry_count + 1} attempts. Last error: {state.errors[-1] if state.errors else 'Unknown error'}"

        return "SQL generation failed: No SQL could be generated"



    def extract_sql(self, llm_response: str) -> str:
        """
        Example:
        ```python
        self.extract_sql("Here's the SQL query in a code block: ```sql\nSELECT * FROM customers\n```")
        ```

        Extracts the SQL query from the LLM response. This is useful in case the LLM response contains other information besides the SQL query.
        Override this function if your LLM responses need custom extraction logic.

        Args:
            llm_response (str): The LLM response.

        Returns:
            str: The extracted SQL query.
        """

        import re
        """
        Extracts the SQL query from the LLM response, handling various formats including:
        - WITH clause
        - SELECT statement
        - CREATE TABLE AS SELECT
        - Markdown code blocks
        """

        # Match CREATE TABLE ... AS SELECT
        sqls = re.findall(r"\bCREATE\s+TABLE\b.*?\bAS\b.*?;", llm_response, re.DOTALL | re.IGNORECASE)
        if sqls:
            sql = sqls[-1]
            self.log(title="Extracted SQL", message=f"{sql}")
            return sql

        # Match WITH clause (CTEs)
        sqls = re.findall(r"\bWITH\b .*?;", llm_response, re.DOTALL | re.IGNORECASE)
        if sqls:
            sql = sqls[-1]
            self.log(title="Extracted SQL", message=f"{sql}")
            return sql

        # Match SELECT ... ;
        sqls = re.findall(r"\bSELECT\b .*?;", llm_response, re.DOTALL | re.IGNORECASE)
        if sqls:
            sql = sqls[-1]
            self.log(title="Extracted SQL", message=f"{sql}")
            return sql

        # Match ```sql ... ``` blocks
        sqls = re.findall(r"```sql\s*\n(.*?)```", llm_response, re.DOTALL | re.IGNORECASE)
        if sqls:
            sql = sqls[-1].strip()
            self.log(title="Extracted SQL", message=f"{sql}")
            return sql

        # Match any ``` ... ``` code blocks
        sqls = re.findall(r"```(.*?)```", llm_response, re.DOTALL | re.IGNORECASE)
        if sqls:
            sql = sqls[-1].strip()
            self.log(title="Extracted SQL", message=f"{sql}")
            return sql

        return llm_response

    def sql_has_oracle_bind_placeholders(self, sql: str) -> bool:
        """
        True when SQL contains :name bind placeholders outside string literals.
        Oracle execution uses cursor.execute(sql) with no bind dict — ORA-01008 if any remain.
        """
        if "ORACLE" not in (self.dialect or "").upper():
            return False
        if not sql or not isinstance(sql, str):
            return False
        stripped = re.sub(r"'(?:[^']|'')*'", "''", sql)
        return bool(re.search(r":\s*[a-zA-Z_][a-zA-Z0-9_]*", stripped))

    def is_sql_valid(self, sql: str) -> bool:
        """
        Example:
        ```python
        self.is_sql_valid("SELECT * FROM customers")
        ```
        Checks if the SQL query is valid. This is usually used to check if we should run the SQL query or not.
        By default it checks if the SQL query is a SELECT statement. You can override this method to enable running other types of SQL queries.

        Args:
            sql (str): The SQL query to check.

        Returns:
            bool: True if the SQL query is valid, False otherwise.
        """

        if self.sql_has_oracle_bind_placeholders(sql):
            return False

        parsed = sqlparse.parse(sql)

        for statement in parsed:
            if statement.get_type() == 'SELECT':
                return True

        return False

    def should_generate_chart(self, df: pd.DataFrame) -> bool:
        """
        Example:
        ```python
        self.should_generate_chart(df)
        ```

        Checks if a chart should be generated for the given DataFrame. By default, it checks if the DataFrame has more than one row and has numerical columns.
        You can override this method to customize the logic for generating charts.

        Args:
            df (pd.DataFrame): The DataFrame to check.

        Returns:
            bool: True if a chart should be generated, False otherwise.
        """

        if len(df) > 1 and df.select_dtypes(include=['number']).shape[1] > 0:
            return True

        return False

    def generate_rewritten_question(self, last_question: str, new_question: str, **kwargs) -> str:
        """
        **Example:**
        ```python
        rewritten_question = self.generate_rewritten_question("Who are the top 5 customers by sales?", "Show me their email addresses")
        ```

        Generate a rewritten question by combining the last question and the new question if they are related. If the new question is self-contained and not related to the last question, return the new question.

        Args:
            last_question (str): The previous question that was asked.
            new_question (str): The new question to be combined with the last question.
            **kwargs: Additional keyword arguments.

        Returns:
            str: The combined question if related, otherwise the new question.
        """
        if last_question is None:
            return new_question

        prompt = [
            self.system_message("Your goal is to combine a sequence of questions into a singular question if they are related. If the second question does not relate to the first question and is fully self-contained, return the second question. Return just the new combined question with no additional explanations. The question should theoretically be answerable with a single SQL statement."),
            self.user_message("First question: " + last_question + "\nSecond question: " + new_question),
        ]

        return self.submit_prompt(prompt=prompt, **kwargs)

    def generate_followup_questions(
        self, question: str, sql: str, columns: list, n_questions: int = 10, **kwargs
    ) -> list:
        """
        **Example:**
        ```python
        self.generate_followup_questions("What are the top 10 customers by sales?", sql, ["customer_name", "total_sales"])
        ```

        Generate a list of followup questions that you can ask Askme360.

        Args:
            question (str): The question that was asked.
            sql (str): The LLM-generated SQL query.
            columns (list): The column names from the SQL query results.
            n_questions (int): Number of follow-up questions to generate.

        Returns:
            list: A list of followup questions that you can ask Askme360.
        """

        message_log = [
            self.system_message(
                f"You are a helpful data assistant. The user asked the question: '{question}'\n\nThe SQL query for this question was: {sql}\n\nThe following are the column names from the query results: \n{columns}\n\n"
            ),
            self.user_message(
                f"Generate a list of {n_questions} followup questions that the user might ask about this data. Respond with a list of questions, one per line. Do not answer with any explanations -- just the questions. Remember that there should be an unambiguous SQL query that can be generated from the question. Prefer questions that are answerable outside of the context of this conversation. Prefer questions that are slight modifications of the SQL query that was generated that allow digging deeper into the data. Each question will be turned into a button that the user can click to generate a new SQL query so don't use 'example' type questions. Each question must have a one-to-one correspondence with an instantiated SQL query." +
                self._response_language()
            ),
        ]

        kwargs["llm_scoped_feature"] = "suggested"
        llm_response = self.submit_prompt(message_log, **kwargs)

        numbers_removed = re.sub(r"^\d+\.\s*", "", llm_response, flags=re.MULTILINE)
        return numbers_removed.split("\n")

    @staticmethod
    def _humanize_column_label(name: Any) -> str:
        """Turn SQL-style or snake_case field names into plain labels for user-facing talkback text."""
        s = str(name).strip()
        if not s:
            return s
        s = " ".join(s.replace("_", " ").split())
        if "_" in str(name) or s.isupper():
            return s.title()
        return s

    def _format_columns_for_continuation_prompt(self, columns: list) -> str:
        """List fields for the LLM: prefer human-readable labels, not SCREAMING_SNAKE_CASE."""
        if not columns:
            return "(no columns listed)\n"
        parts: List[str] = [
            "Human-readable field names (when you mention a data field, use only these phrasings—"
            "plain language, no ALL_CAPS, no underscores, no raw SQL identifiers):"
        ]
        for c in columns:
            h = self._humanize_column_label(c)
            parts.append(f"  - {h}")
        return "\n".join(parts) + "\n"

    def _format_columns_for_continuation_nonqwen_grounding(self, columns: list) -> str:
        """
        For OpenAI / non-Qwen continuation: ground follow-ups in column semantics without
        echoing them in the user-facing four lines.
        """
        col_seq = [c for c in (columns or []) if c is not None and str(c).strip()]
        if not col_seq:
            return ""
        parts: List[str] = [
            "Result schema (internal: use only to know what the answer might cover; do NOT name or copy "
            "any of the words or phrases below into the four user-facing suggestions; use generic questions, "
            "and only the user question, answer text, and SQL style):"
        ]
        for c in col_seq:
            parts.append(f"  - {self._humanize_column_label(c)}")
        return "\n".join(parts) + "\n"

    def _format_df_for_continuation_prompt(
        self, df: pd.DataFrame, *, hide_column_names: bool = False
    ) -> str:
        """Compact tabular snippet so follow-up suggestions can reference actual values.

        When ``hide_column_names`` is True (continuation talkback), column headers are replaced
        with generic placeholders so suggested questions are not copied from result field names.
        """
        if df is None or df.empty:
            return ""
        target_rows = 50
        df_sample = df.head(target_rows).copy()
        col_note = ""
        if len(df_sample.columns) > 20:
            df_sample = df_sample.iloc[:, :20]
            col_note = " (first 20 columns only)"
        if hide_column_names:
            df_sample.columns = [
                f"col_{i + 1}" for i in range(len(df_sample.columns))
            ]
        try:
            markdown_content = df_sample.to_markdown(index=False)
        except Exception:
            markdown_content = str(df_sample)
        if len(markdown_content) > 8000:
            markdown_content = markdown_content[:8000] + "\n... [truncated]"
        value_hint = (
            "Use cell values, categories, and numbers for context only; do not name col_1, col_2, or 'column' headers in the user-facing four suggestions. "
            if hide_column_names
            else
            "Base your follow-up suggestions on this data — reference specific values, categories, or comparisons visible below when possible: "
        )
        return (
            f"\n\nTABULAR RESULT DATA (up to {target_rows} rows){col_note}.\n"
            f"Shape: {df.shape}\n"
            + value_hint
            + "\n"
            f"{markdown_content}"
        )

    def _continuation_sql_api_snippet(
        self,
        sql: Optional[str] = None,
        api_endpoint: Optional[str] = None,
        api_params: Optional[str] = None,
        api_response_validator: Optional[str] = None,
    ) -> str:
        """Optional SQL / API metadata so follow-ups stay consistent with how data was fetched."""
        parts: List[str] = []
        if sql:
            parts.append(
                "GENERATED SQL (propose follow-ups that could be answered by filters, different GROUP BYs, time windows, joins, limits, or drill-downs consistent with this query):\n"
                + sql
            )
        if api_endpoint or api_params or api_response_validator:
            ap = ["API CALL CONTEXT (propose follow-ups that stay consistent with this endpoint and parameters):"]
            if api_endpoint:
                ap.append(f"Endpoint(s):\n{api_endpoint}")
            if api_params:
                ap.append(f"Parameters / request body:\n{api_params}")
            if api_response_validator:
                ap.append(f"Response transformation / validator:\n{api_response_validator}")
            parts.append("\n\n".join(ap))
        return "\n\n".join(parts)

    @staticmethod
    def _polish_continuation_suggestion_line_body(body: str) -> str:
        """
        Make each suggested follow-up copy-paste friendly: remove vague hedges the model
        still sometimes emits (e.g. "any", "particular") on user-facing lines.
        """
        t = (body or "").strip()
        if not t:
            return t
        t = re.sub(r"(?i)\bany\s+months\b", "which months", t)
        t = re.sub(r"(?i)\bany\s+quarters?\b", "which quarters", t)
        t = re.sub(r"(?i)\bany\s+weeks?\b", "which weeks", t)
        t = re.sub(r"(?i)\bany\s+days?\b", "which days", t)
        t = re.sub(r"(?i)\bidentify\s+any\s+", "List ", t)
        t = re.sub(r"(?i)\bfind\s+any\s+", "Find ", t)
        t = re.sub(r"(?i)\bhighlight\s+any\s+", "Highlight ", t)
        t = re.sub(r"(?i)\bshow\s+any\s+", "Show ", t)
        t = re.sub(r"(?i)\bparticular\s+", "", t)
        t = re.sub(r"(?i)\ba\s+particular\s+", "", t)
        t = re.sub(r"  +", " ", t).strip()
        return t

    @classmethod
    def _scrub_continuation_talkback_output(cls, text: str) -> str:
        """Remove placeholder column tokens and backtick-wrapped identifiers from talkback text."""
        t = (text or "").strip()
        t = re.sub(r"`+[^`\r\n]+`+", "", t)
        t = re.sub(r"\bcol_\d+\b", "", t, flags=re.IGNORECASE)
        # Polish numbered suggestion lines for copy-paste ready phrasing
        lines = t.split("\n")
        out_lines: List[str] = []
        for line in lines:
            m = re.match(r"^(\s*(?:\d+)[.)]\s*)(.+)$", line)
            if m:
                prefix, body = m.group(1), m.group(2)
                out_lines.append(prefix + cls._polish_continuation_suggestion_line_body(body))
            else:
                out_lines.append(line)
        t = "\n".join(out_lines)
        t = re.sub(r"  +", " ", t)
        t = re.sub(r"\n{3,}", "\n\n", t)
        return t.strip()

    def generate_continuation_talkback(
        self,
        question: str,
        answer_context: str,
        columns: list,
        n_items: int = 3,
        df: Optional[pd.DataFrame] = None,
        sql: Optional[str] = None,
        api_endpoint: Optional[str] = None,
        api_params: Optional[str] = None,
        api_response_validator: Optional[str] = None,
        use_qwen: bool = False,
        **kwargs,
    ) -> str:
        """
        **Example:**
        ```python
        self.generate_continuation_talkback("What are the top 10 customers?", "Data shown.", ["name", "sales"])
        ```

        Generate a conversational continuation talkback as a single paragraph
        containing proactive suggestions/questions.

        Args:
            question (str): The original user query.
            answer_context (str): A summary of the results or data context.
            columns (list): The list of column names (if available).
            n_items (int): Ignored; exactly three suggestions are always generated.
            df (pd.DataFrame, optional): When ``use_qwen`` is true, result rows (anonymized headers) ground suggestions.
            sql (str, optional): Executed or generated SQL — use so follow-ups fit the same schema/query style.
            api_endpoint, api_params, api_response_validator (str, optional): API call metadata when data came from an API.
            use_qwen (bool): If true, include a tabular data snippet from ``df``; if false, use ``columns``-only
                schema hint (no row table) for continuation talkback.

        Returns:
            str: A conversational string with n_items suggestions; lines avoid naming result columns or schema tokens.
        """
        use_qwen_flag = bool(use_qwen or kwargs.get("use_qwen", False))
        df_effective: Optional[pd.DataFrame] = None
        if use_qwen_flag and df is not None and isinstance(df, pd.DataFrame) and not df.empty:
            df_effective = df

        df_snippet = ""
        if df_effective is not None:
            df_snippet = self._format_df_for_continuation_prompt(
                df_effective, hide_column_names=True
            )

        columns_grounding = ""
        if df_effective is None:
            columns_grounding = self._format_columns_for_continuation_nonqwen_grounding(
                columns
            )

        sql_api_snippet = self._continuation_sql_api_snippet(
            sql=sql,
            api_endpoint=api_endpoint,
            api_params=api_params,
            api_response_validator=api_response_validator,
        )

        n = 3
        if df_effective is not None:
            grounding = (
                "When tabular result data is shown above, each suggestion must be a standalone next question the user could paste into the chat: short, direct, and phrased like a real user query (e.g. starts with Show, Which, How, Compare, or a clear imperative). "
                f"When GENERATED SQL is shown, use it only to infer query style; do not copy table, column, or alias names from it into the {n} user-facing lines. "
                "When API CALL CONTEXT is shown, follow-ups may stay in the same API family without quoting parameter or path tokens as if they were labels. "
                "Do not invent facts that contradict the visible data."
            )
        else:
            grounding = (
                "Rely on the context/answer, the result schema line items (internal), and (if present) SQL or API context. "
                f"The schema list is for your reasoning only: do not copy those labels into the {n} user-facing lines. "
                "When GENERATED SQL is shown, use it only to infer query style; do not copy table, column, or alias names. "
                "When API CALL CONTEXT is shown, follow-ups may stay in the same family without echoing param keys or path segments as if they were user labels. "
                "Do not invent facts that contradict the visible data. "
                f"Each of the {n} suggestions must be a self-contained, paste-ready user message (a single sentence), not a capability description; prefer 'Which month had the lowest sales?' over 'I can help explore reasons for the lowest month.'"
            )

        n_lines = "\n".join(f"{i}. [Suggestion {i}]" for i in range(1, n + 1))
        message_log = [
            self.system_message(
                f"You are a helpful and proactive data assistant. The user previously asked: '{question}'\n\n"
                f"The system returned the following context/answer:\n{answer_context}\n\n"
                f"{df_snippet}{columns_grounding}\n"
                + (f"\n{sql_api_snippet}\n" if sql_api_snippet else "\n")
                + "\nYour task is to generate a continuation response as a single conversational text block. "
                f"The response must start with 'If you want, I can also:' and then provide exactly {n} proactive suggestions or follow-up questions. "
                + grounding
                + "\n\n"
                f"Format the {n} items as a list within the paragraph, separated by newlines as natural. "
                "Use this structure:\n"
                "If you want, I can also:\n"
                "1. [Suggestion 1]\n"
                "2. [Suggestion 2]\n"
                "3. [Suggestion 3]\n\n"
                "Do not return JSON. Just return the plain text conversational paragraph.\n"
                f"Copy-paste requirement: each of the {n} lines (after the numbering) must read as a full next question the user can paste as-is. "
                "Use concrete wording (e.g. Which month, Compare H1 to H2, List months with a spike) — not vague hedges. "
                "Do NOT use the words 'any' or 'particular' in the suggestions (e.g. say 'Which months had unusual spikes' not 'Identify any months'). Do not use 'certain' or 'some' as filler. "
                "Each line must be a direct, specific question or imperative without trailing conversational filler — do NOT include 'e.g.', 'for example', 'such as', or 'if you\\'d like a deeper analysis' inside a suggestion. "
                f"CRITICAL: The {n} suggestions must not name or quote result columns, measures, field labels, or placeholder headers (e.g. col_1, col_2), and must not repeat SQL or API identifier tokens. "
                "If the context/answer text above names specific fields or columns, do not echo those names in the lines. "
                "You may still refer to the topic of the user question, entity names or category *values* that appear in the data (e.g. a region or product name), and generic analysis terms (trend, breakdown, top contributors, year-over-year). "
                "Do not use backticks. Do not use parenthetical example lists like (e.g., ...)."
            ),
            self.user_message(
                f"Generate a proactive continuation talkback with {n} items." +
                self._response_language()
            ),
        ]

        cont_kwargs = dict(kwargs)
        cont_kwargs["llm_scoped_feature"] = "continuation"
        llm_response = self.submit_prompt(message_log, **cont_kwargs)

        return self._scrub_continuation_talkback_output((llm_response or "").strip())

    def generate_questions(self, **kwargs) -> List[str]:
        """
        **Example:**
        ```python
        self.generate_questions()
        ```

        Generate a list of questions that you can ask Askme360.
        """
        question_sql = self.get_similar_question_sql(question="", **kwargs)

        return [q["question"] for q in question_sql]

    def _format_df_excerpt_for_summary(
        self,
        df: pd.DataFrame,
        *,
        max_rows: int = 35,
        max_cols: int = 25,
        max_chars: int = 7500,
    ) -> str:
        """Compact tabular text for LLM summary prompts (OpenAI, OIC, Qwen, fallbacks)."""
        if df is None or df.empty:
            return ""
        sample = df.head(max_rows).copy()
        if len(sample.columns) > max_cols:
            sample = sample.iloc[:, :max_cols]
        try:
            md = sample.to_markdown(index=False)
        except Exception:
            md = str(sample)
        if len(md) > max_chars:
            md = md[:max_chars] + "\n... [truncated]"
        return (
            f"\n\nResult set: {len(df)} row(s), {df.shape[1]} column(s).\n{md}\n"
        )

    def generate_summary(self, question: str, columns, df=None, **kwargs) -> str:
        """
        Generate a summary of the results of a SQL query or API call.
        Single-agent implementation.

        Args:
            question (str): The question that was asked.
            columns (list): List of column names, e.g. ["col1", "col2"]
            df (pd.DataFrame, optional): The actual data to summarize.

        Returns:
            str: The summary of the results of the query.
        """
        # Prepare data-rich info ONLY if use_qwen is explicitly requested
        df_info = ""
        if kwargs.get("use_qwen") and df is not None and not df.empty:
            if len(df) > 50:
                df_sample = df.head(50)
                df_info = f"\n\nDATAFRAME CONTENT (UP TO 50 ROWS):\nDataFrame shape: {df.shape}\nSample data:\n{df_sample.to_markdown(index=False)}"
            else:
                df_info = f"\n\nDATAFRAME CONTENT:\nDataFrame shape: {df.shape}\nData:\n{df.to_markdown(index=False)}"
        
        # Pass data info through dedicated keyword for local Qwen (Ollama) injection
        if df_info:
            kwargs["qwen_data"] = df_info

        # Row-level excerpt for every backend: non-Qwen paths used to be metadata-only, which led to
        # "no data" summaries when the client sent rows but empty/missing column lists or on Qwen fallback.
        use_qwen = kwargs.get("use_qwen", False)
        from app.services.llm_service import _qwen_configured
        
        row_context = (
            self._format_df_excerpt_for_summary(df)
            if df is not None and not df.empty and (use_qwen and _qwen_configured())
            else ""
        )

        has_data = df is not None and not df.empty
        if has_data and not row_context:
            privacy_note = "\nNote: Data was successfully retrieved, but the raw rows are hidden for privacy. You must use the available fields listed above to describe what was found. Do NOT state that no data was returned or no records were found."
        else:
            privacy_note = "\nIf a result table is shown above, your summary must reflect that data; do not state that no data was returned."

        # Base message_log
        message_log = [
            self.system_message(
                f"You are a helpful data assistant. The user asked: '{question}'\n\n"
                f"The query successfully returned records with the following fields: \n{columns}\n"
                f"{row_context}"
                f"{privacy_note}"
            ),
            self.user_message(
                "Summarize the results based on the question and the available fields. "
                "The summary must be comprehensive but concise, between 20 to 50 words strictly. "
                "Do not respond with any additional explanation beyond the summary." +
                self._response_language()
            ),
        ]
        kwargs["llm_scoped_feature"] = "summary"
        summary = self.submit_prompt(message_log, **kwargs)
        return summary

    def generate_summary_multi_agent(self, question: str, columns, df=None, **kwargs) -> str:
        """
        Generate a summary of the results of a SQL query or API call.
        Multi-agent implementation.

        Args:
            question (str): The question that was asked.
            columns (List[dict]): List of dicts with keys "agent_id" and "columns",
                e.g. [{"agent_id": "agent1", "columns": ["col1", "col2"]}, ...]
            df (pd.DataFrame, optional): Combined DataFrame from all agents.

        Returns:
            str: The summary of the results of the query.
        """
        # Multi-agent: format each agent's columns concisely
        agents_summary = []
        for agent_data in columns:
            name = agent_data.get("agent_name") or agent_data.get("agent_id", "Agent")
            cols = agent_data.get("columns", [])
            # Only include up to 15 columns per agent to keep prompt size sane
            col_preview = f"{cols[:15]}..." if len(cols) > 15 else str(cols)
            agents_summary.append(f"  - {name}: {col_preview}")
        
        columns_description = "Agents returned these columns:\n" + "\n".join(agents_summary)

        agent_data_blocks = kwargs.pop("agent_data_blocks", None) or []
        model_override = kwargs.pop("model", None)

        # Detect weak schema overlap (common when datasets are unrelated).
        normalized_sets: List[set] = []
        for agent_data in columns:
            cols = agent_data.get("columns", []) or []
            if isinstance(cols, list):
                normalized_sets.append({str(c).strip().lower() for c in cols if str(c).strip()})
        overlap_ratio = 1.0
        if len(normalized_sets) >= 2:
            union_cols = set().union(*normalized_sets) if normalized_sets else set()
            inter_cols = set.intersection(*normalized_sets) if normalized_sets else set()
            overlap_ratio = (len(inter_cols) / max(1, len(union_cols))) if union_cols else 0.0
        disjoint_mode = overlap_ratio < 0.20
        
        # Prepare data-rich info ONLY for local Qwen path (df/blocks omitted upstream when Qwen unavailable).
        df_info = ""
        if kwargs.get("use_qwen") and df is not None and not df.empty:
            # Row limit (requested: 100)
            target_rows = 100
            df_sample = df.head(target_rows)
            
            # Trim columns if too wide (prevents massive Markdown payloads)
            col_note = ""
            if len(df_sample.columns) > 20:
                df_sample = df_sample.iloc[:, :20]
                col_note = " (TRUNCATED TO FIRST 20 COLUMNS)"
            
            # Format as Markdown with character cap (approx 8000 chars fits well in 8k context)
            markdown_content = df_sample.to_markdown(index=False)
            if len(markdown_content) > 8000:
                markdown_content = markdown_content[:8000] + "\n... [TRUNCATED DUE TO SIZE]"
                
            df_info = f"\n\nCOMBINED DATAFRAME CONTENT (UP TO {target_rows} ROWS){col_note}:\nDataFrame shape: {df.shape}\nData:\n{markdown_content}"

        use_qwen = kwargs.get("use_qwen", False)
        from app.services.llm_service import _qwen_configured
        
        # Add per-agent samples so the model can reason on unrelated datasets safely.
        block_sections: List[str] = []
        if isinstance(agent_data_blocks, list) and (use_qwen and _qwen_configured()):
            for idx, block in enumerate(agent_data_blocks):
                if not isinstance(block, dict):
                    continue
                bid = block.get("agent_id", f"block_{idx+1}")
                rows = block.get("data", []) if isinstance(block.get("data"), list) else []
                rows = [r for r in rows if isinstance(r, dict)]
                if not rows:
                    continue
                headers = block.get("headers", []) if isinstance(block.get("headers"), list) else []
                if not headers:
                    seen = set()
                    for row in rows:
                        for k in row.keys():
                            ks = str(k)
                            if ks not in seen:
                                seen.add(ks)
                                headers.append(ks)
                # Keep prompt bounded.
                sample_df = pd.DataFrame(rows[:40])
                if not sample_df.empty and len(sample_df.columns) > 20:
                    sample_df = sample_df.iloc[:, :20]
                sample_md = sample_df.to_markdown(index=False) if not sample_df.empty else ""
                if len(sample_md) > 3500:
                    sample_md = sample_md[:3500] + "\n... [TRUNCATED DUE TO SIZE]"
                block_sections.append(
                    f"Agent {bid}\nHeaders: {headers[:20]}\nRows: {len(rows)}\nSample:\n{sample_md}"
                )

        blocks_info = ""
        if block_sections:
            blocks_info = "\n\nPER-AGENT DATA BLOCKS:\n" + "\n\n---\n\n".join(block_sections)
            if kwargs.get("use_qwen"):
                if df_info:
                    df_info = (df_info + "\n\n" + blocks_info).strip()
                else:
                    df_info = blocks_info.strip()

        # Pass data info through dedicated keyword for local Qwen (Ollama) injection
        if df_info:
            kwargs["qwen_data"] = df_info

        if block_sections:
            row_context = blocks_info
        elif df is not None and not df.empty and (use_qwen and _qwen_configured()):
            row_context = self._format_df_excerpt_for_summary(df)
        else:
            row_context = ""

        summary_instruction = (
            "The datasets are from different agents and may be unrelated. "
            "Produce one concise insight per agent first, then a combined conclusion. "
            if disjoint_mode
            else "Summarize the combined data from all agents based on the question that was asked. "
            "Highlight insights from each data source effectively. "
        )

        has_data = df is not None and not df.empty
        if has_data and not row_context:
            privacy_note = "\nNote: Data was successfully retrieved, but the raw rows are hidden for privacy. You must use the available agent fields listed above to describe what was found. Do NOT state that no data was returned or no records were found."
        else:
            privacy_note = "\nIf a result table is shown above, your summary must reflect that data; do not state that no data was returned."

        # Base message_log (includes tabular samples for OIC/OpenAI/Gemini, not only Qwen)
        message_log = [
            self.system_message(
                f"You are a helpful data assistant. The user asked: '{question}'\n\n"
                f"The following agents successfully returned records with these fields:\n{columns_description}\n"
                f"{row_context}"
                f"{privacy_note}"
            ),
            self.user_message(
                summary_instruction
                + "The summary must be between 20 to 50 words strictly. "
                "Do not respond with any additional explanation beyond the summary." +
                self._response_language()
            ),
        ]
        # Non-Qwen backends (OIC/OpenAI/Gemini): never append row samples — metadata is in the system message only.
        kwargs["llm_scoped_feature"] = "summary"
        if model_override:
            kwargs["model"] = model_override
        summary = self.submit_prompt(message_log, **kwargs)
        return summary

    # ----------------- Use Any Embeddings API ----------------- #
    @abstractmethod
    def generate_embedding(self, data: str, **kwargs) -> List[float]:
        pass

    # ----------------- Use Any Database to Store and Retrieve Context ----------------- #
    @abstractmethod
    def get_similar_question_sql(self, question: str, **kwargs) -> list:
        """
        This method is used to get similar questions and their corresponding SQL statements.

        Args:
            question (str): The question to get similar questions and their corresponding SQL statements for.

        Returns:
            list: A list of similar questions and their corresponding SQL statements.
        """
        pass

    @abstractmethod
    def get_related_ddl(self, question: str, **kwargs) -> list:
        """
        This method is used to get related DDL statements to a question.

        Args:
            question (str): The question to get related DDL statements for.

        Returns:
            list: A list of related DDL statements.
        """
        pass

    @abstractmethod
    def get_related_documentation(self, question: str, **kwargs) -> list:
        """
        This method is used to get related documentation to a question.

        Args:
            question (str): The question to get related documentation for.

        Returns:
            list: A list of related documentation.
        """
        pass

    @abstractmethod
    def add_question_sql(self, question: str, sql: str, **kwargs) -> str:
        """
        This method is used to add a question and its corresponding SQL query to the training data.

        Args:
            question (str): The question to add.
            sql (str): The SQL query to add.

        Returns:
            str: The ID of the training data that was added.
        """
        pass

    @abstractmethod
    def add_ddl(self, ddl: str, **kwargs) -> str:
        """
        This method is used to add a DDL statement to the training data.

        Args:
            ddl (str): The DDL statement to add.

        Returns:
            str: The ID of the training data that was added.
        """
        pass

    @abstractmethod
    def add_documentation(self, documentation: str, **kwargs) -> str:
        """
        This method is used to add documentation to the training data.

        Args:
            documentation (str): The documentation to add.

        Returns:
            str: The ID of the training data that was added.
        """
        pass

    @abstractmethod
    def get_training_data(self, **kwargs) -> pd.DataFrame:
        """
        Example:
        ```python
        self.get_training_data()
        ```

        This method is used to get all the training data from the retrieval layer.

        Returns:
            pd.DataFrame: The training data.
        """
        pass

    @abstractmethod
    def remove_training_data(self, id: str, **kwargs) -> bool:
        """
        Example:
        ```python
        self.remove_training_data(id="123-ddl")
        ```

        This method is used to remove training data from the retrieval layer.

        Args:
            id (str): The ID of the training data to remove.

        Returns:
            bool: True if the training data was removed, False otherwise.
        """
        pass

    # ----------------- Use Any Language Model API ----------------- #

    @abstractmethod
    def system_message(self, message: str) -> any:
        pass

    @abstractmethod
    def user_message(self, message: str) -> any:
        pass

    @abstractmethod
    def assistant_message(self, message: str) -> any:
        pass

    def str_to_approx_token_count(self, string: str) -> int:
        return len(string) / 4

    def merge_sql_query_summaries_into_details(
        self, response_details: Dict[str, Any], user_question: str
    ) -> None:
        """
        Mutate ``response_details`` with ``sql_query_summaries`` (and optional
        ``sql_query_summaries_error``) from ``details['sql']`` envelope + this agent's LLM
        (``submit_prompt`` / same keys as other AskmeBase LLM calls). Question + SQL only — no row data.
        """
        if not sql_plain_summary_enabled():
            return
        items = database_sql_items_from_details_sql(response_details.get("sql"))
        if not items:
            return
        summaries, err = self._explain_sql_queries_batched(user_question, items)
        response_details["sql_query_summaries"] = summaries or []
        if err:
            response_details["sql_query_summaries_error"] = err

    def _explain_sql_queries_batched(
        self, user_question: str, items: List[Dict[str, str]]
    ) -> Tuple[List[Dict[str, Any]], Optional[str]]:
        """
        One LLM call. Returns (sql_query_summaries, error_message_or_none). Never raises.
        """
        if not items:
            return [], None
        blocks: List[str] = []
        for i, it in enumerate(items, start=1):
            blocks.append(
                f"--- SQL {i} ---\n"
                f"agent_id: {it.get('agent_id', '')}\n"
                f"sub_question: {it.get('sub_question', '')}\n"
                f"sql:\n{it.get('sql_text', '')}\n"
            )
        user_block = (
            f"User question (context only; do not invent facts from data not shown):\n{user_question}\n\n"
            + "\n".join(blocks)
        )
        system = (
            "You explain SQL queries in plain English for non-technical readers. "
            "For each SQL block, identify query type and intent from table/column names. "
            "Use actual values from the SQL (tables, columns, filters, limits) and never placeholders.\n\n"
            'Write the "summary" in this exact flow (content style similar to a business report):\n'
            "1) What This Query Does (Plain English): 1 short paragraph.\n"
            "2) Step sections for complex SQL (e.g., Step 1 - ..., Step 2 - ...). For simple SQL, skip steps.\n"
            "3) If output columns are obvious, add a small table: What it shows | What it means.\n"
            "4) The Final Output: 1 short paragraph.\n"
            "5) Quick Summary Table: markdown table with columns Clause | Value Used | Purpose (only clauses present).\n\n"
            "Safety and style: never invent query results or row values; explain jargon in plain words; "
            "warn if UPDATE/DELETE has no WHERE; warn if DROP TABLE appears; keep tone friendly and concise.\n\n"
            "Return ONLY a JSON array (no markdown fences), same length and order as the SQL blocks. "
            'Each element must be an object with keys: "agent_id" (string), "sub_question" (string, may be empty), '
            '"summary" (string), and "summary_html" (string).\n\n'
            "summary_html must be valid, fully renderable HTML that mirrors summary content. "
            "Target this structure and order: <div> wrapper, <h3>What This Query Does (Plain English)</h3>, "
            "<p>...</p>, optional <hr>, step headings using <h4>, optional explanatory table, "
            "<h3>The Final Output</h3>, and <h3>Quick Summary Table</h3> with a proper "
            "<table><thead><tbody><tr><th><td> structure. "
            "No markdown tokens (**, *, |) may appear in summary_html. Keep HTML clean and properly nested.\n\n"
            "agent_id and sub_question must match the corresponding block."
        )
        try:
            raw = self.submit_prompt(
                [
                    self.system_message(system),
                    self.user_message(user_block),
                ],
                temperature=0.2,
                max_tokens=min(4000, 800 * len(items)),
                generation_name="sql_plain_language_summaries",
            )
        except Exception as e:
            return [], f"llm_call:{e!s}"

        parsed = _sql_summary_parse_json_array_from_llm_text(raw or "")
        if not parsed:
            return [], "json_parse"

        summaries: List[Dict[str, Any]] = []
        for idx, it in enumerate(items):
            row = parsed[idx] if idx < len(parsed) else None
            summary = ""
            summary_html = ""
            if isinstance(row, dict):
                summary = str(row.get("summary") or "").strip()
                summary_html = str(row.get("summary_html") or "").strip()
            if not summary_html and summary:
                summary_html = _sql_summary_markdown_to_html(summary)
            sq = (it.get("sub_question") or "").strip()
            row_out: Dict[str, Any] = {
                "agent_id": it.get("agent_id", ""),
                "sub_question": sq,
                "summary": summary,
                "summary_html": summary_html,
            }
            summaries.append(row_out)
        return summaries, None

    def add_ddl_to_prompt(
        self, initial_prompt: str, ddl_list: list[str], max_tokens: int = 40000
    ) -> str:
        if len(ddl_list) > 0:
            initial_prompt += "\n===Tables \n"

            for ddl in ddl_list:
                if (
                    self.str_to_approx_token_count(initial_prompt)
                    + self.str_to_approx_token_count(ddl)
                    < max_tokens
                ):
                    initial_prompt += f"{ddl}\n\n"

        return initial_prompt

    def add_documentation_to_prompt(
        self,
        initial_prompt: str,
        documentation_list: list[str],
        max_tokens: int = 40000,
    ) -> str:
        if len(documentation_list) > 0:
            initial_prompt += "\n===Additional Context \n\n"
            initial_prompt += "Always follow :- Give a single valid SQL query to answer the question. do not give any explanations or comments.: \n"

            for documentation in documentation_list:
                if (
                    self.str_to_approx_token_count(initial_prompt)
                    + self.str_to_approx_token_count(documentation)
                    < max_tokens
                ):
                    initial_prompt += f"{documentation}\n\n"

        return initial_prompt

    def add_sql_to_prompt(
        self, initial_prompt: str, sql_list: list[str], max_tokens: int = 40000
    ) -> str:
        if len(sql_list) > 0:
            initial_prompt += "\n===Question-SQL Pairs\n\n"

            for question in sql_list:
                if (
                    self.str_to_approx_token_count(initial_prompt)
                    + self.str_to_approx_token_count(question["sql"])
                    < max_tokens
                ):
                    initial_prompt += f"{question['question']}\n{question['sql']}\n\n"

        return initial_prompt

    def get_sql_prompt(
        self,
        initial_prompt : str,
        question: str,
        question_sql_list: list,
        ddl_list: list,
        doc_list: list,
        **kwargs,
    ):
        """
        Example:
        ```python
        self.get_sql_prompt(
            question="What are the top 10 customers by sales?",
            question_sql_list=[{"question": "What are the top 10 customers by sales?", "sql": "SELECT * FROM customers ORDER BY sales DESC LIMIT 10"}],
            ddl_list=["CREATE TABLE customers (id INT, name TEXT, sales DECIMAL)"],
            doc_list=["The customers table contains information about customers and their sales."],
        )

        ```

        This method is used to generate a prompt for the LLM to generate SQL.

        Args:
            question (str): The question to generate SQL for.
            question_sql_list (list): A list of questions and their corresponding SQL statements.
            ddl_list (list): A list of DDL statements.
            doc_list (list): A list of documentation.

        Returns:
            any: The prompt for the LLM to generate SQL.
        """

        if initial_prompt is None:
            initial_prompt = f"You are a {self.dialect} expert. " + \
            "Please help to generate a SQL query to answer the question. Your response should ONLY be based on the given context and follow the response guidelines and format instructions. "

        initial_prompt = self.add_ddl_to_prompt(
            initial_prompt, ddl_list, max_tokens=self.max_tokens
        )

        if self.static_documentation != "":
            doc_list.append(self.static_documentation)

        initial_prompt = self.add_documentation_to_prompt(
            initial_prompt, doc_list, max_tokens=self.max_tokens
        )

        initial_prompt += (
            "===Response Guidelines \n"
            "1. If the provided context is sufficient, please generate a single valid SQL query without any explanations for the question. \n"
            "2. If the provided context is almost sufficient but requires knowledge of a specific string in a particular column, please generate an intermediate SQL query to find the distinct strings in that column. Prepend the query with a comment saying intermediate_sql \n"
            "3. If the provided context is insufficient, please explain why it can't be generated. \n"
            "4. Please use the most relevant table(s). \n"
            "5. If the question has been asked and answered before, please repeat the answer exactly as it was given before. \n"
            f"6. Ensure that the output SQL is {self.dialect}-compliant and executable, and free of syntax errors. \n"
            f"7. If the question does not specify a row limit or pagination, add a default limit of {self.max_rows} rows using the correct syntax for {self.dialect} (e.g., LIMIT {self.max_rows} for Postgres/MySQL/SQLite/BigQuery/Snowflake; TOP {self.max_rows} for SQL Server; FETCH FIRST {self.max_rows} ROWS ONLY for Oracle). Do not add a limit if the query already includes LIMIT/OFFSET, TOP, or FETCH/NEXT. If you add a limit to an unordered query, also include a reasonable ORDER BY (e.g., primary timestamp or id) for deterministic results.\n"
            "8. For status, category, or flag columns: do not assume words in the question map to exact "
            "stored literals unless DDL or examples show them. When excluding a state or using \"not\" "
            "semantics, prefer filters consistent with the schema (e.g. inequality or NOT IN), not guessed synonyms.\n"
        )
        if "ORACLE" in (self.dialect or "").upper():
            initial_prompt += (
                "9. Oracle GROUP BY (ORA-00979): ORDER BY must not reference columns absent from GROUP BY. "
                "For bucketed aggregates, use ORDER BY 1, MIN/MAX of the column used for buckets, or repeat "
                "the exact GROUP BY expression — not ORDER BY CASE that reuses ungrouped base columns.\n"
                "10. Oracle dates: NEVER use EXTRACT(QUARTER FROM date) — it is invalid and causes ORA-00907. "
                "Use TO_CHAR(date_col, 'Q') for calendar quarter (1–4), or TRUNC(date_col, 'Q'); "
                "EXTRACT is only valid for YEAR, MONTH, DAY, etc.\n"
                "11. Oracle ORA-01008: Do NOT use bind variables (:name, :country_name). SQL is run as a plain "
                "string with no bind parameters. Use literal quoted values in WHERE (e.g. "
                "COUNTRY IN ('Sweden','United States')) when the question gives examples; never use :placeholders.\n"
            )

        message_log = [self.system_message(initial_prompt)]

        for example in question_sql_list:
            if example is None:
                print("example is None")
            else:
                if example is not None and "question" in example and "sql" in example:
                    message_log.append(self.user_message(example["question"]))
                    message_log.append(self.assistant_message(example["sql"]))

        message_log.append(self.user_message(question))

        return message_log




    def get_followup_questions_prompt(
        self,
        question: str,
        question_sql_list: list,
        ddl_list: list,
        doc_list: list,
        **kwargs,
    ) -> list:
        initial_prompt = f"The user initially asked the question: '{question}': \n\n"

        initial_prompt = self.add_ddl_to_prompt(
            initial_prompt, ddl_list, max_tokens=self.max_tokens
        )

        initial_prompt = self.add_documentation_to_prompt(
            initial_prompt, doc_list, max_tokens=self.max_tokens
        )

        initial_prompt = self.add_sql_to_prompt(
            initial_prompt, question_sql_list, max_tokens=self.max_tokens
        )

        message_log = [self.system_message(initial_prompt)]
        message_log.append(
            self.user_message(
                "Generate a list of followup questions that the user might ask about this data. Respond with a list of questions, one per line. Do not answer with any explanations -- just the questions."
            )
        )

        return message_log

    def _messages_to_prompt_text(self, prompt) -> str:
        """Convert list of message dicts (role/content) to a single string for OCI chat_completion."""
        if not prompt:
            return ""
        parts = []
        for msg in prompt:
            role = (msg.get("role") or "user").upper()
            content = msg.get("content") or ""
            parts.append(f"{role}:\n{content}")
        return "\n\n".join(parts)

    def _use_oci_llm_for_sql(self) -> bool:
        """Return True when OCI LLM should be used for SQL generation (config flag or initial_config has OCI_*)."""
        if self.config and self.config.get("use_oci_llm"):
            return True
        try:
            from app.core.config import get_initial_config
            cfg = get_initial_config()
            if cfg.get("OCI_COMPARTMENT_ID") and cfg.get("OCI_GENAI_ENDPOINT") and cfg.get("OCI_MODEL_ID"):
                return True
        except Exception:
            pass
        return False

    # def _submit_prompt_or_oci(self, prompt, **kwargs) -> str:
    #     """Call OCI chat_completion when configured for OCI, otherwise submit_prompt."""
    #     try:
    #         from askme_open_agent.askme_oci_llm.oci_llm_service import chat_completion as oci_chat_completion
    #         prompt_text = self._messages_to_prompt_text(prompt)
    #         return oci_chat_completion(prompt=prompt_text)
    #     except Exception as e:
    #         self.log(f"OCI chat_completion failed ({e}), falling back to submit_prompt", title="OCI LLM")
    #     return self.submit_prompt(prompt, **kwargs)

    def get_multi_agent_sql_prompt(self, question: str, merged_training: str, agent_ids: list) -> str:
        """Build the multi-database SQL generation prompt (SIMPLE + ADVANCED merge)."""
        a0 = agent_ids[0] if agent_ids else '1'
        a1 = agent_ids[1] if len(agent_ids) > 1 else '2'
        prompt = f"""
{merged_training}

User Question: {question}

CRITICAL INSTRUCTIONS:
1. Analyze if this question requires data from MULTIPLE databases or just ONE database
2. Table names above are prefixed with agent{{id}}_ for reference only - use ORIGINAL table names in SQL
3. Only query databases that have RELEVANT data for this question
4. Choose between SIMPLE merge (one strategy for all) or ADVANCED merge (different strategies per step)

DECISION LOGIC:
Step 1: Determine which databases contain relevant tables/data for this question
Step 2: If only ONE database is relevant:
   - Return query for ONLY that database
   - Set other database queries to empty string ""
   - Set merge_key to null
   - Set merge_strategy to "none"

Step 3: If MULTIPLE databases are relevant:
   - SIMPLE CASE: All agents can be merged with ONE strategy
     * Use single merge_key and merge_strategy
   - ADVANCED CASE: Different merge strategies needed between different agent pairs
     * Use execution_order and merge_steps with specific strategies per step

OUTPUT FORMATS:

FORMAT 1 - SIMPLE MERGE (Single strategy for all agents):
{{
    "queries": {{
        "agent{a0}": "SELECT ... FROM original_table WHERE ...",
        "agent{a1}": "SELECT ... FROM original_table WHERE ..."
    }},
    "merge_key": "common_column_name_or_null",
    "merge_strategy": "none|independent|concatenate|inner|outer|left|right",
    "reasoning": "Brief explanation of which databases are relevant and why"
}}

FORMAT 2 - ADVANCED MERGE (Different strategies per step):
{{
    "queries": {{
        "agent1": "SELECT ...",
        "agent5": "SELECT ...",
        "agent7": "SELECT ..."
    }},
    "execution_order": ["agent1", "agent5", "agent7"],
    "merge_steps": [
        {{
            "agents": ["agent1", "agent5"],
            "merge_key": "CUSTOMER_NAME",
            "merge_strategy": "outer"
        }},
        {{
            "agents": ["result_1_5", "agent7"],
            "merge_key": "ID",
            "merge_strategy": "inner"
        }}
    ],
    "reasoning": "Explanation of multi-step merge process"
}}

MERGE_STRATEGY OPTIONS:
- "none": Only one database has relevant data
- "independent": Multiple databases queried but results shown separately (NO merge, returns dict of DataFrames)
- "concatenate": Stack rows vertically from multiple agents with SIMILAR schemas (like UNION ALL)
- "inner": Merge ALL agents horizontally on matching keys only (keeps only rows with matches in all DataFrames)
- "outer": Merge ALL agents horizontally keeping all rows from all DataFrames
- "left": Keep all rows from first agent, match with others horizontally
- "right": Keep all rows from last agent, match with others horizontally

WHEN TO USE SIMPLE vs ADVANCED:
- SIMPLE: All agents share same merge key OR all need same strategy
- ADVANCED: Different merge keys/strategies needed (e.g., outer join agent1+agent5, then inner join with agent7)

EXAMPLES:

Example 1 - Single Database (ONE agent):
Question: "Which customers have tanks with warranty expiry dates within the next 30 days?"
Database agent5 has: INV_ITEM_SERIAL_DETAILS_V with warranty_exp_date
Database agent1 has: askme_om (orders - NO warranty data)
Response:
{{
    "queries": {{
        "agent5": "SELECT DISTINCT CUSTOMER_ID, CUSTOMER_NAME, SERIAL_NUMBER, WARRANTY_EXP_DATE FROM INV_ITEM_SERIAL_DETAILS_V WHERE WARRANTY_EXP_DATE BETWEEN CURDATE() AND DATE_ADD(CURDATE(), INTERVAL 30 DAY)",
        "agent1": ""
    }},
    "merge_key": null,
    "merge_strategy": "none",
    "reasoning": "Only agent5 contains warranty data. Agent1 has no relevant tables."
}}

Example 2 - Independent Databases (MULTIPLE agents, NO common join):
Question: "Show me customers with expiring warranties AND total sales for 2025"
Database agent5 has: INV_ITEM_SERIAL_DETAILS_V (warranty data)
Database agent1 has: askme_om (sales data)
No common join key between them
Response:
{{
    "queries": {{
        "agent5": "SELECT DISTINCT CUSTOMER_ID, CUSTOMER_NAME, SERIAL_NUMBER, WARRANTY_EXP_DATE FROM INV_ITEM_SERIAL_DETAILS_V WHERE WARRANTY_EXP_DATE BETWEEN CURDATE() AND DATE_ADD(CURDATE(), INTERVAL 30 DAY)",
        "agent1": "SELECT CUSTOMER_NAME, SUM(ORDERVALUE) AS TOTAL_SALES_2025 FROM askme_om WHERE EXTRACT(YEAR FROM ORDERDATE) = 2025 GROUP BY CUSTOMER_NAME"
    }},
    "merge_key": null,
    "merge_strategy": "independent",
    "reasoning": "Both databases needed but no clear natural join key between warranty data and sales data. Results should be shown separately."
}}

Example 3 - Concatenate Strategy (MULTIPLE agents, SIMILAR schemas):
Question: "Show all customer transactions from all systems"
Database agent1 has: transactions table (transaction_id, customer_name, amount, date, source)
Database agent2 has: legacy_transactions table (transaction_id, customer_name, amount, date, source)
Database agent3 has: partner_transactions table (transaction_id, customer_name, amount, date, source)
All have similar column structure
Response:
{{
    "queries": {{
        "agent1": "SELECT transaction_id, customer_name, amount, date, 'System1' as source FROM transactions WHERE date >= '2025-01-01'",
        "agent2": "SELECT transaction_id, customer_name, amount, date, 'Legacy' as source FROM legacy_transactions WHERE date >= '2025-01-01'",
        "agent3": "SELECT transaction_id, customer_name, amount, date, 'Partner' as source FROM partner_transactions WHERE date >= '2025-01-01'"
    }},
    "merge_key": null,
    "merge_strategy": "concatenate",
    "reasoning": "All three databases have transaction data with similar schemas. Stack them vertically to create one unified transaction list."
}}

Example 4 - Simple Merge (MULTIPLE agents, ONE strategy):
Question: "Show customer names with their order totals and inventory status"
Database agent1 has: customers table with customer_id, customer_name
Database agent2 has: orders table with customer_id, order_total
Database agent3 has: inventory table with customer_id, stock_level
All share customer_id as common key
Response:
{{
    "queries": {{
        "agent1": "SELECT customer_id, customer_name, email FROM customers",
        "agent2": "SELECT customer_id, SUM(order_total) as total_orders FROM orders GROUP BY customer_id",
        "agent3": "SELECT customer_id, SUM(stock_level) as total_inventory FROM inventory GROUP BY customer_id"
    }},
    "merge_key": "customer_id",
    "merge_strategy": "outer",
    "reasoning": "All three databases share customer_id. Using outer join to keep all customers even if they don't have orders or inventory."
}}

Example 5 - Advanced Multi-Step Merge (Different strategies):
Question: "Show customer purchase history with product details and only active loyalty members"
Database agent1 has: customers (customer_id, name) - all customers
Database agent2 has: purchases (customer_id, product_id, amount)
Database agent3 has: products (product_id, product_name, category)
Database agent4 has: loyalty_members (customer_id, tier, points) - subset of customers
Strategy: First outer join customers+purchases (keep all customers), then join with products on product_id, finally inner join with loyalty (only active members)
Response:
{{
    "queries": {{
        "agent1": "SELECT customer_id, name, email FROM customers",
        "agent2": "SELECT customer_id, product_id, SUM(amount) as total_spent FROM purchases GROUP BY customer_id, product_id",
        "agent3": "SELECT product_id, product_name, category FROM products",
        "agent4": "SELECT customer_id, tier, points FROM loyalty_members WHERE status = 'active'"
    }},
    "execution_order": ["agent1", "agent2", "agent3", "agent4"],
    "merge_steps": [
        {{
            "agents": ["agent1", "agent2"],
            "merge_key": "customer_id",
            "merge_strategy": "outer"
        }},
        {{
            "agents": ["result_1_2", "agent3"],
            "merge_key": "product_id",
            "merge_strategy": "left"
        }},
        {{
            "agents": ["result_1_2_3", "agent4"],
            "merge_key": "customer_id",
            "merge_strategy": "inner"
        }}
    ],
    "reasoning": "Multi-step merge: (1) outer join customers+purchases to keep all customers, (2) left join with products to get product details, (3) inner join with loyalty to filter only active members."
}}

Example 6 - Advanced Merge with Different Keys:
Question: "Show regional sales performance merged with headquarters targets"
Database agent1 has: regional_sales (region_id, region_name, sales_amount)
Database agent2 has: regional_inventory (region_id, stock_level)
Database agent3 has: targets (target_id, region_name, target_amount) - uses region_name not region_id
Strategy: First join agent1+agent2 on region_id, then join result with agent3 on region_name
Response:
{{
    "queries": {{
        "agent1": "SELECT region_id, region_name, SUM(sales_amount) as total_sales FROM regional_sales GROUP BY region_id, region_name",
        "agent2": "SELECT region_id, SUM(stock_level) as total_stock FROM regional_inventory GROUP BY region_id",
        "agent3": "SELECT region_name, target_amount FROM targets WHERE year = 2025"
    }},
    "execution_order": ["agent1", "agent2", "agent3"],
    "merge_steps": [
        {{
            "agents": ["agent1", "agent2"],
            "merge_key": "region_id",
            "merge_strategy": "inner"
        }},
        {{
            "agents": ["result_1_2", "agent3"],
            "merge_key": "region_name",
            "merge_strategy": "left"
        }}
    ],
    "reasoning": "Two-step merge with different keys: (1) inner join sales+inventory on region_id, (2) left join result with targets on region_name (different key)."
}}

KEY RULES TO REMEMBER:
1. Query ONLY databases with relevant data - set others to empty string ""
2. Use "none" when only 1 database is relevant
3. Use "independent" when multiple databases are needed but have NO common join key AND different purposes
4. Use "concatenate" when multiple databases have SIMILAR schemas and should be stacked vertically (like UNION)
5. Use SIMPLE merge (single merge_key/merge_strategy) when all agents use same strategy
6. Use ADVANCED merge (execution_order + merge_steps) when:
   - Different merge keys needed between different agent pairs
   - Different merge strategies needed (e.g., outer then inner)
   - Sequential dependencies (result of step 1 used in step 2)
7. In merge_steps, intermediate results are referenced as "result_X_Y" (agents X and Y merged)
8. Always use ORIGINAL table names (not the agent-prefixed versions shown in training)
9. For "concatenate", ensure all queries return SAME column names and types

NOW ANALYZE THE USER'S QUESTION AND GENERATE APPROPRIATE QUERIES:
- First identify which databases have relevant tables
- If only one database is relevant, query only that one (strategy: "none")
- If multiple have similar data to stack, use "concatenate" strategy (vertical stacking)
- If multiple are relevant but unrelated, use "independent" strategy (returns separate results)
- If multiple are relevant with ONE common strategy, use SIMPLE merge format
- If multiple are relevant with DIFFERENT strategies/keys per step, use ADVANCED merge format with execution_order and merge_steps
"""
        return prompt.strip()

    def _extract_json_from_response(self, response: str) -> Optional[dict]:
        """Extract a JSON object from LLM response (handles ```json ... ``` or raw {...})."""
        if not response or not response.strip():
            return None
        text = response.strip()
        # Try ```json ... ``` block first
        match = re.search(r"```(?:json)?\s*(\{[\s\S]*?\})\s*```", text)
        if match:
            try:
                return json.loads(match.group(1))
            except json.JSONDecodeError:
                pass
        # Try first {...} in text
        start = text.find("{")
        if start >= 0:
            depth = 0
            for i in range(start, len(text)):
                if text[i] == "{":
                    depth += 1
                elif text[i] == "}":
                    depth -= 1
                    if depth == 0:
                        try:
                            return json.loads(text[start : i + 1])
                        except json.JSONDecodeError:
                            break
        return None

    def get_merge_strategy_from_result_shapes(
        self,
        question: str,
        queries: Dict[str, str],
        result_shapes: Dict[str, dict],
    ) -> dict:
        """
        Ask LLM to decide how to merge multi-agent result DataFrames based on question,
        per-agent SQL, and result shapes (columns, row_count, dtypes) only — no row data.
        Returns merge_info compatible with run_sql_cached (merge_key, merge_strategy,
        execution_order, merge_steps). On parse failure returns independent (no merge).
        """
        if not queries or not result_shapes:
            return {"queries": queries or {}, "merge_strategy": "independent"}

        # Build shape summary for prompt (no row data); include all columns so LLM can choose merge key
        agent_lines = []
        for agent_id, shape in result_shapes.items():
            cols = shape.get("columns", [])
            row_count = shape.get("row_count", 0)
            dtypes = shape.get("dtypes", {})
            sql = (queries.get(agent_id) or "").strip()
            dtype_str = ", ".join(f"{c}({dtypes.get(c, '?')})" for c in cols)
            agent_lines.append(
                f"Agent {agent_id}: rows={row_count}, columns=[{dtype_str}]\n  SQL: {sql}"
            )
        shapes_block = "\n".join(agent_lines)

        prompt_text = f"""You are a data merge expert. Given a user question and the SQL + result shape (columns, row count, dtypes) for each agent — NO actual row data — decide how to merge the results.

USER QUESTION: {question}

PER-AGENT SQL AND RESULT SHAPES:
{shapes_block}

MERGE STRATEGY OPTIONS:
- "none": Only one agent has data; others empty — not applicable here since we have result shapes for multiple.
- "independent": Keep results separate (return dict of DataFrames). Use when no common join key or different purposes.
- "concatenate": Stack rows vertically when schemas are similar (like UNION ALL). Use when columns/meaning align.
- "inner": Merge horizontally on a common key; keep only rows with matches in all.
- "outer": Merge horizontally; keep all rows from all agents.
- "left" / "right": Merge keeping all rows from first/last agent.

OUTPUT: Respond with a single JSON object only. Use the same structure as below. Preserve agent keys in "queries" (you can pass through: {list(queries.keys())}).

SIMPLE MERGE (one key, one strategy):
{{
  "queries": {{ "agent1": "...", "agent2": "..." }},
  "merge_key": "common_column_name_or_null",
  "merge_strategy": "none|independent|concatenate|inner|outer|left|right",
  "reasoning": "Brief explanation"
}}

ADVANCED MERGE (different keys/strategies per step):
{{
  "queries": {{ "agent1": "...", "agent2": "..." }},
  "execution_order": ["agent1", "agent2"],
  "merge_steps": [
    {{ "agents": ["agent1", "agent2"], "merge_key": "COL", "merge_strategy": "outer" }}
  ],
  "reasoning": "Explanation"
}}

RULES:
1. Choose a merge_key only if that column exists in the result shapes of the agents you are merging.
2. Use "concatenate" when columns are similar and stacking rows answers the question.
3. Use "independent" when there is no sensible merge or no common key.
4. Return ONLY valid JSON, no markdown or extra text."""

        try:
            message_log = [
                self.system_message("You are a data merge expert. Return only a single JSON object."),
                self.user_message(prompt_text),
            ]
            response = self.submit_prompt(message_log, **{})
            parsed = self._extract_json_from_response(response)
            if not parsed or not isinstance(parsed, dict):
                return {"queries": queries, "merge_strategy": "independent"}

            # Build merge_info with pass-through queries; ensure required keys
            merge_info = {
                "queries": parsed.get("queries", queries),
                "merge_key": parsed.get("merge_key"),
                "merge_strategy": parsed.get("merge_strategy", "independent"),
                "execution_order": parsed.get("execution_order", list(queries.keys())),
                "merge_steps": parsed.get("merge_steps") or [],
            }
            if merge_info["merge_strategy"] == "none":
                merge_info["merge_strategy"] = "independent"
            return merge_info
        except Exception as e:
            self.log(f"get_merge_strategy_from_result_shapes failed: {e}", title="MergeStrategy")
            return {"queries": queries, "merge_strategy": "independent"}

    @abstractmethod
    def submit_prompt(self, prompt, **kwargs) -> str:
        """
        Example:
        ```python
        self.submit_prompt(
            [
                self.system_message("The user will give you SQL and you will try to guess what the business question this query is answering. Return just the question without any additional explanation. Do not reference the table name in the question."),
                self.user_message("What are the top 10 customers by sales?"),
            ]
        )
        ```

        This method is used to submit a prompt to the LLM.

        Args:
            prompt (any): The prompt to submit to the LLM.

        Returns:
            str: The response from the LLM.
        """
        pass

    def generate_question(self, sql: str, **kwargs) -> str:
        response = self.submit_prompt(
            [
                self.system_message(
                    "The user will give you SQL and you will try to guess what the business question this query is answering. Return just the question without any additional explanation. Do not reference the table name in the question."
                ),
                self.user_message(sql),
            ],
            **kwargs,
        )

        return response
    
    def generate_init_question(self, sql: str,count :int, **kwargs) -> str:
        response = self.submit_prompt(
            [
                self.system_message(
                     f"The user will give you SQL and you will try to guess what the business questions this query is answering. Return just the list of {count} questions without any additional explanation. Do not reference the table name in the question."
                ),
                self.user_message(sql),
            ],
            **kwargs,
        )
        
        return response
    
    def generate_normal_question(self, query: str, **kwargs) -> str:
        response = self.submit_prompt(
            [
                self.user_message(query),
            ],
            **kwargs,
        )
        return response

    def rephrase_text_for_agent_classification(self, text: str, **kwargs) -> str:
        """
        Rewrite a user's draft **agent description** (typically for a new agent with no
        initial training data) into crisp text so an LLM can route end-user questions
        to the right agent among many: high signal, low ambiguity, clear scope boundaries.

        Returns plain text only — no markdown, labels, or quotes wrapping the answer.
        """
        t = (text or "").strip()
        if not t:
            return ""
        response = self.submit_prompt(
            [
                self.system_message(
                    "You rewrite draft agent descriptions for a multi-agent system. There is no "
                    "training corpus yet—only this text—so the description must stand alone for routing. "
                    "Goals: (1) Let an LLM reliably tell this agent apart from others. (2) Remove "
                    "ambiguity: replace vague phrases with precise scope (domain, data sources, systems, "
                    "org units, geography, time horizons, user audience) only when those facts appear in "
                    "the draft; do not invent capabilities, tools, or data the user did not imply. "
                    "(3) Strip filler, marketing tone, and duplicate ideas. (4) Where the draft is fuzzy, "
                    "state the uncertainty briefly in neutral wording (e.g. 'scope unclear: …') rather "
                    "than guessing. "
                    "Output exactly one tight paragraph (or two short sentences max): the refined "
                    "agent description only. No bullets, no markdown, no preamble or labels."
                ),
                self.user_message("Draft agent description:\n" + t),
            ],
            temperature=0.25,
            max_tokens=640,
            **kwargs,
        )
        out = (response or "").strip()
        if (out.startswith('"') and out.endswith('"')) or (out.startswith("'") and out.endswith("'")):
            if len(out) >= 2:
                out = out[1:-1].strip()
        return out

    def select_datasource_from_intent(
        self, question: str, available_types: List[str], **kwargs
    ) -> Dict[str, Any]:
        """
        LLM routing: pick DATABASE, API, or FILE from available_types for this question.
        Returns dict with keys: datasource (str), reasoning (str), confidence (float).
        On parse failure, returns primary (first of available_types) with low confidence.
        """
        import json as _json
        import re

        if not available_types:
            return {"datasource": "DATABASE", "reasoning": "no options", "confidence": 0.0}
        if len(available_types) == 1:
            return {
                "datasource": available_types[0],
                "reasoning": "single configured datasource",
                "confidence": 1.0,
            }

        desc = {
            "DATABASE": "Structured relational data; SQL queries, tables, aggregates.",
            "API": "External HTTP/SAP APIs, real-time operational data.",
            "FILE": "Documents, spreadsheets, uploads, unstructured text in files.",
        }
        lines = [f"- {t}: {desc.get(t, t)}" for t in available_types]
        user_block = (
            f'Question: "{question}"\n\n'
            f"Agent datasources:\n" + "\n".join(lines) + "\n\n"
            "Keyword hints: SQL/database/table → DATABASE; API/external/real-time → API; "
            "CSV/Excel/file/upload/document → FILE.\n"
            'Return JSON only: {"datasource":"DATABASE|API|FILE","reasoning":"...","confidence":0.0-1.0}'
        )
        try:
            raw = self.submit_prompt(
                [
                    self.system_message(
                        "You route user questions to the best datasource. "
                        "Respond with a single JSON object only, no markdown."
                    ),
                    self.user_message(user_block),
                ],
                temperature=0.2,
                max_tokens=400,
                **kwargs,
            )
        except Exception as e:
            return {
                "datasource": available_types[0],
                "reasoning": f"llm_error:{e}",
                "confidence": 0.3,
            }

        text = (raw or "").strip()
        m = re.search(r"\{[\s\S]*\}", text)
        if m:
            text = m.group(0)
        try:
            data = _json.loads(text)
        except Exception:
            return {
                "datasource": available_types[0],
                "reasoning": "json_parse_failed",
                "confidence": 0.3,
            }

        choice = str(data.get("datasource", "")).strip().upper()
        if choice == "DB":
            choice = "DATABASE"
        if choice not in available_types:
            return {
                "datasource": available_types[0],
                "reasoning": str(data.get("reasoning", "invalid_choice")),
                "confidence": float(data.get("confidence", 0.3) or 0.3),
            }
        try:
            conf = float(data.get("confidence", 0.8))
        except (TypeError, ValueError):
            conf = 0.8
        return {
            "datasource": choice,
            "reasoning": str(data.get("reasoning", "")),
            "confidence": max(0.0, min(1.0, conf)),
        }
    
    def generate_compare_data(self, query: str,df:str, web_data:str, **kwargs) -> str:
        kwargs["llm_scoped_feature"] = "web"
        response = self.submit_prompt(
            [
                self.system_message(
                    f"The user has some data from a database in table format and get some information from the internet. "
                    f"Answer the user's question using both data sources. "
                    f"Do NOT include any sentiment analysis in your response. "
                    f"Do NOT include any follow-up suggestions, offers, or phrases like 'let me know', 'if you need', or 'feel free to ask'. "
                    f"Just provide the factual answer and analysis only."
                ),
                self.user_message(f"User Question: {query}, data from the database: {df} data from the web: {web_data}")
            ],
            **kwargs,
        )
        return response

    def generate_contextual_question(self, previous_questions: list, current_question: str, **kwargs) -> str:
        """
        Combine previous questions with current question to generate contextualized query
        
        Args:
            previous_questions: List of up to 5 previous user questions
            current_question: The new question from user
        
        Returns:
            str: Rewritten contextual question
        """
        if not previous_questions:
            return current_question
        
        context = "\n".join([f"{i+1}. {q}" for i, q in enumerate(previous_questions)])
        
        prompt = [
            self.system_message(
                "You are helping to understand questions in context. "
                "Given previous questions and a new question, rewrite the new question to be self-contained "
                "by incorporating relevant context from previous questions. "
                "Return ONLY the rewritten question, no explanations. "
                "Important: if the latest question explicitly specifies a source/system/channel "
                "(for example: 'from order management', 'from SAP', 'from API', 'from database'), "
                "then treat it as an override and replace conflicting source/system from earlier context. "
                "Do not merge conflicting sources into one question."
            ),
            self.user_message(
                f"Previous questions:\n{context}\n\n"
                f"New question: {current_question}\n\n"
                f"Rewrite the new question with context:"
            )
        ]
        
        return self.submit_prompt(prompt, **kwargs)

    def _add_limit_to_sql(self, sql: str, limit: int) -> str:
        """
        Add appropriate LIMIT clause based on detected SQL dialect
        
        Args:
            sql: Original SQL query
            limit: Number of rows to limit
        
        Returns:
            str: SQL with appropriate limit syntax
        """
        sql = sql.strip().rstrip(';')
        sql_upper = sql.upper()
        
        # Check if limit already exists
        if any(keyword in sql_upper for keyword in ['LIMIT', 'TOP', 'FETCH', 'ROWNUM']):
            return sql
        
        # Detect dialect from self.dialect or infer from SQL
        dialect = self.dialect.upper() if hasattr(self, 'dialect') else 'MYSQL'
        
        # MySQL, PostgreSQL, SQLite, BigQuery, Snowflake
        if any(db in dialect for db in ['MYSQL', 'POSTGRES', 'SQLITE', 'BIGQUERY', 'SNOWFLAKE', 'MARIADB']):
            return f"{sql} LIMIT {limit};"
        
        # SQL Server / T-SQL
        elif any(db in dialect for db in ['SQLSERVER', 'MSSQL', 'T-SQL', 'TSQL']):
            # Check if ORDER BY exists
            if 'ORDER BY' in sql_upper:
                return f"{sql} OFFSET 0 ROWS FETCH NEXT {limit} ROWS ONLY;"
            else:
                # Add TOP after SELECT
                return sql.replace('SELECT', f'SELECT TOP {limit}', 1) + ';'
        
        # Oracle
        elif 'ORACLE' in dialect:
            if 'ORDER BY' in sql_upper:
                return f"{sql} FETCH FIRST {limit} ROWS ONLY;"
            else:
                # Use ROWNUM for older Oracle versions
                return f"SELECT * FROM ({sql}) WHERE ROWNUM <= {limit};"
        
        # DB2
        elif 'DB2' in dialect:
            return f"{sql} FETCH FIRST {limit} ROWS ONLY;"
        
        # Teradata
        elif 'TERADATA' in dialect:
            return f"{sql} SAMPLE {limit};"
        
        # Default fallback (MySQL syntax - most common)
        else:
            return f"{sql} LIMIT {limit};"


    def _remove_limit_from_sql(self, sql: str) -> str:
        """
        Remove LIMIT clause from SQL to restore original query structure
        
        Args:
            sql: SQL with limit
        
        Returns:
            str: SQL without limit clause
        """
        import re
        
        sql = sql.strip().rstrip(';')
        
        # Remove MySQL/PostgreSQL style LIMIT
        sql = re.sub(r'\s+LIMIT\s+\d+\s*$', '', sql, flags=re.IGNORECASE)
        
        # Remove SQL Server style TOP
        sql = re.sub(r'SELECT\s+TOP\s+\d+\s+', 'SELECT ', sql, flags=re.IGNORECASE)
        
        # Remove SQL Server style OFFSET/FETCH
        sql = re.sub(r'\s+OFFSET\s+\d+\s+ROWS\s+FETCH\s+NEXT\s+\d+\s+ROWS\s+ONLY\s*$', '', sql, flags=re.IGNORECASE)
        
        # Remove Oracle style FETCH FIRST
        sql = re.sub(r'\s+FETCH\s+FIRST\s+\d+\s+ROWS\s+ONLY\s*$', '', sql, flags=re.IGNORECASE)
        
        # Remove Oracle ROWNUM wrapper
        match = re.match(r'SELECT\s+\*\s+FROM\s+\((.*)\)\s+WHERE\s+ROWNUM\s*<=\s*\d+', sql, re.IGNORECASE | re.DOTALL)
        if match:
            sql = match.group(1)
        
        # Remove Teradata SAMPLE
        sql = re.sub(r'\s+SAMPLE\s+\d+\s*$', '', sql, flags=re.IGNORECASE)
        
        return sql.strip()
    
    def generate_and_test_followup_question(self, question: str, sql: str, df: pd.DataFrame, 
                                       already_generated: list = None, **kwargs) -> dict:
        """
        Generate ONE follow-up question and test it with LIMIT 1
        
        Args:
            question: Original question
            sql: Original SQL
            df: Original dataframe results
            already_generated: List of already generated questions to avoid duplicates
        
        Returns:
            dict: {'question': str, 'sql': str, 'success': bool}
        """
        try:
            # Prepare exclusion text for prompt
            exclusion_text = ""
            if already_generated and len(already_generated) > 0:
                exclusion_text = f"\n\nDo NOT generate any of these questions:\n" + "\n".join([f"- {q}" for q in already_generated])
            
            # Generate a single follow-up question
            prompt = [
                self.system_message(
                    f"The user asked: '{question}'\n\n"
                    f"SQL query: {sql}\n\n"
                    f"Result columns: {df.columns.tolist()}\n\n"
                    "Generate ONE follow-up question that the user might ask. "
                    "The question should be answerable with a SQL query based on available data. "
                    "Make it DIFFERENT from any previously generated questions. "
                    f"Return ONLY the question, no numbering or explanation.{exclusion_text}"
                ),
                self.user_message("Generate one unique follow-up question:")
            ]
            
            followup_q = self.submit_prompt(prompt, **kwargs).strip()
            
            # Remove any numbering that might appear
            followup_q = re.sub(r'^\d+\.\s*', '', followup_q).strip()
            
            # Check if this question is too similar to already generated ones
            if already_generated:
                for existing_q in already_generated:
                    # Simple similarity check - if 70% of words match, consider it duplicate
                    existing_words = set(existing_q.lower().split())
                    new_words = set(followup_q.lower().split())
                    if len(existing_words) > 0:
                        similarity = len(existing_words & new_words) / len(existing_words | new_words)
                        if similarity > 0.7:
                            return {'question': followup_q, 'sql': None, 'success': False, 'duplicate': True}
            
            # Generate SQL for this question
            followup_sql = self.generate_sql(question=followup_q, skip_node_pipeline=True)
            
            if not followup_sql:
                return {'question': followup_q, 'sql': None, 'success': False}
            
            # Add LIMIT 1 using dialect-aware function (will replace existing limits)
            test_sql = self._add_limit_to_sql(followup_sql, limit=1)
            
            # Validate SQL
            if not self.is_sql_valid(sql=test_sql):
                return {'question': followup_q, 'sql': test_sql, 'success': False}
            
            # Test run with LIMIT 1
            test_df = self.run_sql(sql=test_sql)
            
            if test_df is not None and len(test_df) >= 0:
                # Remove test limit and apply production limit
                base_sql = self._remove_limit_from_sql(test_sql)
                final_sql = self._add_limit_to_sql(base_sql, limit=int(self.max_rows))
                return {'question': followup_q, 'sql': final_sql, 'success': True}
            else:
                return {'question': followup_q, 'sql': test_sql, 'success': False}
                
        except Exception as e:
            print(f"Error generating followup: {e}")
            return {'question': '', 'sql': None, 'success': False}
    

    def generate_error_recovery_questions(self, failed_question: str, max_questions: int = 3, **kwargs) -> list:
        validated_questions = []
        already_generated = []
        attempts = 0
        max_attempts = max_questions * 1  # First phase: context-based attempts
        
        # Phase 1: Try to generate questions based on user's failed question context
        while len(validated_questions) < max_questions and attempts < max_attempts:
            attempts += 1
            
            exclusion_text = ""
            if already_generated:
                exclusion_text = f"\n\nDo NOT generate any of these questions:\n" + "\n".join([f"- {q}" for q in already_generated])
            
            prompt = [
                self.system_message(
                    f"The user asked: '{failed_question}' but it could not be answered.\n\n"
                    "Generate ONE similar but simpler question that CAN be answered with the available data. "
                    "The question should be related to the user's intent but more straightforward. "
                    f"Return ONLY the question, no numbering or explanation.{exclusion_text}"
                ),
                self.user_message("Generate one alternative question:")
            ]
            
            try:
                alt_question = self.submit_prompt(prompt, **kwargs).strip()
                alt_question = re.sub(r'^\d+\.\s*', '', alt_question).strip()
                
                if alt_question in already_generated:
                    continue
                
                already_generated.append(alt_question)
                
                alt_sql = self.generate_sql(question=alt_question, skip_node_pipeline=True)
                
                if not alt_sql or not self.is_sql_valid(sql=alt_sql):
                    continue
                
                test_sql = self._add_limit_to_sql(alt_sql, limit=1)
                test_df = self.run_sql(sql=test_sql)
                
                if test_df is not None and len(test_df) >= 0:
                    base_sql = self._remove_limit_from_sql(test_sql)
                    final_sql = self._add_limit_to_sql(base_sql, limit=int(self.max_rows))
                    validated_questions.append({
                        'question': alt_question,
                        'sql': final_sql
                    })
                    print(f"✓ Context recovery {len(validated_questions)}/{max_questions}: {alt_question}")
            except Exception as e:
                print(f"Context recovery attempt {attempts} failed: {e}")
                continue
        
        # Phase 2: If not enough questions found, generate from training data
        if len(validated_questions) < max_questions:
            print(f"Context-based recovery found only {len(validated_questions)} questions. Trying training data...")
            validated_questions = self._generate_from_training_data(
                failed_question=failed_question,
                existing_questions=validated_questions,
                already_generated=already_generated,
                max_questions=max_questions,
                **kwargs
            )
        
        return validated_questions


    def _generate_training_data_for_suggestion(self, initial_training: str, existing_questions: list, 
                                    already_generated: list, max_questions: int = 3, **kwargs) -> list:

        """
        Generate validated questions from training data when context-based generation fails
        """
        validated_questions = existing_questions.copy()
        attempts = 0
        max_attempts = (max_questions - len(validated_questions)) * 5
        training_context = initial_training
        while len(validated_questions) < max_questions and attempts < max_attempts:
            attempts += 1
            exclusion_text = ""
            all_generated = already_generated + [q['question'] for q in validated_questions]
            if all_generated:
                exclusion_text = f"\n\nDo NOT generate any of these questions:\n" + "\n".join([f"- {q}" for q in all_generated])
            prompt = [
                self.system_message(
                    f"{training_context}\n\n"
                    "Based on the available data schema above, generate ONE simple business question "
                    "that can definitely be answered with SQL. "
                    "The question should be practical and relevant to the data available. "
                    f"Return ONLY the question, no numbering or explanation.{exclusion_text}"
                ),
                self.user_message("Generate one answerable business question based on available data:")
            ]
            try:
                alt_question = self.submit_prompt(prompt, **kwargs).strip()
                alt_question = re.sub(r'^\d+\.\s*', '', alt_question).strip()
                if alt_question in all_generated:
                    continue
                already_generated.append(alt_question)
                alt_sql = self.generate_sql(question=alt_question, skip_node_pipeline=True)
                if not alt_sql or not self.is_sql_valid(sql=alt_sql):
                    continue
                test_sql = self._add_limit_to_sql(alt_sql, limit=1)
                test_df = self.run_sql(sql=test_sql)
                if test_df is not None and len(test_df) >= 0:
                    base_sql = self._remove_limit_from_sql(test_sql)
                    final_sql = self._add_limit_to_sql(base_sql, limit=int(self.max_rows))
                    validated_questions.append({
                        'question': alt_question,
                        'sql': final_sql
                    })
                    print(f"✓ Training data recovery {len(validated_questions)}/{max_questions}: {alt_question}")
            except Exception as e:
                print(f"Training data recovery attempt {attempts} failed: {e}")
                continue
        return validated_questions

    def _generate_from_training_data(self, failed_question: str, existing_questions: list, 
                                    already_generated: list, max_questions: int = 3, **kwargs) -> list:
        """
        Generate validated questions from training data when context-based generation fails
        """
        validated_questions = existing_questions.copy()
        attempts = 0
        max_attempts = (max_questions - len(validated_questions)) * 5
        
        # Get training data for context
        try:
            training_data = self.get_training_data()
            if training_data is None or (hasattr(training_data, 'empty') and training_data.empty):
                training_context = "No specific training data available."
            else:
                # Extract relevant info from training data
                if isinstance(training_data, pd.DataFrame):
                    # Get sample questions/DDL from training data
                    training_sample = training_data.head(20).to_string()
                    training_context = f"Available training data schema:\n{training_sample}"
                else:
                    training_context = str(training_data)[:2000]
        except Exception as e:
            print(f"Error getting training data: {e}")
            training_context = "Training data available but not accessible."
        
        while len(validated_questions) < max_questions and attempts < max_attempts:
            attempts += 1
            
            exclusion_text = ""
            all_generated = already_generated + [q['question'] for q in validated_questions]
            if all_generated:
                exclusion_text = f"\n\nDo NOT generate any of these questions:\n" + "\n".join([f"- {q}" for q in all_generated])
            
            prompt = [
                self.system_message(
                    f"The user asked: '{failed_question}' but we couldn't find related questions.\n\n"
                    f"{training_context}\n\n"
                    "Based on the available data schema above, generate ONE simple business question "
                    "that can definitely be answered with SQL. "
                    "The question should be practical and relevant to the data available. "
                    f"Return ONLY the question, no numbering or explanation.{exclusion_text}"
                ),
                self.user_message("Generate one answerable business question based on available data:")
            ]
            
            try:
                alt_question = self.submit_prompt(prompt, **kwargs).strip()
                alt_question = re.sub(r'^\d+\.\s*', '', alt_question).strip()
                
                if alt_question in all_generated:
                    continue
                
                already_generated.append(alt_question)
                
                alt_sql = self.generate_sql(question=alt_question, skip_node_pipeline=True)
                
                if not alt_sql or not self.is_sql_valid(sql=alt_sql):
                    continue
                
                test_sql = self._add_limit_to_sql(alt_sql, limit=1)
                test_df = self.run_sql(sql=test_sql)
                
                if test_df is not None and len(test_df) >= 0:
                    base_sql = self._remove_limit_from_sql(test_sql)
                    final_sql = self._add_limit_to_sql(base_sql, limit=int(self.max_rows))
                    validated_questions.append({
                        'question': alt_question,
                        'sql': final_sql
                    })
                    print(f"✓ Training data recovery {len(validated_questions)}/{max_questions}: {alt_question}")
            except Exception as e:
                print(f"Training data recovery attempt {attempts} failed: {e}")
                continue
        
        return validated_questions
    
    def generate_initial_question_with_run(self, initial_training, num_questions: int = 25, **kwargs) -> List[dict]:
        """
        Generate initial questions from training data with validated SQL
        
        Args:
            num_questions: Number of questions to generate
            **kwargs: Additional parameters for prompt submission
            
        Returns:
            List of dictionaries with 'question', 'sql', and metadata
        """
        print(f"Starting generation of {num_questions} questions from training data...")
        
        # Generate questions from training data only
        validated_questions = []
        already_generated = []
        
        try:
            validated_questions = self._generate_training_data_for_suggestion(
                initial_training=initial_training,  
                existing_questions=[],
                already_generated=already_generated,
                max_questions=num_questions,
                **kwargs
            )
            
            if validated_questions:
                print(f"✓ Successfully generated {len(validated_questions)} questions from training data")
                # Format questions to include metadata
                formatted_questions = []
                for q in validated_questions:
                    formatted_questions.append({
                        'question': q['question'],
                        'sql': q['sql'],
                        'generated_at': datetime.datetime.now().isoformat(),
                        'status': 'validated',
                        'source': 'training_data'
                    })
                return formatted_questions
            else:
                print("No questions could be generated from training data")
                return []
                
        except Exception as e:
            print(f"Training data generation failed: {e}")
            return []

    def generate_sql_for_question_with_validation(self, question: str, max_retries: int = 1, **kwargs) -> dict:
        """
        Generate SQL for a user-provided question using LLM, then validate by executing.
        Allows up to max_retries re-generation attempts if validation fails.
        
        Args:
            question: Natural language question to generate SQL for
            max_retries: Number of retry attempts if SQL validation fails (default: 1)
        
        Returns:
            dict: {
                'question': str,      # Original question
                'sql': str,           # Validated SQL or empty string
                'success': bool,      # Whether validation passed
                'error': str          # Error message if failed (optional)
            }
        """
        for attempt in range(1 + max_retries):  # 1 initial + max_retries
            try:
                print(f"  Attempt {attempt + 1}/{1 + max_retries} for: {question[:60]}...")
                
                # Step 1: Generate SQL via LLM
                generated_sql = self.generate_sql(question=question)
                
                if not generated_sql:
                    print(f"  ✗ No SQL generated")
                    continue
                
                # Step 2: Check SQL validity (syntax check)
                if not self.is_sql_valid(sql=generated_sql):
                    print(f"  ✗ SQL not valid (syntax)")
                    continue
                
                # Step 3: Test-run with LIMIT 1
                test_sql = self._add_limit_to_sql(generated_sql, limit=1)
                test_df = self.run_sql(sql=test_sql)
                
                if test_df is not None and len(test_df) >= 0:
                    # Success — build production SQL with max_rows limit
                    base_sql = self._remove_limit_from_sql(test_sql)
                    final_sql = self._add_limit_to_sql(base_sql, limit=int(self.max_rows))
                    
                    print(f"  ✓ SQL validated successfully")
                    return {
                        'question': question,
                        'sql': final_sql,
                        'success': True
                    }
                else:
                    print(f"  ✗ SQL execution returned no results")
                    continue
                    
            except Exception as e:
                print(f"  ✗ Attempt {attempt + 1} error: {e}")
                continue
        
        # All attempts exhausted
        return {
            'question': question,
            'sql': '',
            'success': False,
            'error': f'SQL generation/validation failed after {1 + max_retries} attempts'
        }

    def rephrase_question(self, question: str, **kwargs) -> str:
        """
        Use LLM to rephrase a question into a simpler, more database-friendly form.
        Called when the original question fails SQL generation/validation.
        
        Args:
            question: The original question that failed SQL generation
        
        Returns:
            str: A rephrased version of the question
        """
        try:
            prompt = [
                self.system_message(
                    "You are an expert at rephrasing natural language questions to make them "
                    "easier to convert into SQL queries.\n\n"
                    "The following question could not be converted into a valid SQL query. "
                    "Please rephrase it into a simpler, clearer form that:\n"
                    "1. Preserves the original intent and meaning\n"
                    "2. Uses simpler language and structure\n"
                    "3. Is more directly answerable with a SQL query\n"
                    "4. Avoids ambiguity\n\n"
                    "Return ONLY the rephrased question, nothing else. "
                    "Do not add any explanation, numbering, or prefix."
                ),
                self.user_message(
                    f"Rephrase this question for SQL generation:\n\n{question}"
                )
            ]
            
            rephrased = self.submit_prompt(prompt, **kwargs).strip()
            # Clean up any quotes or prefixes the LLM might add
            rephrased = rephrased.strip('"\'')
            import re
            rephrased = re.sub(r'^(Rephrased|Question|Q):\s*', '', rephrased, flags=re.IGNORECASE).strip()
            
            print(f"  🔄 Rephrased: '{question[:50]}...' → '{rephrased[:50]}...'")
            return rephrased
            
        except Exception as e:
            print(f"  ✗ Error rephrasing question: {e}")
            return question  # Return original if rephrasing fails

    def generate_sql_for_question_with_rephrase(self, question: str, max_retries: int = 0, **kwargs) -> dict:
        """
        Generate SQL for a user-provided question with rephrase fallback.
        
        Flow:
        1. Try generate_sql + validate with max_retries on original question
        2. If fails → LLM rephrase the question
        3. Try generate_sql + validate with max_retries on rephrased question
        4. If still fails → return failure
        
        Args:
            question: Natural language question to generate SQL for
            max_retries: Number of retry attempts per phase (default: 1)
        
        Returns:
            dict: {
                'question': str,           # Original question
                'rephrased_question': str,  # Rephrased question (if used), else None
                'sql': str,                # Validated SQL or empty string
                'success': bool,           # Whether validation passed
                'error': str               # Error message if failed (optional)
            }
        """
        # Phase 1: Try with original question
        print(f"\n📝 Phase 1: Trying original question...")
        result = self.generate_sql_for_question_with_validation(
            question=question, max_retries=max_retries, **kwargs
        )
        
        if result.get('success'):
            result['rephrased_question'] = None
            return result
        
        # Phase 2: Rephrase and retry
        print(f"  ⚠ Original question failed. Attempting rephrase...")
        rephrased = self.rephrase_question(question=question, **kwargs)
        
        # Don't retry if rephrase returned the same question
        if rephrased.strip().lower() == question.strip().lower():
            print(f"  ✗ Rephrase returned same question, skipping retry")
            result['rephrased_question'] = None
            return result
        
        print(f"\n📝 Phase 2: Trying rephrased question...")
        rephrased_result = self.generate_sql_for_question_with_validation(
            question=rephrased, max_retries=max_retries, **kwargs
        )
        
        # Build final result — keep original question as primary
        if rephrased_result.get('success'):
            return {
                'question': question,
                'rephrased_question': rephrased,
                'sql': rephrased_result['sql'],
                'success': True
            }
        else:
            return {
                'question': question,
                'rephrased_question': rephrased,
                'sql': '',
                'success': False,
                'error': f'SQL generation failed for both original and rephrased question after {2 * (1 + max_retries)} total attempts'
            }
        
    def generate_validated_followup_questions(self, question: str, sql: str, df: pd.DataFrame, 
                                         max_questions: int = 3, **kwargs) -> list:
        """
        Generate up to 3 DIFFERENT validated follow-up questions
        
        Args:
            question: Original question
            sql: Original SQL
            df: Original dataframe
            max_questions: Maximum number of questions to generate (default 3)
        
        Returns:
            list: List of dicts with {'question': str, 'sql': str}
        """
        validated_questions = []
        already_generated_texts = []  # Track all attempted questions
        attempts = 0
        max_attempts = max_questions * 5  # Increased to 15 attempts for better variety
        
        while len(validated_questions) < max_questions and attempts < max_attempts:
            attempts += 1
            
            # Pass already generated questions to avoid duplicates
            result = self.generate_and_test_followup_question(
                question, 
                sql, 
                df, 
                already_generated=already_generated_texts,
                **kwargs
            )
            
            if result.get('duplicate'):
                print(f"⚠ Skipping duplicate question: {result['question']}")
                continue
            
            # Track this question regardless of success
            if result['question']:
                already_generated_texts.append(result['question'])
            
            if result['success'] and result['question']:
                # Double-check for exact duplicates in validated list
                if not any(vq['question'].lower() == result['question'].lower() for vq in validated_questions):
                    validated_questions.append({
                        'question': result['question'],
                        'sql': result['sql']
                    })
                    print(f"✓ Validated {len(validated_questions)}/{max_questions}: {result['question']}")
            else:
                print(f"✗ Failed validation (attempt {attempts}): {result['question'][:50]}...")
        
        return validated_questions

    def _distribute_questions_across_agents(self, max_questions: int, num_agents: int) -> List[int]:
        """
        Distribute max_questions across num_agents. Earlier agents get at least as many as later.
        E.g. max_questions=3, num_agents=2 -> [2, 1]; max_questions=5, num_agents=2 -> [3, 2].
        """
        if num_agents <= 0:
            return []
        base, remainder = divmod(max_questions, num_agents)
        return [base + (1 if i < remainder else 0) for i in range(num_agents)]

    def generate_multi_validated_followup_questions(
        self,
        question: str,
        sql_list: List[str],
        df_list: Union[List[pd.DataFrame], pd.DataFrame],
        max_questions: int = 3,
        **kwargs
    ) -> List[dict]:
        """
        Generate up to max_questions validated follow-up questions from multiple SQL/DF pairs
        (e.g. one per agent). Questions are distributed across agents: earlier agents get
        at least as many as later (e.g. 3 questions, 2 agents -> 2 from first, 1 from second).

        Args:
            question: Original question
            sql_list: List of SQL strings (one per agent/result)
            df_list: Either a list of DataFrames (one per SQL) or a single DataFrame (e.g. merged
                     result); if single, the same df is used for every SQL when generating follow-ups.
            max_questions: Maximum total questions to generate (default 3)

        Returns:
            list: List of dicts with {'question': str, 'sql': str}
        """
        # Normalize sql_list to list
        if not isinstance(sql_list, list):
            sql_list = [sql_list] if sql_list else []
        if len(sql_list) == 0:
            return []

        # Normalize df_list: accept either a single DataFrame (merged) or list of DataFrames
        if not isinstance(df_list, list):
            # Single DataFrame (merged case): use same df for every SQL
            if df_list is not None and hasattr(df_list, "columns"):
                df_list = [df_list] * len(sql_list)
            else:
                df_list = []
        elif len(df_list) == 1 and len(sql_list) > 1 and hasattr(df_list[0], "columns"):
            # List of one DataFrame (merged): replicate so each SQL has a df
            df_list = df_list * len(sql_list)

        if len(df_list) != len(sql_list):
            print(f"⚠ sql_list length ({len(sql_list)}) != df_list length ({len(df_list)}); using min length.")
        n = min(len(sql_list), len(df_list))
        if n == 0:
            return []

        sql_list = sql_list[:n]
        df_list = df_list[:n]

        quotas = self._distribute_questions_across_agents(max_questions, n)
        validated_questions = []
        already_generated_texts = []

        for agent_idx in range(n):
            need = quotas[agent_idx]
            if need <= 0:
                continue
            sql = sql_list[agent_idx]
            df = df_list[agent_idx]
            if df is None or not hasattr(df, "columns"):
                continue
            attempts = 0
            max_attempts = need * 5

            while len([vq for vq in validated_questions if vq.get("_agent_idx") == agent_idx]) < need and attempts < max_attempts:
                attempts += 1
                result = self.generate_and_test_followup_question(
                    question,
                    sql,
                    df,
                    already_generated=already_generated_texts,
                    **kwargs
                )
                if result.get("duplicate"):
                    print(f"⚠ Skipping duplicate question: {result['question']}")
                    continue
                if result.get("question"):
                    already_generated_texts.append(result["question"])
                if result.get("success") and result.get("question"):
                    agent_validated = [vq for vq in validated_questions if vq.get("_agent_idx") == agent_idx]
                    if not any(vq["question"].lower() == result["question"].lower() for vq in agent_validated):
                        validated_questions.append({
                            "question": result["question"],
                            "sql": result["sql"],
                            "_agent_idx": agent_idx,
                        })
                        print(f"✓ Validated (agent {agent_idx}) {len(validated_questions)}/{max_questions}: {result['question']}")
                else:
                    print(f"✗ Failed validation (agent {agent_idx}, attempt {attempts}): {result.get('question', '')[:50]}...")

        # Drop internal key for return
        return [{"question": vq["question"], "sql": vq["sql"]} for vq in validated_questions]

    @staticmethod
    def _strip_markdown_code_fences(text: str) -> str:
        """
        Remove ``` / ```python fences when the regex in _extract_python_code missed a pair
        (e.g. unclosed fence, or only a closing ``` on the last line) so exec() does not see
        backticks as the first line of Python.
        """
        t = (text or "").strip()
        if not t:
            return t
        if t.startswith("```"):
            first_nl = t.find("\n")
            if first_nl != -1:
                t = t[first_nl + 1 :]
            else:
                t = t[3:].lstrip()
        t = t.rstrip()
        if t.rstrip().endswith("```"):
            t = t.rstrip()[:-3].rstrip()
        return t.strip()

    def _extract_python_code(self, markdown_string: str) -> str:
        # Strip whitespace to avoid indentation errors in LLM-generated code
        markdown_string = markdown_string.strip()

        # Regex pattern to match Python code blocks
        pattern = r"```[\w\s]*python\n([\s\S]*?)```|```([\s\S]*?)```"

        # Find all matches in the markdown string
        matches = re.findall(pattern, markdown_string, re.IGNORECASE)

        # Extract the Python code from the matches
        python_code = []
        for match in matches:
            python = match[0] if match[0] else match[1]
            python_code.append(python.strip())

        if len(python_code) == 0:
            return self._strip_markdown_code_fences(markdown_string)

        # Phi/mini models often emit multiple ``` blocks or a truncated first block
        # (e.g. incomplete `for`), which then fails compile(). Prefer any fenced block
        # that is valid Python; among those, prefer the longest (usually the real chart).
        valid_blocks = [b for b in python_code if self._is_executable_python_chunk(b)]
        if valid_blocks:
            return self._strip_markdown_code_fences(max(valid_blocks, key=len))
        return self._strip_markdown_code_fences(max(python_code, key=len))

    @staticmethod
    def _is_executable_python_chunk(code: str) -> bool:
        if not (code or "").strip():
            return False
        try:
            compile(code, "<llm_python>", "exec")
            return True
        except SyntaxError:
            return False

    def _sanitize_plotly_code(self, raw_plotly_code: str) -> str:
        # Remove the fig.show() statement from the plotly code
        plotly_code = raw_plotly_code.replace("fig.show()", "")

        return plotly_code

    def generate_plotly_code(
        self, question: str = None, sql: str = None, df_metadata: str = None, **kwargs
    ) -> str:
        if question is not None:
            system_msg = f"The following is a pandas DataFrame that contains the results of the query that answers the question the user asked: '{question}'"
        else:
            system_msg = "The following is a pandas DataFrame "

        if sql is not None:
            system_msg += f"\n\nThe DataFrame was produced using this query: {sql}\n\n"

        system_msg += f"The following is information about the resulting pandas DataFrame 'df': \n{df_metadata}"

        message_log = [
            self.system_message(system_msg),
            self.user_message(
                "Can you generate the Python plotly code to chart the results of the dataframe? Assume the data is in a pandas dataframe called 'df'. If there is only one value in the dataframe, use an Indicator. Respond with only Python code. Do not answer with any explanations -- just the code."
            ),
        ]
        

        plotly_code = self.submit_prompt(message_log, **kwargs)


        return self._sanitize_plotly_code(self._extract_python_code(plotly_code))
    
        
    def generate_forecast(self, question: str, df: pd.DataFrame, **kwargs) -> pd.DataFrame:
        """
        Generate forecast dataframe based on user question and historical data.
        
        Args:
            question (str): The forecast question the user asked.
            df (pd.DataFrame): The historical dataframe to base the forecast on.
            
        Returns:
            pd.DataFrame: A new dataframe with forecasted data.
        """
        # Convert dataframe to markdown for the prompt
        df_markdown = df.to_markdown(index=False)
        df_info = f"DataFrame shape: {df.shape}\nColumns: {df.columns.tolist()}\n\nData:\n{df_markdown}"
        
        system_msg = f"The following is a pandas DataFrame with historical data:\n\n{df_info}\n\nThe user wants to create a forecast based on this data."
        
        message_log = [
            self.system_message(system_msg),
            self.user_message(
                f"Based on the historical data provided and the user's forecast request: '{question}', "
                "generate Python code that creates a new pandas DataFrame called 'forecast_df' with forecasted data. "
                "IMPORTANT: You can ONLY use pandas (pd) and numpy (np) - DO NOT use any external libraries like statsmodels, scikit-learn, prophet, etc.\n\n"
                "The code should:\n"
                "1. Analyze the historical data structure\n"
                "2. Generate forecasted values using simple methods like:\n"
                "   - Linear trend projection (np.polyfit, np.polyval)\n"
                "   - Moving averages\n"
                "   - Simple extrapolation\n"
                "   - Basic statistical methods (mean, median, etc.)\n"
                "3. Create a new DataFrame 'forecast_df' with the forecasted data\n"
                "4. Ensure the forecast_df has similar structure to the original df (same columns if applicable)\n\n"
                "If the data has a date/time column, use it for time-series forecasting with simple trend analysis. "
                "Use only pandas and numpy functions. "
                "Respond with ONLY Python code. Do not include any explanations -- just the code that creates 'forecast_df'."
                + self._response_language()
            ),
        ]
        
        forecast_code = self.submit_prompt(message_log, **kwargs)
        python_code = self._extract_python_code(forecast_code)
        
        # Pre-process code to remove problematic imports
        python_code = self._sanitize_forecast_code(python_code)
        
        # Execute the code to generate the forecast dataframe
        try:
            # Create a safe execution environment - only allow pandas and numpy
            import pandas as pd
            import numpy as np
            
            # Create a custom __import__ that blocks external libraries
            import builtins
            _original_import = builtins.__import__
            def safe_import(name, globals=None, locals=None, fromlist=(), level=0):
                allowed_modules = {'pandas', 'numpy', 'pd', 'np', 'datetime', 'date', 'time', 'plotly', 'plotly.express', 'px', 'sklearn', 'joblib', 'scipy', 'statsmodels', 'sm', 'matplotlib', 'plt'}
                module_name = name.split('.')[0]
                if module_name in allowed_modules:
                    return _original_import(name, globals, locals, fromlist, level)
                else:
                    raise ImportError(f"Import of '{name}' is not allowed. Allowed: pandas, numpy, plotly, datetime, sklearn, statsmodels, matplotlib.")
            
            # Create restricted builtins with custom __import__
            safe_builtins = {
                'abs': abs, 'all': all, 'any': any, 'bool': bool, 'dict': dict,
                'float': float, 'int': int, 'len': len, 'list': list, 'max': max,
                'min': min, 'range': range, 'round': round, 'sorted': sorted,
                'str': str, 'sum': sum, 'tuple': tuple, 'zip': zip, 'print': print,
                'enumerate': enumerate, '__import__': safe_import,
                '__builtins__': {
                    'abs': abs, 'all': all, 'any': any, 'bool': bool, 
                    'dict': dict, 'float': float, 'int': int, 'len': len,
                    'list': list, 'max': max, 'min': min, 'range': range,
                    'round': round, 'sorted': sorted, 'str': str, 'sum': sum,
                    'tuple': tuple, 'zip': zip, 'print': print, 'enumerate': enumerate,
                    '__import__': safe_import
                }
            }
            local_vars = {'df': df.copy(), 'pd': pd, 'np': np, 'LinearRegression': LinearRegression}
            exec(python_code, safe_builtins, local_vars)
            
            if 'forecast_df' in local_vars:
                return local_vars['forecast_df']
            else:
                # If forecast_df wasn't created, try to find any DataFrame in local_vars
                for key, value in local_vars.items():
                    if isinstance(value, pd.DataFrame) and key != 'df':
                        return value
                raise ValueError("Forecast code did not create a 'forecast_df' DataFrame")
        except Exception as e:
            self.log(f"Error executing forecast code: {e}\nCode: {python_code}", "Forecast Error")
            # Try simple fallback: extend the dataframe with basic trend
            try:
                import numpy as np
                forecast_df = self._simple_forecast_fallback(df, question)
                if forecast_df is not None:
                    return forecast_df
            except Exception as fallback_error:
                self.log(f"Fallback forecast also failed: {fallback_error}", "Forecast Error")
            # Return original dataframe if all else fails
            return df.copy()

    def _simple_forecast_fallback(self, df: pd.DataFrame, question: str) -> pd.DataFrame:
        """
        Simple fallback forecast method using basic pandas/numpy operations.
        Creates a simple trend extension of numeric columns.
        """
        try:
            import numpy as np
            forecast_df = df.copy()
            # Simple approach: duplicate the last few rows with slight adjustments
            # This is a very basic fallback - better than returning the original df
            if len(forecast_df) > 1:
                last_row = forecast_df.iloc[-1:].copy()
                # Try to extend by 5-10 periods
                num_periods = min(10, max(5, len(forecast_df) // 2))
                extended_rows = []
                for i in range(1, num_periods + 1):
                    new_row = last_row.copy()
                    # For numeric columns, apply a simple trend
                    for col in new_row.select_dtypes(include=[np.number]).columns:
                        if col in forecast_df.columns:
                            # Calculate simple trend from last few values
                            recent_values = forecast_df[col].tail(min(5, len(forecast_df))).values
                            if len(recent_values) > 1:
                                trend = (recent_values[-1] - recent_values[0]) / len(recent_values)
                                new_row[col] = new_row[col].values[0] + (trend * i)
                    extended_rows.append(new_row)
                if extended_rows:
                    forecast_df = pd.concat([forecast_df] + extended_rows, ignore_index=True)
            return forecast_df
        except Exception:
            return None

    def fetch_google_search_context_for_user_query(
        self,
        user_question: str,
        *,
        gemini_api_key: Optional[str] = None,
        gemini_model: str = "gemini-3-flash-preview",
    ) -> str:
        """
        Google Search via Gemini for Qwen follow-ups: **only the user question** is sent to
        Gemini (so DF/table rows stay on the local Ollama path). A separate local Qwen call
        still receives the full dataframe; this is optional web context from search only.
        """
        key = (gemini_api_key or "").strip() or _resolve_gemini_api_key(self.config) or ""
        if not key or not (user_question or "").strip():
            return ""

        system_instruction = (
            "You have access to Google Search. The user will provide a question. "
            "Search the web for current, relevant information and summarize in plain text. "
            "Be concise. Do not output code blocks. You do not have access to the user's private table data; "
            "base your answer on public web results and the question wording only."
        )
        user_text = f"User question:\n{user_question}\n\nSearch as needed and summarize what is relevant."
        try:
            out = _gemini_generate_content(
                user_text=user_text,
                api_key=key,
                system_instruction=system_instruction,
                model=gemini_model,
                google_search=True,
            )
            return (out or "").strip()
        except Exception as e:
            print(f"[GoogleSearchContext] {e}", flush=True)
            try:
                self.log(f"fetch_google_search_context_for_user_query: {e}", title="GoogleSearchContext")
            except Exception:
                pass
            return ""

    def fetch_google_search_context_for_prior_df(
        self,
        user_question: str,
        *args: Any,
        **kwargs: Any,
    ) -> str:
        """
        Legacy name: *args (e.g. old prior-df string) is ignored so table data is never sent to Gemini.
        """
        return self.fetch_google_search_context_for_user_query(
            user_question,
            gemini_api_key=kwargs.get("gemini_api_key"),
            gemini_model=kwargs.get("gemini_model", "gemini-3-flash-preview"),
        )

    def ask_data_analysis(self, question: str, df: pd.DataFrame, **kwargs) -> dict:
        """
        Generate flexible data analysis based on user question.
        Can return text summary, new dataframe, chart code, or combination.
        
        Args:
            question (str): The user's analysis request (e.g., "build a forecast", "do trend analysis", "detailed summary")
            df (pd.DataFrame): The historical dataframe to analyze
            
        Returns:
            dict: Dictionary with keys:
                - 'text': Optional text summary/analysis
                - 'dataframe': Optional new DataFrame (e.g., forecast data)
                - 'chart_code': Optional plotly code for chart generation
                - 'has_dataframe': bool indicating if dataframe should be shown
                - 'has_chart': bool indicating if chart should be shown
                - 'has_text': bool indicating if text should be shown
        """
        use_qwen = kwargs.get("use_qwen", False)
        from app.services.llm_service import _qwen_configured
        
        # Convert dataframe to markdown for the prompt (limit size for large dataframes)
        if not (use_qwen and _qwen_configured()):
            df_info = f"DataFrame shape: {df.shape}\nColumns: {df.columns.tolist()}"
        elif len(df) > 100:
            df_sample = df.head(100)
            df_info = f"DataFrame shape: {df.shape}\nColumns: {df.columns.tolist()}\n\nSample data (first 100 rows):\n{df_sample.to_markdown(index=False)}"
        else:
            df_info = f"DataFrame shape: {df.shape}\nColumns: {df.columns.tolist()}\n\nData:\n{df.to_markdown(index=False)}"
        
        system_msg = f"The following is a pandas DataFrame with data:\n\n{df_info}\n\nThe user wants to perform analysis on this data."

        
        user_msg = (
            f"Based on the data provided and the user's request: '{question}', "
            "determine what type of analysis to perform and generate the appropriate output.\n\n"
            "You can generate:\n"
            "1. TEXT ANALYSIS: A detailed text summary, analysis, or insights (use OUTPUT_TYPE: TEXT)\n"
            "2. NEW DATAFRAME: Transform or create new data (forecast, aggregation, filtering, etc.) (use OUTPUT_TYPE: DATAFRAME)\n"
            "3. CHART: Visualization code for the data (use OUTPUT_TYPE: CHART)\n"
            "4. COMBINATION: Any combination of the above (use OUTPUT_TYPE: COMBINATION)\n\n"
            "For TEXT analysis, provide a detailed analysis/summary in plain text. CRITICAL: DO NOT include meta-explanations, disclaimers, or references to your methodology (e.g., 'Based on the data...', 'generated using linear trend extrapolation', 'Note: These forecasts are...'). Provide ONLY the direct business insights.\n"
            "For DATAFRAME, generate Python code that creates 'result_df' using ONLY pandas (pd) and numpy (np). DO NOT use external libraries. IMPORTANT: Data may contain NaNs/nulls. Always handle missing values (e.g. .fillna(0) or .astype(float)) to avoid 'Cannot convert non-finite values' errors.\n"
            "For CHART, generate Python plotly code that creates 'chart_fig' using plotly express or graph_objects. Handle missing values gracefully.\n"
            "For COMBINATION, provide multiple outputs.\n\n"
            "IMPORTANT: Start your response with OUTPUT_TYPE: [TEXT|DATAFRAME|CHART|COMBINATION]\n\n"
            "Then provide the appropriate output(s):\n"
            "- If TEXT: Provide the analysis text directly with NO methodology explanations and NO disclaimers.\n"
            "- If DATAFRAME: Provide Python code that creates 'result_df'\n"
            "- If CHART: Provide Python plotly code that creates 'chart_fig'\n"
            "- If COMBINATION: Provide multiple sections labeled with TEXT:, DATAFRAME:, CHART:\n\n"
            "Examples:\n"
            "- 'build a forecast' -> OUTPUT_TYPE: DATAFRAME (generate forecast code)\n"
            "- 'do trend analysis' -> OUTPUT_TYPE: COMBINATION (text analysis + chart)\n"
            "- 'detailed summary' -> OUTPUT_TYPE: TEXT\n"
            "- 'show me a chart' -> OUTPUT_TYPE: CHART\n"
            + self._response_language()
        )

        from app.core.config import settings as _app_settings

        # If local Ollama/Qwen is configured and callers want it (@askme path), do not prefer Gemini
        # over the scoped local model — same priority as the user request for @askme.
        _skip_gemini_for_qwen = bool(kwargs.get("use_qwen")) and _app_settings.is_local_qwen_configured()

        gemini_used = False
        from app.services.llm_service import (
            call_oic_anthropic_user_prompt,
            oic_anthropic_integration_enabled,
            record_external_usage,
        )

        if oic_anthropic_integration_enabled():
            full_text = f"{system_msg.strip()}\n\n---\n\n{user_msg.strip()}"
            response = call_oic_anthropic_user_prompt(full_text)
            in_tok = max(1, len(full_text) // 4) if full_text else 0
            out_tok = max(1, len(str(response or "")) // 4) if response else 0
            record_external_usage(
                in_tok,
                out_tok,
                provider="oic_anthropic",
                model="oic_anthropic",
                generation_name="ask_data_analysis",
                estimated=True,
            )
            gemini_used = True

        gemini_key = kwargs.get("gemini_api_key")
        if not gemini_key:
            gemini_key = _resolve_gemini_api_key(self.config)

        # Question-only web snippets via Gemini+Google Search (DF never sent; table stays in Qwen only).
        _web_ctx_injected = False
        _gsearch_for_qwen = bool(
            _skip_gemini_for_qwen
            and gemini_key
            and bool(getattr(_app_settings, "USE_GEMINI_GOOGLE_SEARCH_WITH_QWEN", True))
            and (kwargs.get("gemini_google_search_with_qwen", True) is not False)
        )
        if _gsearch_for_qwen:
            # Only the analysis question goes to Gemini for search — df stays local (Qwen).
            gctx = self.fetch_google_search_context_for_user_query(
                question,
                gemini_api_key=gemini_key,
                gemini_model=kwargs.get("gemini_search_model", "gemini-3-flash-preview"),
            )
            if gctx:
                _web_ctx_injected = True
                system_msg = (
                    system_msg
                    + "\n\n--- Web context (Google Search; only the user question was sent to Gemini) ---\n"
                    + gctx
                )
                print(
                    "[AskDataAnalysis] Google Search snippets added for Qwen (table data not sent to Gemini).",
                    flush=True,
                )
            else:
                print(
                    "[AskDataAnalysis] Google Search for Qwen returned empty (check GEMINI_API_KEY, google-genai, network).",
                    flush=True,
                )

        message_log = [
            self.system_message(system_msg),
            self.user_message(user_msg),
        ]

        if _skip_gemini_for_qwen:
            print(
                "[AskDataAnalysis] Local Qwen (Ollama) configured and selected — using local model for analysis (skipping Gemini).",
                flush=True,
            )
        elif not gemini_used and gemini_key:
            try:
                model = kwargs.get("gemini_model", "gemini-3-flash-preview")
                print(f"[AskDataAnalysis] Using Gemini model={model} (google_search)", flush=True)
                response = _gemini_generate_content(
                    user_text=user_msg,
                    api_key=gemini_key,
                    system_instruction=system_msg,
                    model=model,
                    google_search=kwargs.get("gemini_google_search", True),
                )
                gemini_used = True
            except Exception as e:
                print(f"[AskDataAnalysis] Gemini failed, falling back to OpenAI: {e}", flush=True)
                self.log(f"Gemini ask_data_analysis failed, falling back to default LLM: {e}", title="AskDataAnalysis")
                response = self.submit_prompt(message_log, **kwargs)
        elif not gemini_used:
            print(
                "[AskDataAnalysis] GEMINI_API_KEY not resolved — using default LLM (set GEMINI_API_KEY or add GEMINI_API_KEY to org admin config / initial_config.txt)",
                flush=True,
            )
            response = self.submit_prompt(message_log, **kwargs)

        
        # Strengthening instructions for smaller local models (e.g. Qwen via Ollama)
        extra_instructions = ""
        if kwargs.get("use_qwen"):
            web_rule = ""
            if _web_ctx_injected:
                web_rule = (
                    "\n\nIf the system message includes a 'Web context (Google Search)' section, the user may be asking for "
                    "current prices, rates, or facts that are not in the table. Use that section to answer and say that those "
                    "parts come from web search. You may use OUTPUT_TYPE: TEXT for a primarily web-based answer. Do not refuse "
                    "only because the table lacks fuel/currency/price data if the web context answers the question."
                )
            extra_instructions = (
                "\n\nSTRICT FORMATTING RULE for Qwen / local LLM:\n"
                "You MUST follow this pattern exactly. Do NOT include any preamble, introduction, or conversational filler before the 'OUTPUT_TYPE:' marker. Do NOT repeat the analysis text outside of the 'TEXT:' section.\n"
                "OUTPUT_TYPE: COMBINATION\n"
                "TEXT:\n"
                "(Your concise analysis summary here - focus ONLY on results and insights, do NOT explain methodology or algorithms used)\n"
                "DATAFRAME:\n"
                "```python\n"
                "# pandas, numpy, sklearn, statsmodels (sm), and matplotlib allowed in this section\n"
                "import pandas as pd\n"
                "import numpy as np\n"
                "result_df = ...\n"
                "```\n"
                "CHART:\n"
                "```python\n"
                "# Plotly allowed ONLY in this section\n"
                "import plotly.express as px\n"
                "chart_fig = ...\n"
                "```\n"
                "If you don't need a section, omit it, but ALWAYS start with OUTPUT_TYPE:."
            ) + web_rule

        _user_scope = (
            f"Based on the data provided and the user's request: {question!r}, "
            "determine what type of analysis to perform and generate the appropriate output.\n\n"
        )
        if _web_ctx_injected and kwargs.get("use_qwen"):
            _user_scope += (
                "The system message may include web search results (question-only; no table rows were sent to the search API). "
                "Use that web context for external facts. Combine with the table when the question needs both.\n\n"
            )

        message_log = [
            self.system_message(system_msg),
            self.user_message(
                _user_scope
                + "You can generate:\n"
                "1. TEXT ANALYSIS: A concise summary of results and insights. Do NOT explain the approach, methodology, or models used. (use OUTPUT_TYPE: TEXT)\n"
                "2. NEW DATAFRAME: Provide Python code for 'result_df' (use OUTPUT_TYPE: DATAFRAME)\n"
                "3. CHART: Provide Python plotly code for 'chart_fig' (use OUTPUT_TYPE: CHART)\n"
                "4. COMBINATION: Any combination of the above (use OUTPUT_TYPE: COMBINATION)\n\n"
                "IMPORTANT: Start your response with OUTPUT_TYPE: [TEXT|DATAFRAME|CHART|COMBINATION]\n\n"
                "Label sections with TEXT:, DATAFRAME:, CHART: if using COMBINATION.\n\n"
                + extra_instructions
                + "\n\n"
                + "CRITICAL: You MUST provide the markers OUTPUT_TYPE:, TEXT:, DATAFRAME:, and CHART: exactly as shown above. "
                "The TEXT section should be a concise summary of results. DO NOT include any technical explanation of how the analysis was performed or which algorithms were used.\n"
                "If this is a prediction, you MUST provide a DATAFRAME section with code to generate the future data. "
                "Ensure all arrays used to create the DataFrame (e.g., when calling `pd.DataFrame({...})`) have the same length to avoid 'All arrays must be of the same length' errors.\n\n"
                + self._response_language()
            ),
        ]
        kwargs["llm_scoped_feature"] = "askme"
        if not gemini_used:
            response = self.submit_prompt(message_log, **kwargs)
        self.log(title="Data Analysis Raw Response", message=response)
        
        result = {
            'text': None,
            'dataframe': None,
            'chart_code': None,
            'has_dataframe': False,
            'has_chart': False,
            'has_text': False
        }
        
        # Parse the response - more robustly for Qwen / local LLMs
        response_upper = response.upper()
        
        # Pre-processing: If OUTPUT_TYPE: exists but not at the start, strip the preamble
        first_marker_pos = response_upper.find('OUTPUT_TYPE:')
        if first_marker_pos != -1:
            # Slice response to start from the first marker to avoid preambles
            parsing_response = response[first_marker_pos:]
            parsing_response_upper = response_upper[first_marker_pos:]
        else:
            parsing_response = response
            parsing_response_upper = response_upper

        columns = df.columns.tolist() if hasattr(df, 'columns') else []
        
        # Check for markers anywhere in the parsed response
        is_combination = 'OUTPUT_TYPE: COMBINATION' in parsing_response_upper or 'OUTPUT_TYPE:COMBINATION' in parsing_response_upper
        if not is_combination and ('TEXT:' in parsing_response_upper and ('DATAFRAME:' in parsing_response_upper or 'CHART:' in parsing_response_upper)):
            is_combination = True
            
        if is_combination:
            # Parse combination response from the sliced parsing_response
            text_match = re.search(r'TEXT:\s*(.*?)(?=DATAFRAME:|CHART:|$)', parsing_response, re.DOTALL | re.IGNORECASE)
            dataframe_match = re.search(r'DATAFRAME:\s*(.*?)(?=TEXT:|CHART:|$)', parsing_response, re.DOTALL | re.IGNORECASE)
            chart_match = re.search(r'CHART:\s*(.*?)(?=TEXT:|DATAFRAME:|$)', parsing_response, re.DOTALL | re.IGNORECASE)
            
            if text_match:
                result['text'] = text_match.group(1).strip()
                result['has_text'] = True
            if dataframe_match:
                code = self._extract_python_code(dataframe_match.group(1))
                code = self._sanitize_forecast_code(code)
                result['dataframe'] = self._execute_dataframe_code(code, df)
                
                # Fallback for predictions/forecasts in combination mode if execution failed
                if result['dataframe'] is None and any(word in question.lower() for word in ['predict', 'forecast', 'future', 'trend']):
                    self.log("Forecasting code execution failed or returned None (Combination Mode), using simple fallback", "Data Analysis Fallback")
                    result['dataframe'] = self._simple_forecast_fallback(df, question)
                
                result['has_dataframe'] = result['dataframe'] is not None
            if chart_match:
                # Use generate_json_code_for_chart instead of extracting plotly code
                try:
                    chart_code = self.generate_json_code_for_chart(question=question, columns=columns, model=kwargs.get('model'), use_qwen=kwargs.get('use_qwen', False))
                    result['chart_code'] = chart_code
                    result['has_chart'] = result['chart_code'] is not None
                except Exception as e:
                    # Fallback to original method if generate_json_code_for_chart fails
                    chart_code = chart_match.group(1).strip()
                    result['chart_code'] = self._sanitize_plotly_code(self._extract_python_code(chart_code))
                    result['has_chart'] = result['chart_code'] is not None
                
        elif 'OUTPUT_TYPE: DATAFRAME' in response_upper or 'OUTPUT_TYPE:DATAFRAME' in response_upper:
            # Extract dataframe code
            code = self._extract_python_code(response)
            code = self._sanitize_forecast_code(code)
            result['dataframe'] = self._execute_dataframe_code(code, df)
            
            # Fallback for predictions/forecasts if execution failed
            if result['dataframe'] is None and any(word in question.lower() for word in ['predict', 'forecast', 'future', 'trend']):
                self.log("Forecasting code execution failed or returned None, using simple fallback", "Data Analysis Fallback")
                result['dataframe'] = self._simple_forecast_fallback(df, question)
            
            result['has_dataframe'] = result['dataframe'] is not None
            
        elif 'OUTPUT_TYPE: CHART' in response_upper or 'OUTPUT_TYPE:CHART' in response_upper:
            # Use generate_json_code_for_chart instead of extracting plotly code
            try:
                chart_code = self.generate_json_code_for_chart(question=question, columns=columns, model=kwargs.get('model'), use_qwen=kwargs.get('use_qwen', False))
                result['chart_code'] = chart_code
                result['has_chart'] = result['chart_code'] is not None
            except Exception as e:
                # Fallback to original method if generate_json_code_for_chart fails
                chart_code = self._extract_python_code(response)
                result['chart_code'] = self._sanitize_plotly_code(chart_code)
                result['has_chart'] = result['chart_code'] is not None
            
        elif 'OUTPUT_TYPE: TEXT' in parsing_response_upper or 'OUTPUT_TYPE:TEXT' in parsing_response_upper:
            # Extract text (remove OUTPUT_TYPE line from the SLICED response)
            text = re.sub(r'OUTPUT_TYPE:\s*TEXT\s*', '', parsing_response, flags=re.IGNORECASE).strip()
            result['text'] = text
            result['has_text'] = True
        else:
            # Default: try to detect what was provided
            if 'result_df' in response or 'forecast_df' in response:
                code = self._extract_python_code(response)
                code = self._sanitize_forecast_code(code)
                result['dataframe'] = self._execute_dataframe_code(code, df)
                result['has_dataframe'] = result['dataframe'] is not None
            elif 'chart_fig' in response or 'fig =' in response or 'px.' in response or 'go.' in response or 'chart_json' in response:
                # Use generate_json_code_for_chart instead of extracting plotly code
                try:
                    chart_code = self.generate_json_code_for_chart(question=question, columns=columns, model=kwargs.get('model'), use_qwen=kwargs.get('use_qwen', False))
                    result['chart_code'] = chart_code
                    result['has_chart'] = result['chart_code'] is not None
                except Exception as e:
                    # Fallback to original method if generate_json_code_for_chart fails
                    chart_code = self._extract_python_code(response)
                    result['chart_code'] = self._sanitize_plotly_code(chart_code)
                    result['has_chart'] = result['chart_code'] is not None
            else:
                # Assume it's text analysis
                result['text'] = response.strip()
                result['has_text'] = True
        
        # Keep has_* aligned with actual payload. COMBINATION / Phi 3 paths sometimes populate
        # text or dataframe but leave flags false, which causes API routes to reject the result.
        if result.get("text") is not None and str(result.get("text", "")).strip():
            result["has_text"] = True
        if result.get("dataframe") is not None:
            result["has_dataframe"] = True
        _ccode = result.get("chart_code")
        if _ccode is not None and str(_ccode).strip() not in ("", "N"):
            result["has_chart"] = True

        return result
    
    def generate_json_code_for_chart(self, question: str, columns: List[str], model=None, use_qwen=False) -> str:
        """
        Generate Python code that creates a Plotly-compatible JSON structure from a DataFrame.
        Only sends column names to the LLM, not the actual DataFrame.
        Similar to generate_json_code_cached in AskmeService.
        
        Args:
            question: User's question
            columns: List of column names from the DataFrame
            model: OpenAI model to use (optional)
            
        Returns:
            Python code string that generates Plotly JSON with structure:
            {
                "data": [{"x": [...], "y": [...], "type": "bar", ...}],
                "layout": {"title": "...", "xaxis": {"title": "..."}, "yaxis": {"title": "..."}}
            }
        """
        # Build system message with question and column information
        system_msg = f"The following is a pandas DataFrame that contains the results of the query that answers the question the user asked: '{question}'"
        system_msg += f"\n\nThe DataFrame has the following columns: {', '.join(columns)}"
        
        # Build user message requesting Plotly JSON format
        user_msg = (
            "Generate Python code that creates a Plotly JSON structure from the dataframe. "
            "Assume the data is in a pandas dataframe called 'df'. "
            "The code must create a dictionary in Plotly's native format that can be directly used with react-plotly.js.\n\n"
            "The code must create a dictionary with this exact structure:\n"
            "chart_json = {\n"
            "    'data': [\n"
            "        {\n"
            "            'x': df[column_name].tolist(),  # x-axis values\n"
            "            'y': df[column_name].tolist(),  # y-axis values (if applicable)\n"
            "            'type': 'bar' or 'scatter' or 'line' or 'pie' or 'box' or other plotly types,\n"
            "            'name': 'Series name'  # optional\n"
            "        }\n"
            "    ],\n"
            "    'layout': {\n"
            "        'title': 'Chart title based on the question',\n"
            "        'xaxis': {'title': 'X-axis label'},\n"
            "        'yaxis': {'title': 'Y-axis label'}  # omit for pie charts\n"
            "    }\n"
            "}\n\n"
            "Rules:\n"
            "- Choose the most appropriate chart type ('bar', 'line', 'scatter', 'pie', 'box', etc.) based on the data and question\n"
            "- For bar/line/scatter: use first column for x, second column for y\n"
            "- For pie charts: use first column for labels, second column for values (set 'labels' and 'values' keys instead of x/y)\n"
            "- Avoid hardcoding category filter literals like `df[col] == 'First'` or `df[col] == 2023`. Instead, build traces by iterating over `df[col].dropna().unique()` (or `sorted(...)`) so the code works for any category labels/types.\n"
            "- Create a clear, descriptive title based on the question\n"
            "- Set appropriate axis titles based on column names\n"
            "- If only one column exists, create a bar chart with index as x and the column as y\n"
            "- For multiple columns, you can create multiple traces in the 'data' array\n"
            "\n"
            "- Return only Python code. Do not include any explanations, comments, or natural language -- just the code.\n"
            "- CRITICAL: The code must be valid Python and strictly follow the format: `chart_json = {...}`.\n"
            "- Do not append any descriptive text to variable values."
        )
        
        # Use agent's message methods and submit_prompt
        message_log = [
            self.system_message(system_msg),
            self.user_message(user_msg)
        ]
        
        # Prepare kwargs for submit_prompt
        kwargs = {}
        if model:
            kwargs['model'] = model
        
        if use_qwen:
            kwargs['use_qwen'] = True
            kwargs['llm_scoped_feature'] = 'askme'
        
        json_code = self.submit_prompt(message_log, **kwargs)
        
        # Extract Python code if wrapped in markdown or explanations
        if hasattr(self, '_extract_python_code'):
            json_code = self._extract_python_code(json_code)
        
        json_code = json_code.strip()

        def _chart_json_code_valid(code: str) -> bool:
            if not (code or "").strip():
                return False
            try:
                compile(code, "<chart_json>", "exec")
                return True
            except SyntaxError:
                return False

        # Local Qwen often emits invalid Python (broken strings, truncated fences). Retry once with OpenAI when hybrid mode allows it.
        if not _chart_json_code_valid(json_code) and use_qwen:
            from app.core.config import settings as _settings

            if _settings.llm_mode() != "qwen_only":
                retry_kwargs = dict(kwargs)
                retry_kwargs["use_qwen"] = False
                json_code_fb = self.submit_prompt(message_log, **retry_kwargs)
                if hasattr(self, "_extract_python_code"):
                    json_code_fb = self._extract_python_code(json_code_fb)
                json_code_fb = json_code_fb.strip()
                if _chart_json_code_valid(json_code_fb):
                    json_code = json_code_fb

        return json_code
    
    def _execute_dataframe_code(self, python_code: str, df: pd.DataFrame) -> pd.DataFrame:
        """Execute Python code to generate a dataframe"""
        try:
            python_code = self._strip_markdown_code_fences(python_code or "")

            import pandas as pd
            import numpy as np
            import builtins
            _original_import = builtins.__import__
            
            def safe_import(name, globals=None, locals=None, fromlist=(), level=0):
                allowed_modules = {'pandas', 'numpy', 'pd', 'np', 'datetime', 'date', 'time', 'plotly', 'plotly.express', 'px', 'sklearn', 'joblib', 'scipy', 'statsmodels', 'sm', 'matplotlib', 'plt'}
                module_name = name.split('.')[0]
                if module_name in allowed_modules:
                    return _original_import(name, globals, locals, fromlist, level)
                else:
                    raise ImportError(f"Import of '{name}' is not allowed. Allowed: pandas, numpy, plotly, datetime, sklearn, statsmodels, matplotlib.")
            
            safe_builtins = {
                'abs': abs, 'all': all, 'any': any, 'bool': bool, 'dict': dict,
                'float': float, 'int': int, 'len': len, 'list': list, 'max': max,
                'min': min, 'range': range, 'round': round, 'sorted': sorted,
                'str': str, 'sum': sum, 'tuple': tuple, 'zip': zip, 'print': print,
                'enumerate': enumerate, '__import__': safe_import,
                'pd': pd, 'np': np, 'LinearRegression': LinearRegression, 'sm': sm,
                '__builtins__': {
                    'abs': abs, 'all': all, 'any': any, 'bool': bool, 
                    'dict': dict, 'float': float, 'int': int, 'len': len,
                    'list': list, 'max': max, 'min': min, 'range': range,
                    'round': round, 'sorted': sorted, 'str': str, 'sum': sum,
                    'tuple': tuple, 'zip': zip, 'print': print, 'enumerate': enumerate,
                    '__import__': safe_import
                }
            }
            df_copy = df.copy()
            for col in df_copy.columns:
                if df_copy[col].dtype == 'object':
                    # Try to convert to numeric; invalid values become NaN (errors='coerce' avoids "invalid error value specified" in pandas 2.2+)
                    converted = pd.to_numeric(df_copy[col], errors='coerce')
                    if converted.dtype != 'object':
                        df_copy[col] = converted
            
            local_vars = {
                'df': df_copy,
                'pd': pd,
                'np': np,
                'LinearRegression': LinearRegression,
                'sm': sm
            }
            exec(python_code, safe_builtins, local_vars)
            
            # Look for result_df, forecast_df, or any DataFrame
            for var_name in ['result_df', 'forecast_df']:
                if var_name in local_vars and isinstance(local_vars[var_name], pd.DataFrame):
                    return local_vars[var_name]
            
            # Find any DataFrame in local_vars
            for key, value in local_vars.items():
                if isinstance(value, pd.DataFrame) and key != 'df':
                    return value
                    
            return None
        except Exception as e:
            self.log(f"Error executing dataframe code: {e}", "Data Analysis Error")
            return None

    def _sanitize_forecast_code(self, code: str) -> str:
        """
        Remove or comment out imports of external libraries that are not allowed.
        Only allows pandas and numpy.
        """
        import re
        lines = code.split('\n')
        sanitized_lines = []
        
        for line in lines:
            # Check if line contains an import statement
            if re.match(r'^\s*(from\s+|import\s+)', line):
                # Check if it's trying to import something other than allowed libraries
                if re.search(r'from\s+(prophet|fbprophet|tensorflow|torch)', line, re.IGNORECASE):
                    # Comment out the line
                    sanitized_lines.append(f'# {line}  # Blocked: external library not allowed')
                    continue
                elif re.search(r'import\s+(prophet|fbprophet|tensorflow|torch)', line, re.IGNORECASE):
                    # Comment out the line
                    sanitized_lines.append(f'# {line}  # Blocked: external library not allowed')
                    continue
            
            sanitized_lines.append(line)
        
        return '\n'.join(sanitized_lines)

    # ----------------- Connect to Any Database to run the Generated SQL ----------------- #



    def connect_to_mysql(
        self,
        host: str = None,
        dbname: str = None,
        user: str = None,
        password: str = None,
        port: int = None,
        **kwargs
    ):

        try:
            import pymysql.cursors
        except ImportError:
            raise DependencyError(
                "You need to install required dependencies to execute this method,"
                " run command: \npip install PyMySQL"
            )

        if not host:
            host = os.getenv("HOST")

        if not host:
            raise ImproperlyConfigured("Please set your MySQL host")

        if not dbname:
            dbname = os.getenv("DATABASE")

        if not dbname:
            raise ImproperlyConfigured("Please set your MySQL database")

        if not user:
            user = os.getenv("USER")

        if not user:
            raise ImproperlyConfigured("Please set your MySQL user")

        if not password:
            password = os.getenv("PASSWORD")

        if not password:
            raise ImproperlyConfigured("Please set your MySQL password")

        if not port:
            port = os.getenv("PORT")

        if not port:
            raise ImproperlyConfigured("Please set your MySQL port")

        conn = None

        try:
            conn = pymysql.connect(
                host=host,
                user=user,
                password=password,
                database=dbname,
                port=port,
                cursorclass=pymysql.cursors.DictCursor,
                **kwargs
            )
            self.conn = conn
        except pymysql.Error as e:
            print("Error connecting to MySQL:", e)
            raise ValidationError(e)

        def run_sql_mysql(sql: str) -> Union[pd.DataFrame, None]:
            if conn:
                timeout_ms = self._get_sql_statement_timeout_ms()
                try:
                    conn.ping(reconnect=True)
                    cs = conn.cursor()
                    # MySQL 5.7.8+ / 8.0: abort SELECT after TIMEOUT ms (server-side).
                    try:
                        cs.execute(f"SET SESSION max_execution_time = {int(timeout_ms)}")
                    except Exception:
                        pass
                    cs.execute(sql)
                    results = cs.fetchall()

                    # Create a pandas dataframe from the results
                    df = pd.DataFrame(
                        results, columns=[desc[0] for desc in cs.description]
                    )
                    return df

                except pymysql.Error as e:
                    conn.rollback()
                    err_str = str(e).lower()
                    if (
                        "max_execution_time" in err_str
                        or "query execution was interrupted" in err_str
                    ):
                        raise QueryTimeoutError(
                            f"Query timed out after {self._get_sql_statement_timeout_ms() // 1000} seconds"
                        )
                    raise ValidationError(e)

                except Exception as e:
                    conn.rollback()
                    raise e

        self.run_sql_is_set = True
        self.run_sql = run_sql_mysql
    
        return self.conn

    def connect_to_mssql(
        self,
        server: str = None,
        dbname: str = None,
        user: str = None,
        password: str = None,
        port: int = None,
        **kwargs
    ):

        try:
            import pymssql
        except ImportError:
            raise DependencyError(
                "You need to install required dependencies to execute this method,"
                " run command: \npip install pymssql"
            )

        if not server:
            server = os.getenv("SERVER") or os.getenv("HOST")

        if not server:
            raise ImproperlyConfigured("Please set your SQL Server server (host)")

        if not dbname:
            dbname = os.getenv("DATABASE")

        if not dbname:
            raise ImproperlyConfigured("Please set your SQL Server database")

        if not user:
            user = os.getenv("USER")

        if not user:
            raise ImproperlyConfigured("Please set your SQL Server user")

        if not password:
            password = os.getenv("PASSWORD")

        if not password:
            raise ImproperlyConfigured("Please set your SQL Server password")

        if port is None:
            port = os.getenv("PORT")

        try:
            port_int = int(port) if port not in (None, "") else 1433
        except (TypeError, ValueError):
            port_int = 1433

        conn = None

        try:
            conn = pymssql.connect(
                server=server,
                user=user,
                password=password,
                database=dbname,
                port=port_int,
                **kwargs
            )
            self.conn = conn
        except Exception as e:
            print("Error connecting to SQL Server:", e)
            raise ValidationError(e)

        def run_sql_mssql(sql: str) -> Union[pd.DataFrame, None]:
            if conn:
                try:
                    cs = conn.cursor(as_dict=True)
                    cs.execute(sql)
                    results = cs.fetchall()
                    if not results:
                        return pd.DataFrame()
                    return pd.DataFrame(results)
                except Exception as e:
                    try:
                        conn.rollback()
                    except Exception:
                        pass
                    raise ValidationError(e)

        self.run_sql_is_set = True
        self.run_sql = run_sql_mssql

        return self.conn

    def connect_to_oracle(
        self,
        user: str = None,
        password: str = None,
        dsn: str = None,
        **kwargs
    ):

        """
        Connect to an Oracle db using oracledb package. This is just a helper function to set [`self.run_sql`][askme_openai_base.AskmeBase.run_sql]
        **Example:**
        ```python
        self.connect_to_oracle(
        user="username",
        password="password",
        dsn="host:port/sid",
        )
        ```
        Args:
            USER (str): Oracle db user name.
            PASSWORD (str): Oracle db user password.
            DSN (str): Oracle db host ip - host:port/sid.
        """

        try:
            import oracledb
        except ImportError:

            raise DependencyError(
                "You need to install required dependencies to execute this method,"
                " run command: \npip install oracledb"
            )

        if not dsn:
            dsn = os.getenv("DSN")

        if not dsn:
            raise ImproperlyConfigured("Please set your Oracle dsn which should include host:port/sid")

        if not user:
            user = os.getenv("USER")

        if not user:
            raise ImproperlyConfigured("Please set your Oracle db user")

        if not password:
            password = os.getenv("PASSWORD")

        if not password:
            raise ImproperlyConfigured("Please set your Oracle db password")

        conn = None

        try:
            conn = oracledb.connect(
                user=user,
                password=password,
                dsn=dsn,
                **kwargs
            )
            self.conn = conn
        except oracledb.Error as e:
            raise ValidationError(e)

        def run_sql_oracle(sql: str) -> Union[pd.DataFrame, None]:
            if conn:
                try:
                    sql = sql.rstrip()
                    if sql.endswith(';'): #fix for a known problem with Oracle db where an extra ; will cause an error.
                        sql = sql[:-1]

                    timeout_ms = self._get_sql_statement_timeout_ms()
                    try:
                        conn.call_timeout = int(timeout_ms)
                    except Exception:
                        pass

                    cs = conn.cursor()
                    cs.execute(sql)
                    results = cs.fetchall()

                    # Create a pandas dataframe from the results
                    df = pd.DataFrame(
                        results, columns=[desc[0] for desc in cs.description]
                    )
                    return df

                except oracledb.Error as e:
                    conn.rollback()
                    err_str = str(e).upper()
                    # ORA-01013: user requested cancel of current operation
                    # DPI-1067: call timeout of N ms exceeded
                    if (
                        "ORA-01013" in err_str
                        or "DPI-1067" in err_str
                        or "CALL TIMEOUT" in err_str
                    ):
                        raise QueryTimeoutError(
                            f"Query timed out after {self._get_sql_statement_timeout_ms() // 1000} seconds"
                        )
                    raise ValidationError(e)

                except Exception as e:
                    conn.rollback()
                    raise e
        self.dialect = "Oracle SQL"
        self.run_sql_is_set = True
        self.run_sql = run_sql_oracle

        return self.conn

    def connect_to_sqlite(self, db_path: str = None, **kwargs):
        """
        Connect to a SQLite database file and set ``self.run_sql`` to execute read-only
        SELECT (and WITH … SELECT) queries. Used for internal self-query agents.
        """
        if not db_path or not str(db_path).strip():
            raise ImproperlyConfigured("Please provide db_path for SQLite connection")

        path = str(db_path).strip()
        if path.lower().startswith("sqlite:///"):
            rel = path.replace("sqlite:///", "").lstrip("./")
            path = os.path.abspath(rel)

        if not os.path.isfile(path):
            raise ImproperlyConfigured(f"SQLite db file not found: {path}")

        conn = sqlite3.connect(path, check_same_thread=False, timeout=60)
        conn.row_factory = sqlite3.Row
        self.conn = conn

        _ALLOWED_FIRST = frozenset({"select", "with"})

        def run_sql_sqlite(sql: str) -> Union[pd.DataFrame, None]:
            if not conn:
                return None
            try:
                sql_stripped = (sql or "").strip().rstrip(";")
                if not sql_stripped:
                    raise ValidationError("Empty SQL")
                first_word = sql_stripped.split()[0].lower()
                if first_word not in _ALLOWED_FIRST:
                    raise ValidationError(
                        "Only SELECT queries (including WITH … SELECT) are allowed on the internal database. "
                        f"Got: {first_word.upper()}"
                    )
                cursor = conn.cursor()
                cursor.execute(sql_stripped)
                rows = cursor.fetchall()
                if not rows:
                    if cursor.description:
                        col_names = [d[0] for d in cursor.description]
                        return pd.DataFrame(columns=col_names)
                    return pd.DataFrame()
                col_names = [d[0] for d in cursor.description]
                return pd.DataFrame([dict(r) for r in rows], columns=col_names)
            except ValidationError:
                raise
            except Exception as e:
                raise ValidationError(e)

        self.dialect = "SQLite"
        self.run_sql_is_set = True
        self.run_sql = run_sql_sqlite
        return self.conn

    
    def run_sql(self, sql: str, **kwargs) -> pd.DataFrame:
        """
        Example:
        ```python
        self.run_sql("SELECT * FROM my_table")
        ```

        Run a SQL query on the connected database.

        Args:
            sql (str): The SQL query to run.

        Returns:
            pd.DataFrame: The results of the SQL query.
        """
        raise Exception(
            "You need to connect to a database first by running self.connect_to_snowflake(), self.connect_to_postgres(), similar function, or manually set self.run_sql"
        )

    def ask(
        self,
        question: Union[str, None] = None,
        print_results: bool = True,
        auto_train: bool = True,
        visualize: bool = True,  # if False, will not generate plotly code
        allow_llm_to_see_data: bool = False,
    ) -> Union[
        Tuple[
            Union[str, None],
            Union[pd.DataFrame, None],
            Union[plotly.graph_objs.Figure, None],
        ],
        None,
    ]:
        """
        **Example:**
        ```python
        self.ask("What are the top 10 customers by sales?")
        ```

        Ask Askme360 a question and get the SQL query that answers it.

        Args:
            question (str): The question to ask.
            print_results (bool): Whether to print the results of the SQL query.
            auto_train (bool): Whether to automatically train Askme360 on the question and SQL query.
            visualize (bool): Whether to generate plotly code and display the plotly figure.

        Returns:
            Tuple[str, pd.DataFrame, plotly.graph_objs.Figure]: The SQL query, the results of the SQL query, and the plotly figure.
        """

        if question is None:
            question = input("Enter a question: ")

        try:
            sql = self.generate_sql(question=question, allow_llm_to_see_data=allow_llm_to_see_data)
        except Exception as e:
            print(e)
            return None, None, None

        if print_results:
            try:
                Code = __import__("IPython.display", fromList=["Code"]).Code
                display(Code(sql))
            except Exception as e:
                print(sql)

        if self.run_sql_is_set is False:
            print(
                "If you want to run the SQL query, connect to a database first."
            )

            if print_results:
                return None
            else:
                return sql, None, None

        try:
            df = self.run_sql(sql)

            if print_results:
                try:
                    display = __import__(
                        "IPython.display", fromList=["display"]
                    ).display
                    display(df)
                except Exception as e:
                    print(df)

            if len(df) > 0 and auto_train:
                self.add_question_sql(question=question, sql=sql)
            # Only generate plotly code if visualize is True
            if visualize:
                try:
                    plotly_code = self.generate_plotly_code(
                        question=question,
                        sql=sql,
                        df_metadata=f"Running df.dtypes gives:\n {df.dtypes}",
                    )
                    fig = self.get_plotly_figure(plotly_code=plotly_code, df=df)
                    if print_results:
                        try:
                            display = __import__(
                                "IPython.display", fromlist=["display"]
                            ).display
                            Image = __import__(
                                "IPython.display", fromlist=["Image"]
                            ).Image
                            img_bytes = fig.to_image(format="png", scale=2)
                            display(Image(img_bytes))
                        except Exception as e:
                            fig.show()
                except Exception as e:
                    # Print stack trace
                    traceback.print_exc()
                    print("Couldn't run plotly code: ", e)
                    if print_results:
                        return None
                    else:
                        return sql, df, None
            else:
                return sql, df, None

        except Exception as e:
            print("Couldn't run sql: ", e)
            if print_results:
                return None
            else:
                return sql, None, None
        return sql, df, fig

    def train(
        self,
        question: str = None,
        sql: str = None,
        ddl: str = None,
        documentation: str = None,
        plan: TrainingPlan = None,
    ) -> str:
        """
        **Example:**
        ```python
        self.train()
        ```

        Train Askme360 on a question and its corresponding SQL query.
        If you call it with no arguments, it will check if you connected to a database and it will attempt to train on the metadata of that database.
        If you call it with the sql argument, it's equivalent to [`self.add_question_sql()`][askme_openai_base.AskmeBase.add_question_sql].
        If you call it with the ddl argument, it's equivalent to [`self.add_ddl()`][askme_openai_base.AskmeBase.add_ddl].
        If you call it with the documentation argument, it's equivalent to [`self.add_documentation()`][askme_openai_base.AskmeBase.add_documentation].
        Additionally, you can pass a [`TrainingPlan`][askme_openai_base.types.TrainingPlan] object. Get a training plan with [`self.get_training_plan_generic()`][askme_openai_base.AskmeBase.get_training_plan_generic].

        Args:
            question (str): The question to train on.
            sql (str): The SQL query to train on.
            ddl (str):  The DDL statement.
            documentation (str): The documentation to train on.
            plan (TrainingPlan): The training plan to train on.
        """

        if question and not sql:
            raise ValidationError("Please also provide a SQL query")

        if documentation:
            print("Adding documentation....")
            return self.add_documentation(documentation)

        if sql:
            if question is None:
                question = self.generate_question(sql)
                print("Question generated with sql:", question, "\nAdding SQL...")
            return self.add_question_sql(question=question, sql=sql)

        if ddl:
            print("Adding ddl:", ddl)
            return self.add_ddl(ddl)

        if plan:
            for item in plan._plan:
                if item.item_type == TrainingPlanItem.ITEM_TYPE_DDL:
                    self.add_ddl(item.item_value)
                elif item.item_type == TrainingPlanItem.ITEM_TYPE_IS:
                    self.add_documentation(item.item_value)
                elif item.item_type == TrainingPlanItem.ITEM_TYPE_SQL:
                    self.add_question_sql(question=item.item_name, sql=item.item_value)

    def _get_databases(self) -> List[str]:
        try:
            print("Trying INFORMATION_SCHEMA.DATABASES")
            df_databases = self.run_sql("SELECT * FROM INFORMATION_SCHEMA.DATABASES")
        except Exception as e:
            print(e)
            try:
                print("Trying SHOW DATABASES")
                df_databases = self.run_sql("SHOW DATABASES")
            except Exception as e:
                print(e)
                return []

        return df_databases["DATABASE_NAME"].unique().tolist()

    def _get_information_schema_tables(self, database: str) -> pd.DataFrame:
        df_tables = self.run_sql(f"SELECT * FROM {database}.INFORMATION_SCHEMA.TABLES")

        return df_tables

    def get_training_plan_generic(self, df) -> TrainingPlan:
        """
        This method is used to generate a training plan from an information schema dataframe.

        Basically what it does is breaks up INFORMATION_SCHEMA.COLUMNS into groups of table/column descriptions that can be used to pass to the LLM.

        Args:
            df (pd.DataFrame): The dataframe to generate the training plan from.

        Returns:
            TrainingPlan: The training plan.
        """
        # For each of the following, we look at the df columns to see if there's a match:
        database_column = df.columns[
            df.columns.str.lower().str.contains("database")
            | df.columns.str.lower().str.contains("table_catalog")
        ].to_list()[0]
        schema_column = df.columns[
            df.columns.str.lower().str.contains("table_schema")
        ].to_list()[0]
        table_column = df.columns[
            df.columns.str.lower().str.contains("table_name")
        ].to_list()[0]
        columns = [database_column,
                    schema_column,
                    table_column]
        candidates = ["column_name",
                      "data_type",
                      "comment"]
        matches = df.columns.str.lower().str.contains("|".join(candidates), regex=True)
        columns += df.columns[matches].to_list()

        plan = TrainingPlan([])

        for database in df[database_column].unique().tolist():
            for schema in (
                df.query(f'{database_column} == "{database}"')[schema_column]
                .unique()
                .tolist()
            ):
                for table in (
                    df.query(
                        f'{database_column} == "{database}" and {schema_column} == "{schema}"'
                    )[table_column]
                    .unique()
                    .tolist()
                ):
                    df_columns_filtered_to_table = df.query(
                        f'{database_column} == "{database}" and {schema_column} == "{schema}" and {table_column} == "{table}"'
                    )
                    doc = f"The following columns are in the {table} table in the {database} database:\n\n"
                    doc += df_columns_filtered_to_table[columns].to_markdown()

                    plan._plan.append(
                        TrainingPlanItem(
                            item_type=TrainingPlanItem.ITEM_TYPE_IS,
                            item_group=f"{database}.{schema}",
                            item_name=table,
                            item_value=doc,
                        )
                    )

        return plan

    def get_training_plan_snowflake(
        self,
        filter_databases: Union[List[str], None] = None,
        filter_schemas: Union[List[str], None] = None,
        include_information_schema: bool = False,
        use_historical_queries: bool = True,
    ) -> TrainingPlan:
        plan = TrainingPlan([])

        if self.run_sql_is_set is False:
            raise ImproperlyConfigured("Please connect to a database first.")

        if use_historical_queries:
            try:
                print("Trying query history")
                df_history = self.run_sql(
                    """ select * from table(information_schema.query_history(result_limit => 5000)) order by start_time"""
                )

                df_history_filtered = df_history.query("ROWS_PRODUCED > 1")
                if filter_databases is not None:
                    mask = (
                        df_history_filtered["QUERY_TEXT"]
                        .str.lower()
                        .apply(
                            lambda x: any(
                                s in x for s in [s.lower() for s in filter_databases]
                            )
                        )
                    )
                    df_history_filtered = df_history_filtered[mask]

                if filter_schemas is not None:
                    mask = (
                        df_history_filtered["QUERY_TEXT"]
                        .str.lower()
                        .apply(
                            lambda x: any(
                                s in x for s in [s.lower() for s in filter_schemas]
                            )
                        )
                    )
                    df_history_filtered = df_history_filtered[mask]

                if len(df_history_filtered) > 10:
                    df_history_filtered = df_history_filtered.sample(10)

                for query in df_history_filtered["QUERY_TEXT"].unique().tolist():
                    plan._plan.append(
                        TrainingPlanItem(
                            item_type=TrainingPlanItem.ITEM_TYPE_SQL,
                            item_group="",
                            item_name=self.generate_question(query),
                            item_value=query,
                        )
                    )

            except Exception as e:
                print(e)

        databases = self._get_databases()

        for database in databases:
            if filter_databases is not None and database not in filter_databases:
                continue

            try:
                df_tables = self._get_information_schema_tables(database=database)

                print(f"Trying INFORMATION_SCHEMA.COLUMNS for {database}")
                df_columns = self.run_sql(
                    f"SELECT * FROM {database}.INFORMATION_SCHEMA.COLUMNS"
                )

                for schema in df_tables["TABLE_SCHEMA"].unique().tolist():
                    if filter_schemas is not None and schema not in filter_schemas:
                        continue

                    if (
                        not include_information_schema
                        and schema == "INFORMATION_SCHEMA"
                    ):
                        continue

                    df_columns_filtered_to_schema = df_columns.query(
                        f"TABLE_SCHEMA == '{schema}'"
                    )

                    try:
                        tables = (
                            df_columns_filtered_to_schema["TABLE_NAME"]
                            .unique()
                            .tolist()
                        )

                        for table in tables:
                            df_columns_filtered_to_table = (
                                df_columns_filtered_to_schema.query(
                                    f"TABLE_NAME == '{table}'"
                                )
                            )
                            doc = f"The following columns are in the {table} table in the {database} database:\n\n"
                            doc += df_columns_filtered_to_table[
                                [
                                    "TABLE_CATALOG",
                                    "TABLE_SCHEMA",
                                    "TABLE_NAME",
                                    "COLUMN_NAME",
                                    "DATA_TYPE",
                                    "COMMENT",
                                ]
                            ].to_markdown()

                            plan._plan.append(
                                TrainingPlanItem(
                                    item_type=TrainingPlanItem.ITEM_TYPE_IS,
                                    item_group=f"{database}.{schema}",
                                    item_name=table,
                                    item_value=doc,
                                )
                            )

                    except Exception as e:
                        print(e)
                        pass
            except Exception as e:
                print(e)

        return plan

    def get_plotly_figure(
        self, plotly_code: str, df: pd.DataFrame, dark_mode: bool = True
    ) -> plotly.graph_objs.Figure:
        """
        **Example:**
        ```python
        fig = self.get_plotly_figure(
            plotly_code="fig = px.bar(df, x='name', y='salary')",
            df=df
        )
        fig.show()
        ```
        Get a Plotly figure from a dataframe and Plotly code.

        Args:
            df (pd.DataFrame): The dataframe to use.
            plotly_code (str): The Plotly code to use.

        Returns:
            plotly.graph_objs.Figure: The Plotly figure.
        """
        # Ensure common symbols exist for LLM-generated code.
        # Do not rely on module globals() here; keep the exec namespace explicit.
        import numpy as np
        ldict = {"df": df, "px": px, "go": go, "pd": pd, "np": np}
        try:
            exec(plotly_code, {}, ldict)
            print("i am LLM figure")
            fig = ldict.get("fig", None)
        except Exception as e:
            # Inspect data types
            print("i am custom figure")
            numeric_cols = df.select_dtypes(include=["number"]).columns.tolist()
            categorical_cols = df.select_dtypes(
                include=["object", "category"]
            ).columns.tolist()

            # Decision-making for plot type
            if len(numeric_cols) >= 2:
                # Use the first two numeric columns for a scatter plot
                fig = px.scatter(df, x=numeric_cols[0], y=numeric_cols[1])
            elif len(numeric_cols) == 1 and len(categorical_cols) >= 1:
                # Use a bar plot if there's one numeric and one categorical column
                fig = px.bar(df, x=categorical_cols[0], y=numeric_cols[0])
            elif len(categorical_cols) >= 1 and df[categorical_cols[0]].nunique() < 10:
                # Use a pie chart for categorical data with fewer unique values
                fig = px.pie(df, names=categorical_cols[0])
            else:
                # Default to a simple line plot if above conditions are not met
                fig = px.line(df)

        if fig is None:
            return None

        if dark_mode:
            fig.update_layout(template="plotly_dark")

        return fig

    # --- Reconcile / merge (schema info, prompts, deterministic merge) ---
    MAX_ROWS_FOR_LLM_DATA = 1000

    @staticmethod
    def get_schema_info(
        dataframe: pd.DataFrame,
        include_all_data: bool = False,
        max_rows: Optional[int] = None,
    ) -> Dict[str, Any]:
        """Return schema info (columns, shape, nulls, sample/full rows) for LLM prompt."""
        sample_head = dataframe.head(5).to_dict("records") if len(dataframe) > 0 else []
        sample_tail = dataframe.tail(5).to_dict("records") if len(dataframe) > 5 else []
        out = {
            "columns": dataframe.dtypes.astype(str).to_dict(),
            "shape": dataframe.shape,
            "nulls": dataframe.isnull().sum().to_dict(),
            "sample_head": sample_head,
            "sample_tail": sample_tail,
        }
        if include_all_data and len(dataframe) > 0:
            cap = max_rows if max_rows is not None else AskmeBase.MAX_ROWS_FOR_LLM_DATA
            full = dataframe.head(cap).to_dict("records")
            out["full_data"] = full
            out["full_data_row_count"] = len(full)
        return out

    _RECONCILE_OUTPUT_RULES = """
## result_df rules
1. **Join key (always)**: Identify the column(s) that link the two datasets (e.g. month, id, date). If the same key has different names (e.g. df has "ORDER_MONTH", excel_df has "Sales month"), merge on the matching values and keep exactly ONE key column in result_df. The key column MUST be present in result_df and MUST NOT be NaN for any row.
2. **Answer the user's question**: If the user asks for analysis or calculations (e.g. "percentage change", "average monthly", "growth", "trend", "compare and show difference", "per order", "month-over-month"), include those derived columns in result_df so the answer is clear. If the user only wants a simple compare/reconcile with no specific metrics, keep result_df to the join key and the data columns from both sides (use suffixes _df, _excel when the same name exists in both). Do not add unrelated derived columns when the question does not ask for them.
3. **Key never dropped**: Never drop the join key from result_df; never leave the key column as all-NaN.
"""

    @staticmethod
    def _build_merge_prompt_schema_only(
        df_info: Dict[str, Any],
        excel_info: Dict[str, Any],
        question: str,
    ) -> str:
        """Prompt when do_df=False: only schema and sample rows are provided. LLM infers from samples."""
        return f"""
You are a data analyst. You will receive only the schema and a small sample (first/last 5 rows) of two DataFrames. Your code will run on the full data at execution time.

## User Question:
"{question}"

## Your task
From the schema and samples below, write pandas code that answers the question. Since you do not see all rows, write generic code that works on the full `df` and `excel_df` (e.g. use column names and logical conditions, not hardcoded values from the sample). Prefer robust logic (filters by column types, merge keys, aggregates) over assumptions about specific values.

## DataFrame `df` — schema and sample only
- Columns & types: {df_info['columns']}
- Shape: {df_info['shape']}
- Nulls: {df_info['nulls']}
- First 5 rows: {df_info.get('sample_head', [])}
- Last 5 rows: {df_info.get('sample_tail', [])}

## DataFrame `excel_df` — schema and sample only
- Columns & types: {excel_info['columns']}
- Shape: {excel_info['shape']}
- Nulls: {excel_info['nulls']}
- First 5 rows: {excel_info.get('sample_head', [])}
- Last 5 rows: {excel_info.get('sample_tail', [])}
{AskmeBase._RECONCILE_OUTPUT_RULES}

## Output requirements (your code MUST define these)
1. **result_df**: A single pandas DataFrame that answers the user's question. Always include the join key column (non-NaN). If the question asks for percentage change, averages, growth, or other metrics, add those columns. Use pd.merge with suffixes (_df, _excel) when the same column name exists in both. Add _source and _reconciled only if the question is about reconciliation.
2. **result_dfs** (optional): List of DataFrames if you produce more than one, e.g. result_dfs = [result_df].
3. **reconciliation_result**: Dict with at least "summary" and "status"; add "missing_rows", "extra_rows", "value_mismatches", "column_differences", "recommendations" if relevant.

DataFrames are already loaded as `df` and `excel_df`. Use only pandas (pd) and standard Python.
Return ONLY executable Python code. No JSON, markdown, or explanations.
"""

    @staticmethod
    def _build_merge_prompt_full_data(
        df_info: Dict[str, Any],
        excel_info: Dict[str, Any],
        question: str,
    ) -> str:
        """Prompt when do_df=True: full data is provided so the LLM can produce the best result_df from actual values."""
        n_df = df_info.get("full_data_row_count", len(df_info.get("full_data", [])))
        n_ex = excel_info.get("full_data_row_count", len(excel_info.get("full_data", [])))
        df_data = json.dumps(df_info.get("full_data", []), default=str)
        excel_data = json.dumps(excel_info.get("full_data", []), default=str)

        return f"""
You are a data analyst. You are given the FULL data of both DataFrames below so you can see every row. Use this to produce the best possible result_df that exactly answers the user question.

## User Question:
"{question}"

## Your task
You see all the data. Use it to:
- Choose the correct join keys (same values in both tables; if column names differ e.g. ORDER_MONTH vs Sales month, merge on values and keep one key column in result_df).
- Build result_df that answers the user's question: if they ask for percentage change, average monthly, growth, or other analysis, include those columns; otherwise a clean merge (key + data columns) is enough.
- The same full data is already in memory as `df` and `excel_df` when your code runs — your code will execute on this full data.

## Full data for `df` — {n_df} rows
{df_data}

## Full data for `excel_df` — {n_ex} rows
{excel_data}
{AskmeBase._RECONCILE_OUTPUT_RULES}

## Output requirements (your code MUST define these)
1. **result_df**: The single DataFrame that answers the user's question. Join key column MUST be present and non-NaN. If the question asks for percentage change, averages, growth, or other metrics, include those columns. Use pd.merge with suffixes (_df, _excel) when the same column name exists in both. Add _source and _reconciled only if the question is about reconciliation.
2. **result_dfs** (optional): List of DataFrames if you produce more than one, e.g. result_dfs = [result_df].
3. **reconciliation_result**: Dict with at least "summary" and "status"; add "missing_rows", "extra_rows", "value_mismatches", "column_differences", "recommendations" if relevant.

DataFrames are already loaded as `df` and `excel_df`. Use only pandas (pd) and standard Python.
Return ONLY executable Python code. No JSON, markdown, or explanations.
"""

    @staticmethod
    def build_merge_prompt(
        df_info: Dict[str, Any],
        excel_info: Dict[str, Any],
        question: str = "",
        do_df: bool = False,
    ) -> str:
        """Build LLM prompt: uses full-data prompt when do_df=True, schema-only prompt when do_df=False."""
        if do_df and df_info.get("full_data") and excel_info.get("full_data"):
            return AskmeBase._build_merge_prompt_full_data(df_info, excel_info, question)
        return AskmeBase._build_merge_prompt_schema_only(df_info, excel_info, question)

    def reconcile_and_merge(
        self,
        df: pd.DataFrame,
        excel_df: pd.DataFrame,
        question: str = "",
        do_df: bool = True,
        model: Optional[str] = None,
        max_rows: Optional[int] = None,
    ) -> Tuple[pd.DataFrame, Dict[str, Any]]:
        """
        Generate merge code via LLM from schema/full data, execute on df and excel_df,
        return merged result_df and reconciliation_result dict.
        """
        def _normalise_value(val):
            s = str(val).strip()
            if s in ("nan", "None", ""):
                return ""
            try:
                num = float(s)
                return str(round(num, 2))
            except ValueError:
                return s

        print(f"[reconcile] do_df: {do_df}")
        df = df.copy()
        excel_df = excel_df.copy()
        for col in df.columns:
            df[col] = df[col].apply(_normalise_value)
        for col in excel_df.columns:
            excel_df[col] = excel_df[col].apply(_normalise_value)

        df_info = AskmeBase.get_schema_info(df, include_all_data=do_df, max_rows=max_rows)
        excel_info = AskmeBase.get_schema_info(excel_df, include_all_data=do_df, max_rows=max_rows)
        user_msg = AskmeBase.build_merge_prompt(df_info, excel_info, question=question, do_df=do_df)
        message_log = [
            self.system_message("You are a data engineering expert. Return only executable Python code."),
            self.user_message(user_msg),
        ]
        kwargs = {}
        if model:
            kwargs["model"] = model
        raw_response = self.submit_prompt(message_log, **kwargs)
        merge_code = self._extract_python_code(raw_response) if hasattr(self, "_extract_python_code") else raw_response.strip()
        print(f"[reconcile] LLM generated merge code:\n{merge_code[:500]}")

        try:
            exec_globals = {"pd": pd, "__builtins__": __builtins__}
            exec_locals = {"df": df, "excel_df": excel_df, "pd": pd}
            exec(merge_code, exec_globals, exec_locals)
            result_df = exec_locals.get("result_df", pd.DataFrame())
            reconciliation_result = exec_locals.get("reconciliation_result", {})
        except Exception as exec_err:
            print(f"[reconcile] LLM code exec failed: {exec_err}")
            result_df = pd.DataFrame()
            reconciliation_result = {}

        if result_df.empty and not df.empty and not excel_df.empty:
            print("[reconcile] LLM merge produced empty result — using deterministic fallback")
            result_df, reconciliation_result = AskmeBase._deterministic_merge(df, excel_df)

        return result_df, reconciliation_result

    @staticmethod
    def _find_key_by_value_match(
        df: pd.DataFrame, excel_df: pd.DataFrame
    ) -> Optional[Tuple[str, str]]:
        """If no common column names, find (df_col, excel_col) with same set of string values."""
        for c_df in df.columns:
            set_df = set(df[c_df].fillna("").astype(str).str.strip())
            for c_ex in excel_df.columns:
                set_ex = set(excel_df[c_ex].fillna("").astype(str).str.strip())
                if set_df == set_ex and len(set_df) > 0:
                    return (c_df, c_ex)
        return None

    @staticmethod
    def _deterministic_merge(
        df: pd.DataFrame, excel_df: pd.DataFrame
    ) -> Tuple[pd.DataFrame, Dict[str, Any]]:
        """Deterministic outer-merge fallback — no LLM. Keeps one join key column, builds reconciliation report."""
        common_cols = [c for c in df.columns if c in excel_df.columns]
        df_work = df.copy()
        ex_work = excel_df.copy()
        join_keys = []

        if common_cols:
            best_key = max(common_cols, key=lambda c: df_work[c].nunique())
            join_keys = [best_key]
            key_df_col, key_ex_col = best_key, best_key
        else:
            pair = AskmeBase._find_key_by_value_match(df_work, ex_work)
            if pair:
                key_df_col, key_ex_col = pair
                ex_work = ex_work.rename(columns={key_ex_col: key_df_col})
                join_keys = [key_df_col]
            else:
                df_copy = df_work.copy()
                df_copy["_source"] = "df_only"
                df_copy["_reconciled"] = False
                ex_copy = ex_work.copy()
                ex_copy["_source"] = "excel_only"
                ex_copy["_reconciled"] = False
                result_df = pd.concat([df_copy, ex_copy], ignore_index=True)
                reconciliation_result = {
                    "status": "FAIL",
                    "summary": "No common columns found — datasets concatenated",
                    "missing_rows": [],
                    "extra_rows": [],
                    "value_mismatches": [],
                    "column_differences": {
                        "in_excel_not_df": [c for c in excel_df.columns if c not in df.columns],
                        "in_df_not_excel": [c for c in df.columns if c not in excel_df.columns],
                    },
                    "recommendations": "Ensure both datasets share at least one common key column.",
                }
                return result_df, reconciliation_result

        non_key_common = [c for c in df_work.columns if c in ex_work.columns and c not in join_keys]
        suffixes = ("_df", "_excel") if non_key_common else ("", "")

        merged = pd.merge(
            df_work, ex_work,
            on=join_keys,
            how="outer",
            suffixes=suffixes,
            indicator=True,
        )
        source_map = {"both": "both", "left_only": "df_only", "right_only": "excel_only"}
        merged["_source"] = merged["_merge"].map(source_map)
        merged["_reconciled"] = merged["_source"] == "both"
        merged.drop(columns=["_merge"], inplace=True)

        result_df = merged
        missing = result_df.loc[result_df["_source"] == "excel_only", join_keys].values.tolist()
        extra = result_df.loc[result_df["_source"] == "df_only", join_keys].values.tolist()

        value_mismatches = []
        both_rows = result_df[result_df["_source"] == "both"]
        for col in non_key_common:
            col_df = f"{col}_df"
            col_ex = f"{col}_excel"
            if col_df in result_df.columns and col_ex in result_df.columns:
                diff = both_rows[both_rows[col_df].astype(str) != both_rows[col_ex].astype(str)]
                for _, row in diff.iterrows():
                    value_mismatches.append({
                        "row_id": str(row[join_keys[0]]),
                        "column": col,
                        "df_value": str(row[col_df]),
                        "excel_value": str(row[col_ex]),
                    })

        status = "PASS" if not missing and not extra and not value_mismatches else "FAIL"
        reconciliation_result = {
            "status": status,
            "summary": f"Merged on {join_keys}. {len(missing)} missing, {len(extra)} extra, {len(value_mismatches)} mismatches.",
            "missing_rows": [str(m) for m in missing],
            "extra_rows": [str(e) for e in extra],
            "value_mismatches": value_mismatches,
            "column_differences": {
                "in_excel_not_df": [c for c in excel_df.columns if c not in df.columns],
                "in_df_not_excel": [c for c in df.columns if c not in excel_df.columns],
            },
            "recommendations": "Review mismatches and align data sources." if status == "FAIL" else "Data is fully reconciled.",
        }
        return result_df, reconciliation_result

    # def get_gemini_answer(self, prompt: str ="", **kwargs) -> str:
    #     try:
    #         from google import genai

    #         client = genai.Client(api_key="AIzaSyDNp71YLBrC3jSJWIPICx9cLrd52qmhvKI")

    #         response = client.models.generate_content(
    #         model="gemini-3-pro-preview",
    #         contents=prompt if prompt else "Explain how AI works in a few words",
    #         )
    #         print("Gemini answer:", response.text)
    #         return response.text
          
    #     except Exception as e:
    #         self.log(title="Error", message=f"Failed to get Gemini answer: {str(e)}")
    #         return None


def _resolve_gemini_api_key(config: Optional[Dict[str, Any]] = None) -> Optional[str]:
    """
    Resolve Gemini API key: agent config, env, initial_config.txt, then Settings.
    initial_config keys: gemini_api_key, GEMINI_API_KEY, or GOOGLE_API_KEY (see app.core.config).
    """
    cfg = config or {}
    for key in ("gemini_api_key", "GEMINI_API_KEY", "GOOGLE_API_KEY", "google_api_key"):
        val = str(cfg.get(key) or "").strip()
        if val:
            return val

    for env_name in ("GEMINI_API_KEY", "GOOGLE_API_KEY"):
        val = (os.environ.get(env_name) or "").strip()
        if val:
            return val

    try:
        from app.core.config import get_initial_config

        ic = get_initial_config()
        for key in ("gemini_api_key", "GEMINI_API_KEY", "GOOGLE_API_KEY"):
            val = str(ic.get(key) or "").strip()
            if val:
                return val
    except Exception:
        pass

    try:
        from app.core.config import settings

        val = str(getattr(settings, "GEMINI_API_KEY", None) or "").strip()
        if val:
            return val
    except Exception:
        pass

    return None


def _gemini_generate_content(
    *,
    user_text: str,
    api_key: str,
    system_instruction: str,
    model: str = "gemini-3-flash-preview",
    google_search: bool = True,
) -> str:
    """
    Call Gemini via google-genai SDK (pip install google-genai).
    Used by AskmeBase.ask_data_analysis for the first-pass analysis response.

    When ``oci_call`` is enabled in ``initial_config.txt``, uses OIC only (no Gemini).
    """
    from app.services.llm_service import (
        call_oic_anthropic_user_prompt,
        oic_anthropic_integration_enabled,
        record_external_usage,
    )

    if oic_anthropic_integration_enabled():
        _ = (api_key, model, google_search)
        full = f"{(system_instruction or '').strip()}\n\n---\n\n{(user_text or '').strip()}".strip()
        out_text = call_oic_anthropic_user_prompt(full)
        in_tok = max(1, len(full) // 4) if full else 0
        out_tok = max(1, len(str(out_text or "")) // 4) if out_text else 0
        record_external_usage(
            in_tok,
            out_tok,
            provider="oic_anthropic",
            model="oic_anthropic",
            generation_name="gemini_generate_content",
            estimated=True,
        )
        return out_text

    try:
        from google import genai
        from google.genai import types
    except ImportError as e:
        raise ImportError(
            "Gemini support requires the google-genai package. Install with: pip install google-genai"
        ) from e

    client = genai.Client(api_key=api_key)
    config_kwargs: Dict[str, Any] = {"system_instruction": system_instruction}
    if google_search:
        config_kwargs["tools"] = [types.Tool(google_search=types.GoogleSearch())]

    config = types.GenerateContentConfig(**config_kwargs)
    response = client.models.generate_content(
        model=model,
        contents=user_text,
        config=config,
    )
    text = getattr(response, "text", None)
    if text is not None:
        in_tok = max(1, len(str((system_instruction or "") + "\n" + (user_text or ""))) // 4) if (system_instruction or user_text) else 0
        out_tok = max(1, len(str(text or "")) // 4) if text else 0
        record_external_usage(
            in_tok,
            out_tok,
            provider="gemini",
            model=model,
            generation_name="gemini_generate_content",
            estimated=True,
        )
        return text
    if getattr(response, "candidates", None):
        parts = []
        for c in response.candidates:
            if not c.content or not getattr(c.content, "parts", None):
                continue
            for p in c.content.parts:
                if getattr(p, "text", None):
                    parts.append(p.text)
        out_text = "".join(parts).strip()
        in_tok = max(1, len(str((system_instruction or "") + "\n" + (user_text or ""))) // 4) if (system_instruction or user_text) else 0
        out_tok = max(1, len(str(out_text or "")) // 4) if out_text else 0
        record_external_usage(
            in_tok,
            out_tok,
            provider="gemini",
            model=model,
            generation_name="gemini_generate_content",
            estimated=True,
        )
        return out_text
    return ""
