"""Core pipeline orchestrator — the brain of V3Q.AI.

Pipeline:
  User Question
    → Query Router (classify intent)
    → Context Builder (assemble metadata)
    → Security Gateway → Claude (generate SQL)
    → SQL Validator
    → ** USER CONFIRMATION ** (shows SQL, user approves/rejects)
    → Execute locally on PostgreSQL
    → IF ERROR: Local LLM → abstract feedback → retry
    → Plausibility Check
    → Store Results
    → Summarize (local LLM)
    → Chart Advisor
    → Stream response
"""

from __future__ import annotations

import asyncio
import logging
import time
from typing import Any, AsyncGenerator

from api_client.claude_client import ClaudeClient
from config import AppSettings
from conversations.conversation_manager import ConversationManager
from conversations.result_storage import ResultStorage
from database.connection_manager import ConnectionManager
from database.local_pg import LocalPG
from knowledge.context_builder import ContextBuilder
from knowledge.query_patterns import QueryPatterns
from local_llm.chart_advisor import ChartAdvisor
from local_llm.error_analyzer import ErrorAnalyzer
from local_llm.ollama_client import OllamaClient
from local_llm.plausibility import PlausibilityChecker
from local_llm.summarizer import Summarizer
from orchestrator.query_router import QueryRouter
from orchestrator.retry_logic import RetryLoop
from security.audit_log import AuditLog
from security.gateway import SecurityGateway
from security.sql_validator import validate_sql, extract_tables_from_sql

logger = logging.getLogger(__name__)


