# ==== CREATED AND OWNED BY askme360 (HIPL) ====
"""
Node-based SQL Generation Architecture

This module implements a pipeline approach to SQL generation:
- PlannerNode: Creates a structured SQL plan from context
- GeneratorNode: Generates SQL based on plan with retry logic
"""

from __future__ import annotations

import json
import re
from abc import ABC, abstractmethod
from dataclasses import dataclass, field
from enum import Enum
from typing import Any, Callable, Dict, List, Optional, TYPE_CHECKING

if TYPE_CHECKING:
    from .askme_base import AskmeBase


class NodeStatus(Enum):
    """Status of a node in the pipeline"""
    PENDING = "pending"
    IN_PROGRESS = "in_progress"
    SUCCESS = "success"
    FAILED = "failed"
    RETRY = "retry"


# Applies to any schema: business wording in questions often does not match exact DB literals.
_DOMAIN_LITERAL_FILTER_HINT = (
    "Categorical and status-like columns: words in the user's question may not exist as exact "
    "stored string literals. Rely on DDL and similar examples for real values. When the intent is "
    "exclusion or \"everything except\" a named state, prefer comparisons consistent with how the "
    "schema encodes data (e.g. `<>` , `!=` , `NOT IN` against documented values) rather than "
    "inventing a literal that only paraphrases natural language."
)


@dataclass
class SQLContext:
    """Context retrieved from ChromaDB for SQL generation"""
    question_sql_list: List[Dict[str, str]] = field(default_factory=list)
    ddl_list: List[str] = field(default_factory=list)
    doc_list: List[str] = field(default_factory=list)


@dataclass
class SQLPlan:
    """Structured plan for SQL generation"""
    tables: List[str] = field(default_factory=list)
    columns: List[str] = field(default_factory=list)
    joins: List[str] = field(default_factory=list)
    filters: List[str] = field(default_factory=list)
    aggregations: List[str] = field(default_factory=list)
    reasoning: str = ""
    raw_response: str = ""


@dataclass
class SQLGenerationResult:
    """Result of SQL generation attempt"""
    sql: str = ""
    is_valid: bool = False
    execution_success: bool = False
    error_message: str = ""
    attempt_number: int = 0


@dataclass
class SQLGenerationState:
    """Main state flowing through the pipeline"""
    question: str
    context: SQLContext = field(default_factory=SQLContext)
    plan: Optional[SQLPlan] = None
    generation_result: Optional[SQLGenerationResult] = None
    retry_count: int = 0
    max_retries: int = 2
    errors: List[str] = field(default_factory=list)
    node_statuses: Dict[str, NodeStatus] = field(default_factory=dict)

    def add_error(self, error: str) -> None:
        """Add an error to the error list"""
        self.errors.append(error)

    def get_error_context(self) -> str:
        """Get formatted error context for retry prompts"""
        if not self.errors:
            return ""
        return "\n".join([f"- Attempt {i+1}: {err}" for i, err in enumerate(self.errors)])


class BaseNode(ABC):
    """Base class for all nodes in the SQL generation pipeline"""

    def __init__(self, askme_instance: 'AskmeBase'):
        self.askme = askme_instance

    @property
    @abstractmethod
    def name(self) -> str:
        """Node name for logging and status tracking"""
        pass

    @abstractmethod
    def process(self, state: SQLGenerationState) -> SQLGenerationState:
        """Process the state and return updated state"""
        pass

    def log(self, message: str, title: str = "Info") -> None:
        """Log a message using the askme instance's log method"""
        self.askme.log(message=message, title=f"[{self.name}] {title}")


