# #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 Infrastructure
# @RELATION DEPENDS_ON -> [ConfigManager]
# @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.

import json
import time
from typing import Any
import uuid

from ...core.config_manager import ConfigManager
from ...core.logger import belief_scope, logger
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: SupersetClient | None = None
        self._database_id: int | None = None
        self._database_backend: str | None = None
        self._database_name: str | None = 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: str | None = None,
        target_database_id: str | None = 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")
                    self._database_name = result.get("database_name", "")
                    logger.reason("Resolved database ID by target_database_id", {
                        "target_database_id": target_database_id,
                        "database_id": self._database_id,
                        "database_name": self._database_name,
                        "backend": self._database_backend,
                        "all_keys": list(result.keys()),
                    })
                    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")
                        self._database_name = db.get("database_name", "")
                        logger.reason("Resolved database ID by name", {
                            "database_name": database_name,
                            "database_id": self._database_id,
                            "actual_name": self._database_name,
                            "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")
            self._database_name = databases[0].get("database_name", "")
            logger.reason("Using default database", {
                "database_name": self._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) -> str | None:
        return self._database_backend
    # endregion get_database_backend

    # region get_database_name [TYPE Function]
    # @PURPOSE Return the cached database name.
    # @POST Returns database name string or None if not yet resolved.
    def get_database_name(self) -> str | None:
        return self._database_name
    # endregion get_database_name

    # 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: int | None = 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)
            }

            # Use the canonical Superset SQL Lab execute endpoint.
            # NOTE: Some older docs reference /sql_lab/execute/ (with underscore),
            # but this endpoint does NOT exist in current Superset versions.
            # The browser UI uses /api/v1/sqllab/execute/ — confirmed by
            # fetch() calls in Superset SQL Lab.
            endpoint = "/api/v1/sqllab/execute/"
            try:
                response = client.network.request(
                    method="POST",
                    endpoint=endpoint,
                    data=json.dumps(payload),
                    headers={"Content-Type": "application/json"},
                )
                logger.reason("SQL Lab endpoint succeeded", extra={
                    "endpoint": endpoint,
                })
            except Exception as ep_err:
                logger.explore("SQL Lab execute failed", extra={
                    "error": str(ep_err),
                    "endpoint": endpoint,
                })
                raise ValueError(
                    f"Superset SQL Lab execute failed: {ep_err}"
                ) from ep_err

            # 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: int | None = 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
