"""
Flyte Python task that triggers a Databricks job via the Jobs API.

Authentication uses the Databricks SDK OIDC flow: KubernetesIdTokenSource supplies
the Kubernetes service account token (from /var/run/secrets/kubernetes.io/serviceaccount/token),
and the SDK exchanges it for a Databricks access token via the workspace OIDC endpoint.
The task uses the kubernetes_oidc_strategy credentials strategy with Config.

Execution: always calls jobs.run_now() with idempotency_token derived from the Flyte
execution task ID. If wait_for_completion is true, then waits via
wait_get_run_job_terminated_or_skipped(); on timeout the run is canceled and an
error is raised. If wait_for_completion is false, returns immediately after triggering.

Authentication: set DATABRICKS_PERSONAL_ACCESS_TOKEN for token-based auth, or
DATABRICKS_SERVICE_PRINCIPAL_CLIENT_ID for OIDC (K8s service account token exchange).
The Databricks workspace must have OIDC federation configured when using OIDC (e.g. correct audience).
"""

import logging as logger
import os
from datetime import timedelta
from pathlib import Path
from typing import Any, Callable, Dict, Union

from databricks.sdk import WorkspaceClient, oidc
from databricks.sdk.core import (
    Config,
    CredentialsProvider,
    credentials_strategy,
    oidc_credentials_provider,
)
from databricks.sdk.service.jobs import TerminationCodeCode
from flytekit import PythonFunctionTask, current_context
from flytekit.configuration import SerializationSettings
from flytekit.extend import TaskPlugins
from flytekit.extend.backend.base_agent import AsyncAgentExecutorMixin

from truefoundry.deploy.v2.lib.patched_models import DatabricksJobTaskConfig
from truefoundry.workflow.constants import (
    IDEMPOTENCY_TOKEN_ALLOWED_CHARS,
    IDEMPOTENCY_TOKEN_MAX_LEN,
    SERVICE_ACCOUNT_TOKEN_PATH,
)


def _idempotency_token(execution_id: str) -> str:
    """Build a run-now idempotency token from the Flyte execution ID (max 64 safe chars)."""
    raw = str(execution_id)
    safe = "".join(
        c if c in IDEMPOTENCY_TOKEN_ALLOWED_CHARS else "_" for c in raw
    ).strip("_")
    if len(safe) > IDEMPOTENCY_TOKEN_MAX_LEN:
        safe = safe[:IDEMPOTENCY_TOKEN_MAX_LEN]
        logger.debug(
            "Idempotency token truncated from %s chars to %s",
            len(raw),
            IDEMPOTENCY_TOKEN_MAX_LEN,
        )
    return safe


class KubernetesIdTokenSource(oidc.IdTokenSource):
    """
    IdTokenSource that reads the Kubernetes service account JWT from the default
    mount path (or a custom path) for use in Databricks OIDC token exchange.
    """

    def __init__(
        self,
        client_id: str,
        token_path: Union[str, Path] = SERVICE_ACCOUNT_TOKEN_PATH,
    ):
        self.client_id = client_id
        self.token_path = (
            Path(token_path) if not isinstance(token_path, Path) else token_path
        )

    def id_token(self) -> oidc.IdToken:
        if not self.token_path.exists():
            raise FileNotFoundError(
                f"Service account token not found at {self.token_path}. "
                "Ensure the task pod runs with a service account and the token is mounted."
            )
        token = self.token_path.read_text().strip()
        return oidc.IdToken(jwt=token)


@credentials_strategy("kubernetes-oidc", [])
def kubernetes_oidc_strategy(cfg: Config) -> CredentialsProvider:
    """Credentials strategy that uses the K8s service account token for OIDC exchange."""
    id_source = KubernetesIdTokenSource(client_id=cfg.client_id)
    return oidc_credentials_provider(cfg, id_source)