class PlannerNode(BaseNode):
    """
    Node 1: Creates a structured SQL plan from context.

    This node analyzes the question and context to create a plan
    identifying tables, columns, joins, filters, and aggregations
    needed for the SQL query.
    """

    @property
    def name(self) -> str:
        return "PlannerNode"

    def _build_planning_prompt(self, state: SQLGenerationState) -> List[Dict[str, Any]]:
        """Build the prompt for SQL planning"""
        dialect = self.askme.dialect

        # Build context summary
        context_parts = []

        if state.context.ddl_list:
            context_parts.append("=== Available Tables (DDL) ===")
            for ddl in state.context.ddl_list:
                context_parts.append(ddl)

        if state.context.doc_list:
            context_parts.append("\n=== Documentation ===")
            for doc in state.context.doc_list:
                context_parts.append(doc)

        if state.context.question_sql_list:
            context_parts.append("\n=== Similar Question-SQL Examples ===")
            for qs in state.context.question_sql_list[:5]:  # Limit to 5 examples
                if qs and "question" in qs and "sql" in qs:
                    context_parts.append(f"Q: {qs['question']}\nSQL: {qs['sql']}\n")

        context_str = "\n".join(context_parts)

        oracle_order_by_guidance = (
            "7. Oracle: For GROUP BY / bucket queries, do not plan ORDER BY that reuses ungrouped base columns; "
            "use ordering by MIN/MAX of the bucket key, by select position, or the same bucket expression as in GROUP BY.\n"
            "8. Oracle: There is no EXTRACT(QUARTER FROM date). Plan quarter as TO_CHAR(date_col, 'Q') or TRUNC(date_col, 'Q').\n"
            "9. Oracle: Do not plan bind variables (:name); SQL runs without binds — use literal IN lists for example values.\n"
            if "ORACLE" in (dialect or "").upper()
            else "\n"
        )

        system_prompt = f"""You are a {dialect} expert SQL planner. Your task is to analyze a question and create a structured plan for generating SQL.

{context_str}

Analyze the user's question and create a SQL plan. DO NOT generate the actual SQL query.
Instead, identify and output a JSON object with the following structure:
{{
    "tables": ["list of tables needed"],
    "columns": ["list of columns to select"],
    "joins": ["list of join conditions if any"],
    "filters": ["list of WHERE conditions"],
    "aggregations": ["list of aggregations like COUNT, SUM, AVG if needed"],
    "reasoning": "brief explanation of your approach"
}}

Important guidelines:
1. Only use tables and columns that exist in the provided DDL—never invent column names.
2. Consider the similar examples for patterns
3. Be specific about column names including table prefixes if there might be ambiguity
4. For filters, describe the condition clearly (e.g., "status = 'active'")
5. Keep reasoning concise (1-2 sentences)
6. {_DOMAIN_LITERAL_FILTER_HINT} Reflect this in the filters array when planning status, category, or flag conditions.
{oracle_order_by_guidance}Respond ONLY with the JSON object, no additional text."""

        return [
            self.askme.system_message(system_prompt),
            self.askme.user_message(f"Create a SQL plan for this question: {state.question}")
        ]

    def _parse_plan_response(self, response: str) -> SQLPlan:
        """Parse the LLM response into a SQLPlan object"""
        plan = SQLPlan(raw_response=response)

        # Try to extract JSON from the response
        try:
            # Try to find JSON in the response
            json_match = re.search(r'\{[\s\S]*\}', response)
            if json_match:
                json_str = json_match.group()
                data = json.loads(json_str)

                plan.tables = data.get("tables", [])
                plan.columns = data.get("columns", [])
                plan.joins = data.get("joins", [])
                plan.filters = data.get("filters", [])
                plan.aggregations = data.get("aggregations", [])
                plan.reasoning = data.get("reasoning", "")
        except (json.JSONDecodeError, AttributeError) as e:
            # If JSON parsing fails, try basic extraction
            self.log(f"JSON parsing failed: {e}, using raw response", title="Warning")
            plan.reasoning = response

        return plan

    def process(self, state: SQLGenerationState) -> SQLGenerationState:
        """Create SQL plan from question and context"""
        state.node_statuses[self.name] = NodeStatus.IN_PROGRESS
        self.log(f"Planning SQL for: {state.question}", title="Start")

        try:
            # Build and submit planning prompt
            prompt = self._build_planning_prompt(state)
            self.log(f"Planning prompt built", title="Prompt")

            # Use the askme instance's prompt submission
            response = self.askme.submit_prompt(prompt)
            self.log(f"Received planning response", title="Response")

            # Parse the response into a plan
            plan = self._parse_plan_response(response)
            state.plan = plan

            self.log(f"Plan created - Tables: {plan.tables}, Columns: {plan.columns}", title="Success")
            state.node_statuses[self.name] = NodeStatus.SUCCESS

        except Exception as e:
            error_msg = f"Planning failed: {str(e)}"
            self.log(error_msg, title="Error")
            state.add_error(error_msg)
            state.node_statuses[self.name] = NodeStatus.FAILED

        return state


