# #region SupersetSqlLabExecutor [C:4] [TYPE Module] [SEMANTICS translate, superset, sqllab, execute, query]
# @BRIEF Submit SQL to Superset SQL Lab API and poll execution status.
# @LAYER: Infra
# @RELATION DEPENDS_ON -> [SupersetClient]
# @RELATION DEPENDS_ON -> [ConfigManager]
# @PRE: Valid Superset environment configuration and authenticated client.
# @POST: SQL is submitted to Superset SQL Lab; execution reference is returned.
# @SIDE_EFFECT: Makes HTTP calls to Superset /api/v1/sqllab/execute/; polls status.
# @RATIONALE: Direct SQL Lab API submission provides audit trail and execution monitoring within Superset.
# @REJECTED: Direct database connection bypass — would skip Superset's SQL Lab audit and RBAC.

from typing import Any, Dict, List, Optional, Tuple
from datetime import datetime, timezone
import json
import time
import uuid

from ...core.logger import logger, belief_scope
from ...core.config_manager import ConfigManager
from ...core.superset_client import SupersetClient


# #region SupersetSqlLabExecutor [C:4] [TYPE Class]
# @BRIEF Submit SQL to Superset SQL Lab API with polling and status tracking.
# @PRE: Valid environment ID and ConfigManager.
# @POST: SQL submitted to Superset; execution reference recorded.
# @SIDE_EFFECT: Makes HTTP POST to /api/v1/sqllab/execute/; polls GET for status.
class SupersetSqlLabExecutor:

    def __init__(self, config_manager: ConfigManager, env_id: str):
        self.config_manager = config_manager
        self.env_id = env_id
        self._client: Optional[SupersetClient] = None
        self._database_id: Optional[int] = None
        self._database_backend: Optional[str] = None

    # region _get_client [TYPE Function]
    # @PURPOSE: Lazy-initialize SupersetClient for the configured environment.
    # @PRE: env_id must correspond to a valid environment config.
    # @POST: Returns authenticated SupersetClient.
    # @SIDE_EFFECT: Authenticates against Superset API.
    def _get_client(self) -> SupersetClient:
        with belief_scope("SupersetSqlLabExecutor._get_client"):
            if self._client is not None:
                return self._client

            environments = self.config_manager.get_environments()
            env_config = next((e for e in environments if e.id == self.env_id), None)
            if not env_config:
                raise ValueError(f"Superset environment '{self.env_id}' not found")

            self._client = SupersetClient(env_config)
            return self._client
    # endregion _get_client

    # region resolve_database_id [TYPE Function]
    # @PURPOSE: Resolve the target database ID from the environment and fetch its backend engine.
    # @PRE: database_name or database_id should be known.
    # @POST: Returns database_id integer or raises ValueError. Also sets self._database_backend.
    # @SIDE_EFFECT: Fetches databases list from Superset; optionally fetches single DB for backend.
    def resolve_database_id(
        self,
        database_name: Optional[str] = None,
        target_database_id: Optional[str] = None,
    ) -> int:
        with belief_scope("SupersetSqlLabExecutor.resolve_database_id"):
            client = self._get_client()

            # If a specific target_database_id is provided, use it directly
            if target_database_id:
                try:
                    db_id = int(target_database_id)
                    # Fetch full DB info to get backend
                    db_info = client.get_database(db_id)
                    result = db_info.get("result", db_info)
                    self._database_id = result.get("id", db_id)
                    self._database_backend = result.get("backend") or result.get("engine")
                    logger.reason("Resolved database ID by target_database_id", {
                        "target_database_id": target_database_id,
                        "database_id": self._database_id,
                        "backend": self._database_backend,
                    })
                    return self._database_id
                except Exception as e:
                    logger.explore("Failed to resolve by target_database_id, falling back", {
                        "target_database_id": target_database_id,
                        "error": str(e),
                    })
                    # Fall through to default resolution

            # Fetch full database list with backend info
            _, databases = client.get_databases(
                query={"columns": ["id", "database_name", "backend"]}
            )
            if not databases:
                raise ValueError("No databases found in Superset environment")

            if database_name:
                for db in databases:
                    if db.get("database_name", "").lower() == database_name.lower():
                        self._database_id = db["id"]
                        self._database_backend = db.get("backend")
                        logger.reason("Resolved database ID by name", {
                            "database_name": database_name,
                            "database_id": self._database_id,
                            "backend": self._database_backend,
                        })
                        return self._database_id
                raise ValueError(
                    f"Database '{database_name}' not found in Superset environment"
                )

            # Default: use first database
            self._database_id = databases[0]["id"]
            self._database_backend = databases[0].get("backend")
            logger.reason("Using default database", {
                "database_name": databases[0].get("database_name"),
                "database_id": self._database_id,
                "backend": self._database_backend,
            })
            return self._database_id
    # endregion resolve_database_id

    # region get_database_backend [TYPE Function]
    # @PURPOSE: Return the cached database backend/engine string.
    # @POST: Returns backend string or None if not yet resolved.
    def get_database_backend(self) -> Optional[str]:
        return self._database_backend
    # endregion get_database_backend

    # region execute_sql [TYPE Function]
    # @PURPOSE: Submit SQL to Superset SQL Lab and return execution tracking info.
    # @PRE: sql is valid SQL string. database_id is a valid Superset DB ID.
    # @POST: Returns execution result dict with query_id and status.
    # @SIDE_EFFECT: Makes HTTP POST to /api/v1/sqllab/execute/.
    def execute_sql(
        self,
        sql: str,
        database_id: Optional[int] = None,
        run_async: bool = True,
    ) -> Dict[str, Any]:
        with belief_scope("SupersetSqlLabExecutor.execute_sql"):
            client = self._get_client()
            db_id = database_id or self._database_id
            if not db_id:
                db_id = self.resolve_database_id()

            logger.reason("Submitting SQL to Superset SQL Lab", {
                "database_id": db_id,
                "sql_length": len(sql),
                "run_async": run_async,
            })

            payload = {
                "database_id": db_id,
                "sql": sql,
                "runAsync": run_async,
                "schema": None,
                "tab": "translation-insert",
                "client_id": f"trl-{uuid.uuid4().hex[:7]}",  # max 11 chars for Superset varchar(11)
            }

            # Try multiple SQL Lab endpoints — some Superset versions use /sqllab/execute/,
            # others use /sql_lab/execute/ (with underscore)
            candidate_endpoints = ["/api/v1/sqllab/execute/", "/api/v1/sql_lab/execute/"]
            response = None
            last_error = None
            for endpoint in candidate_endpoints:
                try:
                    response = client.network.request(
                        method="POST",
                        endpoint=endpoint,
                        data=json.dumps(payload),
                        headers={"Content-Type": "application/json"},
                    )
                    if isinstance(response, dict) and response:
                        logger.reason(f"SQL Lab endpoint succeeded", extra={
                            "endpoint": endpoint,
                        })
                        break
                except Exception as ep_err:
                    last_error = ep_err
                    logger.explore(f"SQL Lab endpoint failed, trying next", extra={
                        "error": str(ep_err),
                        "endpoint": endpoint,
                    })
                    response = None

            if response is None:
                error_msg = f"All SQL Lab endpoints failed. Last error: {last_error}"
                logger.explore("SQL Lab execute failed", extra={"error": error_msg})
                raise ValueError(f"Superset SQL Lab execute failed: {error_msg}")

            # Parse response — try multiple known Superset response formats
            result = response if isinstance(response, dict) else {}

            # Superset may return query_id at top level, nested in "result", as "id", or in "query" block
            query_id = (
                result.get("query_id")
                or result.get("id")
                or (result.get("result") or {}).get("query_id")
                or (result.get("result") or {}).get("id")
                or (result.get("query") or {}).get("query_id")
                or (result.get("query") or {}).get("id")
            )
            status = result.get("status", "unknown")

            # Log full response for debugging if query_id is missing
            if not query_id:
                logger.explore("No query_id from SQL Lab execute", extra={
                    "database_id": db_id,
                    "status": status,
                    "raw_response_keys": list(result.keys()),
                    "raw_response_preview": json.dumps(result)[:2000],
                })
                logger.reason("SQL Lab execute response (no query_id)", extra={
                    "database_id": db_id,
                    "status": status,
                    "response_keys": list(result.keys()),
                    "has_result": "result" in result,
                    "has_query": "query" in result,
                    "raw_preview": json.dumps(result)[:2000],
                })

            logger.reason("SQL Lab execute response", {
                "query_id": query_id,
                "status": status,
                "response_keys": list(result.keys()),
            })

            return {
                "query_id": query_id,
                "status": status,
                "raw_response": result,
                "database_id": db_id,
            }
    # endregion execute_sql

    # region poll_execution_status [TYPE Function]
    # @PURPOSE: Poll Superset for SQL execution status until completion or timeout.
    # @PRE: query_id is a valid Superset query ID.
    # @POST: Returns final execution status dict.
    # @SIDE_EFFECT: Makes HTTP GET requests to Superset.
    def poll_execution_status(
        self,
        query_id: str,
        max_polls: int = 60,
        poll_interval_seconds: float = 2.0,
    ) -> Dict[str, Any]:
        with belief_scope("SupersetSqlLabExecutor.poll_execution_status"):
            client = self._get_client()

            logger.reason("Polling SQL execution status", {
                "query_id": query_id,
                "max_polls": max_polls,
                "interval": poll_interval_seconds,
            })

            for attempt in range(max_polls):
                try:
                    response = client.network.request(
                        method="GET",
                        endpoint=f"/api/v1/query/{query_id}",
                    )
                    raw = response if isinstance(response, dict) else {}
                    # Superset wraps query data in a nested "result" field
                    result = raw.get("result", raw)
                    status = result.get("status", "unknown")
                    state = result.get("state", result.get("status", ""))

                    # Terminal states
                    if state in ("success", "finished", "completed"):
                        logger.reason("SQL execution completed", {
                            "query_id": query_id,
                            "attempt": attempt + 1,
                            "rows_affected": result.get("rows"),
                        })
                        return {
                            "query_id": query_id,
                            "status": "success",
                            "state": state,
                            "rows_affected": result.get("rows"),
                            "error_message": result.get("error_message"),
                            "results": result.get("results"),
                            "completed_on": result.get("completed_on"),
                        }

                    if state in ("failed", "error", "stopped"):
                        error_msg = result.get("error_message", "Unknown error")
                        logger.explore("SQL execution failed", {
                            "query_id": query_id,
                            "state": state,
                            "error": error_msg,
                        })
                        return {
                            "query_id": query_id,
                            "status": "failed",
                            "state": state,
                            "error_message": error_msg,
                            "results": None,
                        }

                    if state in ("pending", "running", "started"):
                        time.sleep(poll_interval_seconds)
                        continue

                    # Unknown state — treat as still running
                    time.sleep(poll_interval_seconds)

                except Exception as e:
                    logger.explore("Polling error, retrying", {
                        "query_id": query_id,
                        "attempt": attempt + 1,
                        "error": str(e),
                    })
                    time.sleep(poll_interval_seconds)
                    continue

            # Timeout
            logger.explore("SQL execution polling timed out", {
                "query_id": query_id,
                "max_polls": max_polls,
            })
            return {
                "query_id": query_id,
                "status": "timeout",
                "error_message": f"Polling timed out after {max_polls} attempts",
                "results": None,
            }
    # endregion poll_execution_status

    # region execute_and_poll [TYPE Function]
    # @PURPOSE: Execute SQL and wait for completion. One-shot convenience method.
    # @PRE: sql is valid SQL.
    # @POST: Returns final execution result.
    # @SIDE_EFFECT: Makes HTTP calls to Superset API.
    def execute_and_poll(
        self,
        sql: str,
        database_id: Optional[int] = None,
        max_polls: int = 60,
        poll_interval_seconds: float = 2.0,
    ) -> Dict[str, Any]:
        with belief_scope("SupersetSqlLabExecutor.execute_and_poll"):
            exec_result = self.execute_sql(
                sql=sql,
                database_id=database_id,
                run_async=False,  # Sync mode — Superset returns result directly
            )

            status = exec_result.get("status", "unknown")

            # Sync mode: response already has the final status and data
            if status == "success":
                logger.reason("SQL execution completed (sync)", {
                    "query_id": exec_result.get("query_id"),
                })
                return {
                    "query_id": exec_result.get("query_id"),
                    "status": "success",
                    "state": "success",
                    "rows_affected": exec_result.get("raw_response", {}).get("query", {}).get("rows"),
                    "error_message": None,
                    "results": exec_result.get("raw_response", {}).get("data"),
                    "completed_on": exec_result.get("raw_response", {}).get("query", {}).get("endDttm"),
                }

            # For async mode (fallback): poll for completion
            query_id = exec_result.get("query_id")
            if not query_id:
                logger.explore("No query_id from SQL Lab execute", {
                    "raw_response": exec_result.get("raw_response"),
                })
                return {
                    "status": "failed",
                    "error_message": "No query_id returned from SQL Lab",
                    "query_id": None,
                }

            return self.poll_execution_status(
                query_id=str(query_id),
                max_polls=max_polls,
                poll_interval_seconds=poll_interval_seconds,
            )
    # endregion execute_and_poll

    # region get_query_results [TYPE Function]
    # @PURPOSE: Fetch the results of a completed query.
    # @PRE: query_id is a valid Superset query ID.
    # @POST: Returns query results if available.
    # @SIDE_EFFECT: Makes HTTP GET to Superset.
    def get_query_results(self, query_id: str) -> Dict[str, Any]:
        with belief_scope("SupersetSqlLabExecutor.get_query_results"):
            client = self._get_client()
            try:
                response = client.network.request(
                    method="GET",
                    endpoint=f"/api/v1/query/{query_id}/results",
                )
                result = response if isinstance(response, dict) else {}
                return {
                    "query_id": query_id,
                    "status": "success" if result.get("results") else "no_results",
                    "results": result.get("results"),
                }
            except Exception as e:
                logger.explore("Failed to fetch query results", {
                    "query_id": query_id,
                    "error": str(e),
                })
                return {
                    "query_id": query_id,
                    "status": "failed",
                    "error_message": str(e),
                    "results": None,
                }
    # endregion get_query_results


# #endregion SupersetSqlLabExecutor
# #endregion SupersetSqlLabExecutor