class Orchestrator:
    """Main orchestrator that ties all components together."""

    def __init__(self, settings: AppSettings, cm: ConnectionManager) -> None:
        self._settings = settings
        self._cm = cm
        self._local_pg = LocalPG(cm)

        # Local LLM
        self._ollama = OllamaClient(settings.ollama)
        self._router = QueryRouter(self._ollama)
        self._plausibility = PlausibilityChecker(self._ollama)
        self._summarizer = Summarizer(self._ollama)
        self._chart_advisor = ChartAdvisor(self._ollama)
        self._error_analyzer = ErrorAnalyzer(self._ollama)

        # Security
        self._audit = AuditLog(cm)
        self._gateway = SecurityGateway(settings, self._audit)

        # Claude
        self._claude = ClaudeClient(self._gateway)

        # Knowledge
        self._context_builder = ContextBuilder(cm)
        self._patterns = QueryPatterns(cm)

        # Conversations & Results
        self._conversations = ConversationManager(cm)
        self._results = ResultStorage(settings, cm, self._local_pg)

        # Retry
        self._retry = RetryLoop(self._error_analyzer, settings.security.max_retry_attempts)

        # Confirmation state — used by WebSocket handler
        self._pending_confirmation: asyncio.Event | None = None
        self._confirmation_result: bool = False

    async def process_question(
        self,
        question: str,
        conversation_id: str | None = None,
    ) -> AsyncGenerator[dict, None]:
        """Process a user question through the full pipeline.

        Yields status events as dicts for WebSocket streaming:
        - {"type": "stream_start", "conversation_id": str}
        - {"type": "stream_chunk", "content": str}
        - {"type": "status", "message": str}
        - {"type": "confirm_api", "description": str, "metadata_preview": str}
        - {"type": "confirm_sql", "query": str, "tables": [str]}
        - {"type": "sql", "query": str}
        - {"type": "results", "message_id": str, "row_count": int, ...}
        - {"type": "stream_end"}
        - {"type": "error", "message": str}

        The WebSocket handler should call confirm() when user approves,
        or reject() when user rejects a confirmation step.
        """
        start_time = time.time()

        # Ensure conversation exists
        if not conversation_id:
            conv = await self._conversations.create_conversation(question[:80])
            conversation_id = str(conv["id"])
        yield {"type": "stream_start", "conversation_id": conversation_id}

        # Save user message
        await self._conversations.add_message(conversation_id, "user", question)

        # 1. Classify intent
        yield {"type": "status", "message": "Analyzing question..."}
        try:
            intent = await self._router.classify(question)
        except Exception:
            intent = "sql_query"  # Default to SQL if classification fails

        if intent == "conversational":
            response = await self._handle_conversational(question)
            await self._conversations.add_message(conversation_id, "assistant", response)
            yield {"type": "stream_chunk", "content": response}
            yield {"type": "stream_end"}
            return

        if intent == "clarification_needed":
            response = (
                "I'd like to help you query the data, but could you be more specific? "
                "For example, you could ask:\n"
                "- 'How many orders were placed last month?'\n"
                "- 'What are the top 10 customers by revenue?'\n"
                "- 'Show me all products with low stock'"
            )
            await self._conversations.add_message(conversation_id, "assistant", response)
            yield {"type": "stream_chunk", "content": response}
            yield {"type": "stream_end"}
            return

        # 2. Build schema context (metadata only)
        yield {"type": "status", "message": "Building schema context..."}
        schema_context = await self._context_builder.build()

        # Get conversation history for context
        history = await self._conversations.get_conversation_history_for_claude(
            conversation_id, limit=6
        )

        # === CONFIRMATION STEP 1: Show user what will be sent to Claude API ===
        table_count = schema_context.count("TABLE:")
        yield {
            "type": "confirm_api",
            "description": "Send metadata to Claude API to generate SQL query?",
            "metadata_preview": f"{table_count} table schemas · No data values sent",
            "full_payload": {
                "system_prompt": schema_context,
                "user_message": question,
                "conversation_history": history[:-1] if history else [],
            },
        }

        # Wait for user confirmation
        confirmed = await self._wait_for_confirmation()
        if not confirmed:
            await self._conversations.add_message(
                conversation_id, "assistant", "API call cancelled by user."
            )
            yield {"type": "stream_chunk", "content": "API call cancelled."}
            yield {"type": "stream_end"}
            return

        # 3. Generate SQL with Claude
        yield {"type": "status", "message": "Generating SQL query with Claude..."}
        feedback = None
        sql = None
        rows = []
        columns = []
        attempts = 0
        max_retries = self._settings.security.max_retry_attempts

        for attempt in range(1, max_retries + 1):
            attempts = attempt

            # === CONFIRM API CALL on retries (attempt > 1) ===
            if attempt > 1:
                yield {
                    "type": "confirm_api",
                    "description": f"SQL failed. Send retry feedback to Claude API? (attempt {attempt}/{max_retries})",
                    "metadata_preview": f"Retry #{attempt} · {table_count} table schemas · Feedback: {feedback[:100] if feedback else 'N/A'}",
                    "full_payload": {
                        "system_prompt": schema_context,
                        "user_message": question,
                        "feedback": feedback,
                        "conversation_history": history[:-1] if history else [],
                    },
                }
                confirmed = await self._wait_for_confirmation()
                if not confirmed:
                    await self._conversations.add_message(
                        conversation_id, "assistant", "Retry cancelled by user.",
                    )
                    yield {"type": "stream_chunk", "content": "Retry cancelled."}
                    yield {"type": "stream_end"}
                    return
                yield {"type": "status", "message": f"Retrying SQL generation with Claude (attempt {attempt})..."}

            # Generate SQL
            try:
                sql = await self._claude.generate_sql(
                    schema_context=schema_context,
                    user_question=question,
                    conversation_history=history[:-1] if history else None,
                    feedback=feedback,
                )
            except Exception as e:
                logger.error(f"SQL generation failed: {e}")
                yield {"type": "stream_chunk", "content": "Failed to generate SQL query."}
                yield {"type": "stream_end"}
                return

            # Validate SQL
            try:
                validated_sql = validate_sql(sql)
            except Exception as e:
                feedback = f"The generated SQL failed validation: {str(e)}"
                logger.warning(f"SQL validation failed (attempt {attempt}): {e}")
                continue

            # === CONFIRMATION STEP 2: Show user the SQL before executing ===
            tables_used = extract_tables_from_sql(validated_sql)
            yield {
                "type": "confirm_sql",
                "query": validated_sql,
                "tables": tables_used,
                "attempt": attempt,
            }

            confirmed = await self._wait_for_confirmation()
            if not confirmed:
                await self._conversations.add_message(
                    conversation_id, "assistant",
                    f"Query execution cancelled by user.\n\nGenerated SQL:\n```sql\n{validated_sql}\n```",
                    sql_query=validated_sql,
                )
                yield {"type": "stream_chunk", "content": "Query cancelled. The SQL was not executed."}
                yield {"type": "sql", "query": validated_sql}
                yield {"type": "stream_end"}
                return

            # Execute SQL (read-only, local PostgreSQL only)
            yield {"type": "status", "message": "Executing query on local database..."}
            try:
                rows, columns = await self._local_pg.execute_readonly(validated_sql)
                sql = validated_sql
                break  # Success
            except Exception as e:
                error_msg = str(e)
                logger.warning(f"SQL execution failed (attempt {attempt}): {error_msg}")

                try:
                    feedback = await self._error_analyzer.analyze(validated_sql, error_msg)
                except Exception:
                    feedback = ErrorAnalyzer.quick_analyze(error_msg)

                yield {"type": "status", "message": f"Query failed, retrying... ({feedback})"}

        if not rows and not columns:
            error_msg = f"I wasn't able to generate a valid query. {feedback or ''}"
            await self._conversations.add_message(
                conversation_id, "assistant", error_msg, sql_query=sql,
            )
            yield {"type": "stream_chunk", "content": error_msg}
            yield {"type": "stream_end"}
            return

        yield {"type": "sql", "query": sql}
        yield {"type": "status", "message": f"Query returned {len(rows)} rows. Analyzing results..."}

        # 4. Plausibility check (local LLM)
        try:
            plausibility = await self._plausibility.check(
                question=question,
                sql=sql,
                columns=columns,
                row_count=len(rows),
                sample_rows=rows[:5],
            )
            if not plausibility["plausible"]:
                logger.warning(f"Plausibility warning: {plausibility['notes']}")
        except Exception as e:
            logger.warning(f"Plausibility check failed: {e}")

        # 5. Store results
        msg = await self._conversations.add_message(
            conversation_id, "assistant", "",
            sql_query=sql,
            has_results=True,
            row_count=len(rows),
            execution_time_ms=int((time.time() - start_time) * 1000),
        )
        message_id = str(msg["id"])

        storage = await self._results.store_results(message_id, rows, columns)

        await self._cm.execute_write(
            "UPDATE v3q_results.messages SET result_csv_path = :csv, result_table = :tbl WHERE id = :id",
            {"csv": storage["csv_path"], "tbl": storage.get("result_table"), "id": message_id},
        )

        # 6. Summarize (local LLM) — include column descriptions so LLM explains columns
        yield {"type": "status", "message": "Generating summary..."}
        column_descriptions = await self._load_column_descriptions(columns, sql)
        try:
            summary = await self._summarizer.summarize(
                question=question,
                columns=columns,
                row_count=len(rows),
                sample_rows=rows[:10],
                column_descriptions=column_descriptions,
            )
        except Exception as e:
            logger.warning(f"Summarization failed: {e}")
            summary = f"Query returned {len(rows)} rows with columns: {', '.join(columns)}."

        await self._cm.execute_write(
            "UPDATE v3q_results.messages SET content = :content WHERE id = :id",
            {"content": summary, "id": message_id},
        )

        # 7. Chart advisor (local LLM)
        chart_config = None
        try:
            chart_config = await self._chart_advisor.suggest(
                question=question,
                columns=columns,
                row_count=len(rows),
                sample_rows=rows[:5],
            )
            if chart_config:
                import json
                await self._cm.execute_write(
                    "UPDATE v3q_results.messages SET chart_config = :config WHERE id = :id",
                    {"config": json.dumps(chart_config), "id": message_id},
                )
        except Exception as e:
            logger.warning(f"Chart advisor failed: {e}")

        # 8. Store successful pattern
        try:
            tables_used = extract_tables_from_sql(sql)
            await self._patterns.store(question, sql, tables_used)
        except Exception as e:
            logger.warning(f"Pattern storage failed: {e}")

        # 9. Stream results
        yield {"type": "stream_chunk", "content": summary}
        yield {
            "type": "results",
            "message_id": message_id,
            "row_count": len(rows),
            "columns": columns,
            "chart_config": chart_config,
            "csv_path": storage["csv_path"],
        }
        yield {"type": "stream_end"}

    async def _wait_for_confirmation(self, timeout: float = 300) -> bool:
        """Wait for user to confirm or reject an action.

        The WebSocket handler calls confirm() or reject() to resolve this.
        Returns True if confirmed, False if rejected or timed out.
        """
        self._pending_confirmation = asyncio.Event()
        self._confirmation_result = False

        try:
            await asyncio.wait_for(self._pending_confirmation.wait(), timeout=timeout)
        except asyncio.TimeoutError:
            return False
        finally:
            self._pending_confirmation = None

        return self._confirmation_result

    def confirm(self) -> None:
        """User approved the pending action."""
        self._confirmation_result = True
        if self._pending_confirmation:
            self._pending_confirmation.set()

    def reject(self) -> None:
        """User rejected the pending action."""
        self._confirmation_result = False
        if self._pending_confirmation:
            self._pending_confirmation.set()

    async def _load_column_descriptions(self, columns: list[str], sql: str) -> dict[str, str]:
        """Load business descriptions for the columns in the result set.

        Looks up column_registry to find display_name/business_description
        for each column used in the query.
        """
        descriptions: dict[str, str] = {}
        try:
            # Extract table names from SQL to narrow the lookup
            from security.sql_validator import extract_tables_from_sql
            tables_used = extract_tables_from_sql(sql)
            # tables_used format: ["v3q_data.tablename", ...]
            table_names = [t.split(".")[-1] for t in tables_used]

            if not table_names:
                return descriptions

            for table_name in table_names:
                # Look up the table in registry
                table_row = await self._cm.fetch_one(
                    "SELECT id FROM v3q_config.table_registry WHERE local_table = :name",
                    {"name": table_name},
                )
                if not table_row:
                    continue

                # Get column descriptions
                col_rows = await self._cm.execute(
                    "SELECT column_name, display_name, business_description FROM v3q_config.column_registry WHERE table_id = :tid",
                    {"tid": str(table_row["id"])},
                )
                for cr in col_rows:
                    col_name = cr["column_name"]
                    if col_name in columns:
                        desc = cr.get("business_description") or cr.get("display_name") or ""
                        if desc:
                            descriptions[col_name] = desc
        except Exception as e:
            logger.warning(f"Could not load column descriptions: {e}")

        return descriptions

    async def _handle_conversational(self, question: str) -> str:
        """Handle conversational (non-SQL) questions using local LLM."""
        try:
            return await self._ollama.generate(
                prompt=f"The user said: \"{question}\"\n\nYou are V3Q.AI, a business intelligence assistant. Respond helpfully and briefly.",
                system="You are V3Q.AI, a helpful business intelligence assistant that helps users query their business data using natural language.",
                temperature=0.5,
            )
        except Exception:
            return "Hello! I'm V3Q.AI, your business intelligence assistant. Ask me questions about your data and I'll help you find answers."