class GeneratorNode(BaseNode):
    """
    Node 2: Generates SQL based on plan with retry logic.

    This node takes the SQL plan and generates actual SQL,
    then validates and tests it with LIMIT 1. On failure,
    it retries up to max_retries times with error context.
    """

    @property
    def name(self) -> str:
        return "GeneratorNode"

    def _build_generation_prompt(self, state: SQLGenerationState) -> List[Dict[str, Any]]:
        """Build the prompt for SQL generation"""
        dialect = self.askme.dialect
        max_rows = self.askme.max_rows

        # Build context from plan and DDL
        context_parts = []

        if state.context.ddl_list:
            context_parts.append("=== Available Tables (DDL) ===")
            for ddl in state.context.ddl_list:
                context_parts.append(ddl)

        if state.context.doc_list:
            context_parts.append("\n=== Documentation ===")
            for doc in state.context.doc_list:
                context_parts.append(doc)

        # Add plan information
        if state.plan:
            context_parts.append("\n=== SQL Plan ===")
            context_parts.append(f"Tables to use: {', '.join(state.plan.tables) if state.plan.tables else 'Not specified'}")
            context_parts.append(f"Columns to select: {', '.join(state.plan.columns) if state.plan.columns else 'Not specified'}")
            if state.plan.joins:
                context_parts.append(f"Joins: {', '.join(state.plan.joins)}")
            if state.plan.filters:
                context_parts.append(f"Filters: {', '.join(state.plan.filters)}")
            if state.plan.aggregations:
                context_parts.append(f"Aggregations: {', '.join(state.plan.aggregations)}")
            if state.plan.reasoning:
                context_parts.append(f"Reasoning: {state.plan.reasoning}")

        # Add error context for retries
        error_context = ""
        if state.errors:
            error_context = f"\n=== Previous Errors (Fix these issues) ===\n{state.get_error_context()}\n"

        context_str = "\n".join(context_parts)

        system_prompt = f"""You are a {dialect} expert. Generate a SQL query based on the provided plan and context.

{context_str}
{error_context}

Guidelines:
1. Generate a single, valid {dialect} SQL query
2. Follow the plan's structure (tables, columns, joins, filters, aggregations)
3. Ensure the query is syntactically correct and executable
4. **Identifiers (critical):** Use **only** table and column names that appear in the DDL/context
   above. Never invent, abbreviate, or guess names. If the user asks for a breakdown and
   no matching column exists in DDL, use only documented numeric columns—often a single documented
   cost column for "total cost"—rather than adding nonexistent columns.
5. If no row limit is specified, add a default LIMIT of {max_rows} rows
6. Use appropriate ORDER BY for deterministic results when adding LIMIT
7. Return ONLY the SQL query, no explanations
8. {_DOMAIN_LITERAL_FILTER_HINT}
9. ORACLE-SPECIFIC RULES (apply when dialect contains "ORACLE"):
   - ORA-00979 (GROUP BY): After GROUP BY, every non-aggregated column referenced in SELECT or
     ORDER BY must appear in GROUP BY. In particular, do NOT put ungrouped base columns inside
     ORDER BY CASE (e.g. CASE WHEN overdue_days BETWEEN ... — if only a bucket CASE is in
     GROUP BY, you cannot reference overdue_days again in ORDER BY).
     Safe patterns: ORDER BY 1 (sort by first selected expression); ORDER BY MIN(driving_column)
     or MAX(driving_column) where driving_column is the same column used to define buckets;
     or ORDER BY with the exact same CASE expression as in GROUP BY/SELECT (verbatim).
   - Do not reference a SELECT list alias inside CASE in ORDER BY (separate Oracle limitation).
   - When SELECT uses CASE ... AS bucket and GROUP BY uses the same CASE, the GROUP BY CASE
     must match SELECT character-for-character.
   - Use FETCH FIRST {max_rows} ROWS ONLY (not LIMIT) for row limits.
   - Dates / quarters (ORA-00907): Oracle does NOT support EXTRACT(QUARTER FROM date). Valid
     EXTRACT fields include YEAR, MONTH, DAY, HOUR, MINUTE, SECOND — not QUARTER. For calendar
     quarter use TO_CHAR(date_col, 'Q') (returns '1'–'4') or TRUNC(date_col, 'Q'); group/filter
     by the same expression you use in SELECT.
   - Top-N per group (e.g. highest country per quarter): prefer inner aggregate (GROUP BY country,
     year, quarter) then outer query with RANK() or ROW_NUMBER() OVER (PARTITION BY year, quarter
     ORDER BY sum_col DESC); avoid invalid EXTRACT in PARTITION BY/ORDER BY.
   - ORA-01008 (binds): Do NOT use :variable or :name bind placeholders — the query is executed
     without bind values. Use literal quoted filters (e.g. COUNTRY IN ('Sweden','United States'))
     when the user names example countries; never emit :country_name style placeholders.
   - ORA-01722 (invalid number): Do not use TO_NUMBER() on VARCHAR code/identifier columns
     (names ending in _CODE, carrier codes, SKUs, status text). Those are not amounts. For
     freight/shipping **cost** or monetary totals, SUM only columns the DDL shows as numeric
     money/amount fields (e.g. freight or cost amounts). Never pretend a CODE column is a cost.
   - ORA-00904 (invalid identifier): A referenced column or table does not exist. Use identifiers
     **exactly** as in the DDL (spelling/case). Do not create new column names. If the failed SQL
     used an invented freight/cost column, remove it and use only DDL-listed columns.

If there were previous errors, carefully address each issue in your new query.
When fixing ORA-00979: remove ungrouped columns from ORDER BY; prefer ORDER BY 1,
MIN/MAX of the grouped dimension, or repeat the grouped CASE verbatim — never a new CASE
that still references columns absent from GROUP BY.
When fixing ORA-01722: remove TO_NUMBER from code/identifier columns; use the correct
numeric cost/amount column from the DDL for sums.
When fixing ORA-00904: delete any column/table not in the DDL;
use only columns present in DDL."""

        # Add similar examples as conversation history
        message_log = [self.askme.system_message(system_prompt)]

        for example in state.context.question_sql_list[:3]:  # Limit to 3 examples
            if example and "question" in example and "sql" in example:
                message_log.append(self.askme.user_message(example["question"]))
                message_log.append(self.askme.assistant_message(example["sql"]))

        message_log.append(self.askme.user_message(state.question))

        return message_log

    def _validate_sql(self, sql: str) -> bool:
        """Validate SQL syntax using askme's is_sql_valid method"""
        return self.askme.is_sql_valid(sql)

    def _has_limit_clause(self, sql: str) -> bool:
        """
        Check if SQL already has a limit clause.

        Checks for various SQL dialects:
        - LIMIT (MySQL, PostgreSQL, SQLite)
        - TOP (SQL Server)
        - FETCH FIRST/NEXT (Oracle, DB2)
        - ROWNUM (Oracle)
        - SAMPLE (Teradata)
        """
        sql_upper = sql.upper()
        limit_keywords = ['LIMIT', 'TOP', 'FETCH', 'ROWNUM', 'SAMPLE']
        return any(keyword in sql_upper for keyword in limit_keywords)

    def _apply_max_rows_guardrail(self, sql: str) -> str:
        """
        Apply max_rows limit as a guardrail if SQL doesn't have a limit clause.

        Args:
            sql: The SQL query to check and potentially modify

        Returns:
            str: SQL with max_rows limit applied if needed
        """
        if self._has_limit_clause(sql):
            self.log(f"SQL already has limit clause, skipping guardrail", title="Guardrail")
            return sql

        max_rows = int(self.askme.max_rows)
        limited_sql = self.askme._add_limit_to_sql(sql, limit=max_rows)
        self.log(f"Applied max_rows guardrail: LIMIT {max_rows}", title="Guardrail")
        return limited_sql

    def _test_sql_with_limit(self, sql: str) -> tuple[bool, str]:
        """
        Test SQL execution with LIMIT 1 to verify it runs.

        Returns:
            tuple: (success: bool, error_message: str)
        """
        try:
            # Add LIMIT 1 for testing
            test_sql = self.askme._add_limit_to_sql(sql, limit=1)

            # Run the test query
            result = self.askme.run_sql(test_sql)

            if result is not None:
                return True, ""
            else:
                return False, "Query returned None"

        except Exception as e:
            return False, str(e)

    def _generate_sql_attempt(self, state: SQLGenerationState) -> SQLGenerationResult:
        """Attempt to generate and validate SQL"""
        result = SQLGenerationResult(attempt_number=state.retry_count + 1)

        try:
            # Build and submit generation prompt
            prompt = self._build_generation_prompt(state)
            response = self.askme.submit_prompt(prompt)

            # Extract SQL from response
            sql = self.askme.extract_sql(response)
            result.sql = sql

            if not sql:
                result.error_message = "No SQL extracted from response"
                return result

            if self.askme.sql_has_oracle_bind_placeholders(sql):
                result.error_message = (
                    "ORA-01008: Bind placeholders (:name) are not supported — SQL runs without bind parameters. "
                    "Rewrite with literal quoted values in WHERE (e.g. COUNTRY IN ('Sweden','United States')) "
                    "when the question lists example countries; do not use :variables."
                )
                return result

            # Validate syntax
            if not self._validate_sql(sql):
                result.error_message = "SQL syntax validation failed - not a valid SELECT statement"
                return result

            result.is_valid = True

            # Test execution with LIMIT 1
            success, error = self._test_sql_with_limit(sql)
            if success:
                # Apply max_rows guardrail if no limit clause exists
                sql = self._apply_max_rows_guardrail(sql)
                result.sql = sql
                result.execution_success = True
                self.log(f"SQL validated and tested successfully", title="Success")
            else:
                result.error_message = f"Execution test failed: {error}"

        except Exception as e:
            result.error_message = f"Generation error: {str(e)}"

        return result

    def process(self, state: SQLGenerationState) -> SQLGenerationState:
        """Generate SQL with retry logic"""
        state.node_statuses[self.name] = NodeStatus.IN_PROGRESS
        self.log(f"Generating SQL for: {state.question}", title="Start")

        while state.retry_count <= state.max_retries:
            attempt_num = state.retry_count + 1
            self.log(f"Attempt {attempt_num}/{state.max_retries + 1}", title="Attempt")

            # Generate SQL
            result = self._generate_sql_attempt(state)
            state.generation_result = result

            if result.execution_success:
                state.node_statuses[self.name] = NodeStatus.SUCCESS
                self.log(f"SQL generation successful on attempt {attempt_num}", title="Success")
                return state

            # Handle failure — include full SQL so the LLM can see exactly which
            # clause caused the error (e.g. ORDER BY alias-in-CASE for ORA-00979).
            error_msg = result.error_message or "Unknown error"
            self.log(f"Attempt {attempt_num} failed: {error_msg}", title="Retry")
            if result.sql:
                state.add_error(
                    f"{error_msg}\nFailed SQL:\n```sql\n{result.sql}\n```"
                )
            else:
                state.add_error(error_msg)

            state.retry_count += 1

            if state.retry_count > state.max_retries:
                state.node_statuses[self.name] = NodeStatus.FAILED
                self.log(f"All {state.max_retries + 1} attempts failed - No SQL generated", title="Failed")
                # Set result to indicate no SQL was generated after all attempts
                state.generation_result = SQLGenerationResult(
                    sql="",
                    is_valid=False,
                    execution_success=False,
                    error_message="No SQL generated after all retry attempts",
                    attempt_number=state.retry_count
                )
                break

            state.node_statuses[self.name] = NodeStatus.RETRY

        return state