class TfyDatabricksFunctionTask(
    AsyncAgentExecutorMixin, PythonFunctionTask[DatabricksJobTaskConfig]
):
    """
    Flyte plugin that runs a Databricks job via the Databricks SDK.

    Authenticates using the kubernetes-oidc credentials strategy (K8s service account
    token exchanged for a Databricks access token). Triggers the job with run_now()
    using an idempotency token derived from the Flyte execution task ID. When
    wait_for_completion is true, waits with wait_get_run_job_terminated_or_skipped();
    on timeout the run is canceled and a RuntimeError is raised.
    """

    _DATABRICKS_TASK_TYPE = "python-task"

    def __init__(
        self,
        task_config: DatabricksJobTaskConfig,
        task_function: Callable,
        **kwargs,
    ):
        super().__init__(
            task_config=task_config,
            task_type=self._DATABRICKS_TASK_TYPE,
            task_function=task_function,
            **kwargs,
        )

    def get_custom(self, settings: SerializationSettings) -> Dict[str, Any]:
        return {"truefoundry": self._task_config.dict()}

    def execute(self, **kwargs) -> Any:
        cfg = self._task_config
        logger.info(
            "Databricks task execute: workspace_host=%s job_id=%s",
            cfg.workspace_host,
            cfg.job_id,
        )
        ctx = current_context()
        execution_id = ctx.execution_id.name

        idem_token = _idempotency_token(execution_id)
        logger.debug(
            "Flyte execution_id=%s idempotency_token=%s", execution_id, idem_token
        )

        # Get the personal access token from the environment variable
        # If not found, default to using OAuth token federation
        databricks_personal_access_token = os.environ.get(
            "DATABRICKS_PERSONAL_ACCESS_TOKEN"
        )
        if databricks_personal_access_token:
            workspace_client = WorkspaceClient(
                host=cfg.workspace_host,
                token=databricks_personal_access_token,
            )
        elif oidc_client_id := os.environ.get("DATABRICKS_SERVICE_PRINCIPAL_CLIENT_ID"):
            config = Config(
                host=cfg.workspace_host,
                client_id=oidc_client_id,
                credentials_strategy=kubernetes_oidc_strategy,
            )
            workspace_client = WorkspaceClient(config=config)
        else:
            raise RuntimeError(
                "DATABRICKS_SERVICE_PRINCIPAL_CLIENT_ID or DATABRICKS_PERSONAL_ACCESS_TOKEN environment variable is required for using databricks task"
            )
        run = workspace_client.jobs.run_now(
            job_id=int(cfg.job_id),
            idempotency_token=idem_token,
            job_parameters=cfg.job_parameters,
        )

        if not cfg.skip_wait_for_completion:
            timeout = (
                timedelta(seconds=cfg.timeout_seconds)
                if cfg.timeout_seconds is not None
                else timedelta(minutes=20)
            )
            try:
                poll = workspace_client.jobs.wait_get_run_job_terminated_or_skipped(
                    run_id=run.run_id,
                    timeout=timeout,
                )
                if poll.status.termination_details.code != TerminationCodeCode.SUCCESS:
                    logger.error("Databricks job run failed: run_id=%s", run.run_id)
                    raise RuntimeError(
                        f"Databricks job run failed: run_id={run.run_id}"
                    )
                logger.info("Databricks job run completed: run_id=%s", run.run_id)
            except TimeoutError as e:
                workspace_client.jobs.cancel_run(run_id=run.run_id)
                logger.warning(
                    "Canceled Databricks run_id=%s after timeout", run.run_id
                )
                logger.error("Databricks job run timed out: %s", e)
                raise RuntimeError(f"Databricks job run timed out: {e}") from e
        else:
            logger.info(
                "Databricks job triggered (wait_for_completion=false): run_id=%s",
                run.run_id,
            )
        return PythonFunctionTask.execute(self, **kwargs)


TaskPlugins.register_pythontask_plugin(
    DatabricksJobTaskConfig, TfyDatabricksFunctionTask
)