class SQLGenerationPipeline:
    """
    Orchestrates context retrieval and node execution for SQL generation.

    Pipeline Flow:
    1. Context Retrieval (similar questions, DDL, documentation)
    2. PlannerNode (creates SQL plan)
    3. GeneratorNode (generates and tests SQL with retry)
    """

    def __init__(self, askme_instance: 'AskmeBase', max_retries: int = 2):
        self.askme = askme_instance
        self.max_retries = max_retries
        self.planner = PlannerNode(askme_instance)
        self.generator = GeneratorNode(askme_instance)

    def log(self, message: str, title: str = "Info") -> None:
        """Log a message using the askme instance's log method"""
        self.askme.log(message=message, title=f"[Pipeline] {title}")

    def _retrieve_context(self, question: str, **kwargs) -> SQLContext:
        """Retrieve context from ChromaDB, or use injected ddl_list/question_sql_list/doc_list from kwargs when provided (e.g. per-agent context)."""
        context = SQLContext()

        if "question_sql_list" in kwargs:
            qs = kwargs["question_sql_list"]
            context.question_sql_list = qs if isinstance(qs, list) else []
        else:
            try:
                context.question_sql_list = self.askme.get_similar_question_sql(question, **kwargs) or []
            except Exception as e:
                self.log(f"Failed to get similar questions: {e}", title="Warning")

        if "ddl_list" in kwargs:
            ddl = kwargs["ddl_list"]
            context.ddl_list = ddl if isinstance(ddl, list) else [ddl] if ddl else []
        else:
            try:
                context.ddl_list = self.askme.get_related_ddl(question, **kwargs) or []
            except Exception as e:
                self.log(f"Failed to get related DDL: {e}", title="Warning")

        if "doc_list" in kwargs:
            doc = kwargs["doc_list"]
            context.doc_list = doc if isinstance(doc, list) else [doc] if doc else []
        else:
            try:
                context.doc_list = self.askme.get_related_documentation(question, **kwargs) or []
            except Exception as e:
                self.log(f"Failed to get documentation: {e}", title="Warning")

        # Add static documentation if available
        if self.askme.static_documentation:
            context.doc_list.append(self.askme.static_documentation)

        return context

    def run(self, question: str, **kwargs) -> SQLGenerationState:
        """
        Run the full SQL generation pipeline.

        Args:
            question: The natural language question to convert to SQL
            **kwargs: Additional arguments passed to context retrieval

        Returns:
            SQLGenerationState: The final state with generated SQL or errors
        """
        self.log(f"Starting pipeline for: {question}", title="Start")

        # Initialize state
        state = SQLGenerationState(
            question=question,
            max_retries=self.max_retries
        )

        # Step 1: Retrieve context
        self.log("Retrieving context...", title="Context")
        state.context = self._retrieve_context(question, **kwargs)
        self.log(
            f"Context retrieved - DDL: {len(state.context.ddl_list)}, "
            f"Docs: {len(state.context.doc_list)}, "
            f"Examples: {len(state.context.question_sql_list)}",
            title="Context"
        )

        # Step 2: Run PlannerNode
        self.log("Running PlannerNode...", title="Planning")
        state = self.planner.process(state)

        if state.node_statuses.get(self.planner.name) == NodeStatus.FAILED:
            self.log("Planning failed, attempting generation without plan", title="Warning")

        # Step 3: Run GeneratorNode
        self.log("Running GeneratorNode...", title="Generation")
        state = self.generator.process(state)

        # Log final status
        final_status = state.node_statuses.get(self.generator.name, NodeStatus.FAILED)
        if final_status == NodeStatus.SUCCESS:
            self.log(f"Pipeline completed successfully", title="Complete")
        else:
            self.log(f"Pipeline completed with status: {final_status.value}", title="Complete")

        return state

    def get_sql(self, question: str, **kwargs) -> str:
        """
        Convenience method to run pipeline and return just the SQL.

        Args:
            question: The natural language question to convert to SQL
            **kwargs: Additional arguments passed to context retrieval

        Returns:
            str: The generated SQL query or error message
        """
        state = self.run(question, **kwargs)

        if state.generation_result and state.generation_result.sql:
            if state.generation_result.execution_success:
                return state.generation_result.sql
            else:
                # Return SQL even if execution test failed (might still be useful)
                return state.generation_result.sql

        # Return error message if no SQL was generated
        if state.errors:
            return f"SQL generation failed: {state.errors[-1]}"

        return "SQL generation failed: Unknown error"