"""Databricks Serverless Jobs executor."""

import copy
import json
import logging
import pathlib
import shutil
import subprocess
import sys
import tempfile
import time
import uuid
from typing import Dict, Any

from .base_executor import BaseExecutor
from ..models import CodeResult
from ..settings import settings
from ..kpi_extractor import extract_kpi_from_stdout


class DatabricksExecutor(BaseExecutor):
    """Executor for running Python projects on Databricks Jobs."""

    @property
    def platform_name(self) -> str:
        return "databricks"

    def __init__(self, python_path: str = "python", config: Dict[str, Any] = None):
        super().__init__(python_path, config)

        # Support both flat mode config and nested databricks config.
        if config and isinstance(config.get("databricks"), dict):
            databricks_config = config["databricks"]
        elif config and config.get("mode", "").lower() == "databricks":
            databricks_config = config
        elif config and config.get("databricks"):
            databricks_config = config
        else:
            databricks_config = {}

        self.job_config = databricks_config.get("job", {})
        self.cli = self._resolve_cli_binary(str(databricks_config.get("cli", "databricks")).strip())
        raw_volume_uri = databricks_config.get("volume_uri", "dbfs:/Volumes/workspace/default/volume")
        if raw_volume_uri.startswith("dbfs:/Workspace/"):
            logging.warning(
                "volume_uri '%s' uses dbfs: prefix with a Workspace path. "
                "Normalizing to '%s'. Consider using a Unity Catalog Volume "
                "(dbfs:/Volumes/<catalog>/<schema>/<volume>) for best compatibility.",
                raw_volume_uri,
                raw_volume_uri[len("dbfs:"):],
            )
            self.volume_uri = raw_volume_uri[len("dbfs:"):]
        else:
            self.volume_uri = raw_volume_uri
        self.timeout = databricks_config.get("timeout", f"{settings.script_execution_timeout}s")
        self.cleanup_remote_files = bool(databricks_config.get("cleanup_remote_files", False))

    def execute(self, payload: str | Dict[str, Any]) -> CodeResult:
        """Execute payload on Databricks Jobs.

        Supported payloads:
        - {"project_dir": "...", "entry_command": "..."} (primary path)
        - raw Python code as string (legacy fallback)
        """
        t0 = time.time()
        try:
            if isinstance(payload, dict):
                project_dir = payload.get("project_dir")
                entry_command = payload.get("entry_command")
                if not project_dir or not entry_command:
                    raise ValueError("Databricks payload must include 'project_dir' and 'entry_command'.")
                return self._execute_project(project_dir=project_dir, entry_command=entry_command, t0=t0)

            if isinstance(payload, str) and "\n" not in payload and payload.startswith("co-datascientist-"):
                runtime_ms = int((time.time() - t0) * 1000)
                return CodeResult(
                    stdout=None,
                    stderr=(
                        "Databricks executor received a Docker image tag. "
                        "Databricks mode expects a stitched project payload instead."
                    ),
                    return_code=1,
                    runtime_ms=runtime_ms,
                )

            return self._execute_inline_python(code=str(payload), t0=t0)
        except FileNotFoundError as e:
            runtime_ms = int((time.time() - t0) * 1000)
            msg = (
                "Databricks CLI executable not found. "
                "Install/authenticate Databricks CLI and/or set `databricks.cli` in config.yaml "
                f"to an absolute executable path. Underlying error: {e}"
            )
            logging.error(msg)
            return CodeResult(stdout=None, stderr=msg, return_code=1, runtime_ms=runtime_ms)
        except Exception as e:
            runtime_ms = int((time.time() - t0) * 1000)
            logging.error(f"Failed to execute Databricks job: {e}")
            return CodeResult(
                stdout=None,
                stderr=f"Failed to execute Databricks job: {e}",
                return_code=1,
                runtime_ms=runtime_ms,
            )

    def _execute_project(self, project_dir: str, entry_command: str, t0: float) -> CodeResult:
        project_path = pathlib.Path(project_dir)
        if not project_path.exists():
            raise FileNotFoundError(f"Project directory does not exist: {project_dir}")

        run_suffix = f"{int(t0)}-{uuid.uuid4().hex[:8]}"
        remote_run_dir = f"{self.volume_uri.rstrip('/')}/runs/{run_suffix}"
        remote_project_zip_uri = f"{remote_run_dir}/project.zip"
        remote_launcher_uri = f"{remote_run_dir}/launcher.py"
        project_zip_runtime_path = self._dbfs_uri_to_runtime_path(remote_project_zip_uri)

        # 1) Upload stitched project as a single zip (more robust across Windows/DBFS)
        self._make_remote_dir(remote_run_dir)
        zip_dir = pathlib.Path(tempfile.mkdtemp(prefix="tropiflo-project-zip-"))
        zip_base = zip_dir / "project_bundle"
        zip_path = pathlib.Path(shutil.make_archive(str(zip_base), "zip", root_dir=str(project_path)))
        try:
            self._upload_file(str(zip_path), remote_project_zip_uri)
        finally:
            try:
                shutil.rmtree(zip_dir, ignore_errors=True)
            except Exception:
                pass

        # 2) Upload launcher script that runs configured entry command.
        #    For `python script.py` commands, run in-process so the user's code
        #    inherits the Databricks runtime context (Spark Connect session, etc.).
        #    Fall back to subprocess for non-Python entry commands.
        launcher_code = (
            "import os\n"
            "import re\n"
            "import sys\n"
            "import tempfile\n"
            "import zipfile\n\n"
            f"PROJECT_ZIP = {json.dumps(project_zip_runtime_path)}\n"
            f"ENTRY_COMMAND = {json.dumps(entry_command)}\n\n"
            "EXTRACT_DIR = tempfile.mkdtemp(prefix='tropiflo-project-')\n"
            "with zipfile.ZipFile(PROJECT_ZIP, 'r') as zf:\n"
            "    zf.extractall(EXTRACT_DIR)\n"
            "os.chdir(EXTRACT_DIR)\n"
            "sys.path.insert(0, EXTRACT_DIR)\n"
            "sys.stdout.flush()\n\n"
            "_py_match = re.match(r'^python[23]?\\s+(.+)', ENTRY_COMMAND.strip())\n"
            "if _py_match:\n"
            "    import shlex\n"
            "    _parts = shlex.split(_py_match.group(1))\n"
            "    sys.argv = _parts\n"
            "    _script = _parts[0]\n"
            "    with open(_script) as _fh:\n"
            "        _code = compile(_fh.read(), _script, 'exec')\n"
            "    _g = dict(globals())\n"
            "    _g.update({'__name__': '__main__', '__file__': os.path.abspath(_script)})\n"
            "    exec(_code, _g)\n"
            "else:\n"
            "    import subprocess as _sp\n"
            "    _r = _sp.run(ENTRY_COMMAND, shell=True, stderr=_sp.STDOUT)\n"
            "    if _r.returncode != 0:\n"
            "        sys.exit(_r.returncode)\n"
        )

        with tempfile.NamedTemporaryFile("w", delete=False, suffix=".py") as f:
            f.write(launcher_code)
            launcher_local = pathlib.Path(f.name)

        self._upload_file(str(launcher_local), remote_launcher_uri)
        launcher_local.unlink(missing_ok=True)

        # 3) Submit and wait for one-time Databricks run
        submit_response = self._submit_run(remote_python_file=remote_launcher_uri, script_stem="launcher", project_dir=project_path)
        run_id = submit_response["run_id"]
        timeout_seconds = self._timeout_to_seconds(self.timeout)
        final_run = self._wait_for_run_completion(run_id=run_id, timeout_seconds=timeout_seconds)

        # 4) Fetch logs/output
        task_run_id = run_id
        tasks = final_run.get("tasks") or []
        if tasks and isinstance(tasks[0], dict) and tasks[0].get("run_id"):
            task_run_id = tasks[0]["run_id"]

        output_json = self._run_cli_checked("jobs", "get-run-output", str(task_run_id), "--output", "json")
        out = json.loads(output_json.stdout)

        result_state = (final_run.get("state") or {}).get("result_state", "FAILED")
        life_cycle = (final_run.get("state") or {}).get("life_cycle_state")
        state_message = (final_run.get("state") or {}).get("state_message")
        logs = out.get("logs")
        error_trace = out.get("error_trace")
        out_error = out.get("error")

        rc = 0 if result_state == "SUCCESS" else 1
        stderr_parts = []
        if rc != 0:
            stderr_parts.append(f"Databricks task failed: result_state={result_state}, life_cycle_state={life_cycle}")
            if state_message:
                stderr_parts.append(state_message)
            if out_error:
                stderr_parts.append(str(out_error))
            if error_trace:
                stderr_parts.append(str(error_trace))

        if self.cleanup_remote_files:
            self._remove_remote(remote_run_dir)

        runtime_ms = int((time.time() - t0) * 1000)
        kpi = extract_kpi_from_stdout(logs) if logs else None
        return CodeResult(
            stdout=logs,
            stderr="\n".join(stderr_parts) if stderr_parts else None,
            return_code=rc,
            runtime_ms=runtime_ms,
            kpi=kpi,
        )

    def _execute_inline_python(self, code: str, t0: float) -> CodeResult:
        with tempfile.NamedTemporaryFile("w", delete=False, suffix=".py") as f:
            f.write(code)
            local_tmp = pathlib.Path(f.name)
        remote_uri = f"{self.volume_uri.rstrip('/')}/{local_tmp.name}"
        self._upload_file(str(local_tmp), remote_uri)
        local_tmp.unlink(missing_ok=True)

        submit_response = self._submit_run(remote_python_file=remote_uri, script_stem=local_tmp.stem, project_dir=None)
        run_id = submit_response["run_id"]
        final_run = self._wait_for_run_completion(run_id=run_id, timeout_seconds=self._timeout_to_seconds(self.timeout))
        task_run_id = ((final_run.get("tasks") or [{}])[0]).get("run_id", run_id)

        output_json = self._run_cli_checked("jobs", "get-run-output", str(task_run_id), "--output", "json")
        out = json.loads(output_json.stdout)
        result_state = (final_run.get("state") or {}).get("result_state", "FAILED")

        runtime_ms = int((time.time() - t0) * 1000)
        inline_logs = out.get("logs")
        kpi = extract_kpi_from_stdout(inline_logs) if inline_logs else None
        return CodeResult(
            stdout=inline_logs,
            stderr=None if result_state == "SUCCESS" else str(out.get("error") or "Task failed"),
            return_code=0 if result_state == "SUCCESS" else 1,
            runtime_ms=runtime_ms,
            kpi=kpi,
        )

    def _submit_run(self, remote_python_file: str, script_stem: str, project_dir: pathlib.Path | None) -> Dict[str, Any]:
        ts = str(int(time.time()))
        name_template = self.job_config.get("name", "run-<script-stem>-<timestamp>")
        run_name = name_template.replace("<script-stem>", script_stem).replace("<timestamp>", ts)

        default_task = {
            "task_key": "t",
            "spark_python_task": {"python_file": remote_python_file},
            "environment_key": "default",
        }
        tasks = copy.deepcopy(self.job_config.get("tasks", [default_task]))
        if tasks and isinstance(tasks[0], dict):
            spark_task = tasks[0].setdefault("spark_python_task", {})
            spark_task["python_file"] = remote_python_file

        dependencies = self._infer_dependencies(project_dir) if project_dir else []
        uses_existing_cluster = (
            tasks
            and isinstance(tasks[0], dict)
            and tasks[0].get("existing_cluster_id")
        )

        if uses_existing_cluster:
            # existing_cluster_id is incompatible with environment_key / serverless environments.
            # Route pip dependencies through task-level libraries instead.
            tasks[0].pop("environment_key", None)
            if dependencies:
                pip_libs = [{"pypi": {"package": dep}} for dep in dependencies]
                tasks[0].setdefault("libraries", []).extend(pip_libs)
            environments = []
        else:
            default_spec = {"client": "1"}
            if dependencies:
                default_spec["dependencies"] = dependencies
            default_environments = [{"environment_key": "default", "spec": default_spec}]
            environments = copy.deepcopy(self.job_config.get("environments", default_environments))

        submit_payload = {
            "run_name": run_name,
            "tasks": tasks,
        }
        if environments:
            submit_payload["environments"] = environments

        with tempfile.NamedTemporaryFile("w", delete=False, suffix=".json") as spec:
            json.dump(submit_payload, spec)
            spec_path = spec.name

        submit = self._run_cli_checked("jobs", "submit", "--json", f"@{spec_path}", "--output", "json")
        pathlib.Path(spec_path).unlink(missing_ok=True)
        return json.loads(submit.stdout)

    def _wait_for_run_completion(self, run_id: int, timeout_seconds: int) -> Dict[str, Any]:
        start = time.time()
        sleep_seconds = 3
        while True:
            run = self._run_cli_checked("jobs", "get-run", str(run_id), "--output", "json")
            run_json = json.loads(run.stdout)
            state = run_json.get("state", {})
            life_cycle = state.get("life_cycle_state")
            if life_cycle in {"TERMINATED", "SKIPPED", "INTERNAL_ERROR"}:
                return run_json
            if time.time() - start > timeout_seconds:
                raise TimeoutError(f"Databricks run {run_id} timed out after {timeout_seconds}s")
            time.sleep(sleep_seconds)

    def _infer_dependencies(self, project_dir: pathlib.Path) -> list[str]:
        req_file = project_dir / "requirements.txt"
        if not req_file.exists():
            return []
        deps = []
        for line in req_file.read_text(encoding="utf-8").splitlines():
            stripped = line.strip()
            if stripped and not stripped.startswith("#"):
                deps.append(stripped)
        return deps

    @property
    def _is_workspace_storage(self) -> bool:
        return self.volume_uri.startswith("/Workspace/")

    def _make_remote_dir(self, remote_dir: str) -> None:
        if self._is_workspace_storage:
            self._run_cli_checked("workspace", "mkdirs", remote_dir)
        else:
            self._run_cli_checked("fs", "mkdir", remote_dir)

    def _upload_file(self, local_path: str, remote_uri: str) -> None:
        if self._is_workspace_storage:
            # --format AUTO treats .zip as a DBC archive and extracts it into a
            # folder instead of preserving the binary. Use RAW for zip files.
            fmt = "RAW" if local_path.endswith(".zip") else "AUTO"
            self._run_cli_checked(
                "workspace", "import", remote_uri,
                "--file", local_path,
                "--format", fmt,
                "--overwrite",
            )
        else:
            self._run_cli_checked(
                "fs", "cp", local_path, remote_uri,
                "--overwrite", "--output", "json",
            )

    def _remove_remote(self, remote_dir: str) -> None:
        if self._is_workspace_storage:
            subprocess.run(
                [self.cli, "workspace", "delete", remote_dir, "--recursive"],
                capture_output=True, text=True,
            )
        else:
            subprocess.run(
                [self.cli, "fs", "rm", "-r", remote_dir],
                capture_output=True, text=True,
            )

    def _dbfs_uri_to_runtime_path(self, dbfs_uri: str) -> str:
        if dbfs_uri.startswith("/Workspace/"):
            return dbfs_uri
        if dbfs_uri.startswith("dbfs:/Volumes/"):
            return dbfs_uri.replace("dbfs:/", "/", 1)
        if dbfs_uri.startswith("dbfs:/"):
            return "/dbfs/" + dbfs_uri[len("dbfs:/"):].lstrip("/")
        return dbfs_uri

    def _timeout_to_seconds(self, timeout: Any) -> int:
        if timeout is None:
            return int(settings.script_execution_timeout)
        if isinstance(timeout, (int, float)):
            return int(timeout)
        s = str(timeout).strip().lower()
        if s.endswith("ms"):
            return max(1, int(float(s[:-2]) / 1000))
        if s.endswith("s"):
            return int(float(s[:-1]))
        if s.endswith("m"):
            return int(float(s[:-1]) * 60)
        if s.endswith("h"):
            return int(float(s[:-1]) * 3600)
        return int(float(s))

    def _resolve_cli_binary(self, cli_value: str) -> str:
        """Resolve Databricks CLI binary path across platforms.

        On Windows, `databricks` can live as `databricks.exe` or `databricks.cmd`
        under the active venv's Scripts directory.
        """
        # If caller passed an explicit existing path, use it.
        candidate = pathlib.Path(cli_value)
        if candidate.exists():
            return str(candidate)

        # Try direct lookup first.
        direct = shutil.which(cli_value)
        if direct:
            return direct

        # Try common executable suffixes.
        if not cli_value.lower().endswith(".exe"):
            exe = shutil.which(f"{cli_value}.exe")
            if exe:
                return exe
        if not cli_value.lower().endswith(".cmd"):
            cmd = shutil.which(f"{cli_value}.cmd")
            if cmd:
                return cmd

        # Windows fallback: look in current interpreter's Scripts folder.
        if os_name_is_windows():
            scripts_dir = pathlib.Path(sys.executable).resolve().parent
            for name in ("databricks.exe", "databricks.cmd", "databricks"):
                local_candidate = scripts_dir / name
                if local_candidate.exists():
                    return str(local_candidate)

        # Leave original value so subprocess returns a standard FileNotFoundError.
        return cli_value

    def _run_cli_checked(self, *args: str) -> subprocess.CompletedProcess:
        """Run Databricks CLI command and include stderr/stdout on failures."""
        try:
            return subprocess.run(
                [self.cli, *args],
                check=True,
                capture_output=True,
                text=True,
            )
        except subprocess.CalledProcessError as e:
            stderr = (e.stderr or "").strip()
            stdout = (e.stdout or "").strip()
            details = stderr or stdout or str(e)
            raise RuntimeError(
                f"Databricks CLI command failed: {self.cli} {' '.join(args)}\n{details}"
            ) from e

    def is_available(self) -> bool:
        """Check if Databricks CLI is available and configured."""
        try:
            result = subprocess.run(
                [self.cli, "--version"],
                capture_output=True,
                timeout=10,
            )
            return result.returncode == 0
        except (subprocess.SubprocessError, FileNotFoundError):
            return False


def os_name_is_windows() -> bool:
    return sys.platform.startswith("win")



# OLD CODE but maybe usefull later!
    # async def _save_checkpoint_to_databricks_volume(self, code_content: str, meta_content: str, base_filename: str, config: dict):
    #     """Save checkpoint files directly to Databricks volume using CLI (following existing upload pattern)."""
    #     try:
    #         # Extract databricks configuration (same pattern as _databricks_run_python_code)
    #         if isinstance(config.get('databricks'), dict):
    #             databricks_config = config['databricks']
    #         else:
    #             databricks_config = config
            
    #         CLI = databricks_config.get('cli', "databricks")
    #         VOLUME_URI = databricks_config.get('volume_uri', "dbfs:/Volumes/workspace/default/volume")
            
    #         # Ensure checkpoints directory exists
    #         checkpoints_dir = f"{VOLUME_URI}/{CHECKPOINTS_FOLDER}"
    #         mkdir_result = subprocess.run([CLI, "fs", "mkdir", checkpoints_dir], 
    #                                     capture_output=True, text=True)
    #         # mkdir is okay to fail if directory already exists
            
    #         # Create remote paths
    #         remote_code_path = f"{checkpoints_dir}/{base_filename}.py"
    #         remote_meta_path = f"{checkpoints_dir}/{base_filename}.json"
            
    #         # Save code file using temp file + CLI upload pattern (following existing code)
    #         with tempfile.NamedTemporaryFile(delete=False, suffix=".py") as f:
    #             f.write(code_content.encode())
    #             local_tmp_code = pathlib.Path(f.name)
            
    #         # Try uploading code file
    #         result = subprocess.run([CLI, "fs", "cp", str(local_tmp_code), remote_code_path,
    #                        "--overwrite", "--output", "json"], capture_output=True, text=True)
    #         if result.returncode != 0:
    #             print(f"Failed to upload checkpoint code: {result.stderr}")
    #             return
    #         os.unlink(local_tmp_code)
            
    #         # Save metadata file using temp file + CLI upload pattern
    #         with tempfile.NamedTemporaryFile(delete=False, suffix=".json") as f:
    #             f.write(meta_content.encode())
    #             local_tmp_meta = pathlib.Path(f.name)
            
    #         # Try uploading metadata file
    #         result = subprocess.run([CLI, "fs", "cp", str(local_tmp_meta), remote_meta_path,
    #                        "--overwrite", "--output", "json"], capture_output=True, text=True)
    #         if result.returncode != 0:
    #             print(f"Failed to upload checkpoint metadata: {result.stderr}")
    #             return
    #         os.unlink(local_tmp_meta)
            
    #         # print(f"Checkpoint uploaded to: {VOLUME_URI}/{CHECKPOINTS_FOLDER}/{base_filename}.*")
            
    #     except Exception as e:
    #         print(f"Checkpoint upload error: {e}")

    # async def _save_current_run_to_databricks_volume(self, code_content: str, meta_content: str, config: dict, unique_id: str, timestamp: str):
    #     """Save current run files directly to Databricks volume under `current_runs` directory."""
    #     try:
    #         if isinstance(config.get('databricks'), dict):
    #             databricks_config = config['databricks']
    #         else:
    #             databricks_config = config

    #         CLI = databricks_config.get('cli', "databricks")
    #         VOLUME_URI = databricks_config.get('volume_uri', "dbfs:/Volumes/workspace/default/volume")

    #         # Ensure current_runs directory exists
    #         current_dir = f"{VOLUME_URI}/{CURRENT_RUNS_FOLDER}"
    #         subprocess.run([CLI, "fs", "mkdir", current_dir], capture_output=True, text=True)

    #         remote_code_path = f"{current_dir}/latest.py"
    #         remote_meta_path = f"{current_dir}/latest.json"
    #         uid_safe = _make_filesystem_safe(unique_id)
    #         ts_safe = _make_filesystem_safe(timestamp)
    #         remote_code_uid_path = f"{current_dir}/run_{ts_safe}_{uid_safe}.py"
    #         remote_meta_uid_path = f"{current_dir}/run_{ts_safe}_{uid_safe}.json"

    #         # Upload code
    #         with tempfile.NamedTemporaryFile(delete=False, suffix=".py") as f:
    #             f.write(code_content.encode())
    #             local_tmp_code = pathlib.Path(f.name)
    #         result = subprocess.run([CLI, "fs", "cp", str(local_tmp_code), remote_code_path,
    #                                  "--overwrite", "--output", "json"], capture_output=True, text=True)
    #         os.unlink(local_tmp_code)
    #         if result.returncode != 0:
    #             print(f"Failed to upload current run code: {result.stderr}")
    #             return

    #         # Upload code (UUID version)
    #         with tempfile.NamedTemporaryFile(delete=False, suffix=".py") as f:
    #             f.write(code_content.encode())
    #             local_tmp_code_uid = pathlib.Path(f.name)
    #         result = subprocess.run([CLI, "fs", "cp", str(local_tmp_code_uid), remote_code_uid_path,
    #                                  "--overwrite", "--output", "json"], capture_output=True, text=True)
    #         os.unlink(local_tmp_code_uid)
    #         if result.returncode != 0:
    #             print(f"Failed to upload current run code (uuid): {result.stderr}")
    #             return

    #         # Upload metadata
    #         with tempfile.NamedTemporaryFile(delete=False, suffix=".json") as f:
    #             f.write(meta_content.encode())
    #             local_tmp_meta = pathlib.Path(f.name)
    #         result = subprocess.run([CLI, "fs", "cp", str(local_tmp_meta), remote_meta_path,
    #                                  "--overwrite", "--output", "json"], capture_output=True, text=True)
    #         os.unlink(local_tmp_meta)
    #         if result.returncode != 0:
    #             print(f"Failed to upload current run metadata: {result.stderr}")
    #             return

    #         # Upload metadata (UUID version)
    #         with tempfile.NamedTemporaryFile(delete=False, suffix=".json") as f:
    #             f.write(meta_content.encode())
    #             local_tmp_meta_uid = pathlib.Path(f.name)
    #         result = subprocess.run([CLI, "fs", "cp", str(local_tmp_meta_uid), remote_meta_uid_path,
    #                                  "--overwrite", "--output", "json"], capture_output=True, text=True)
    #         os.unlink(local_tmp_meta_uid)
    #         if result.returncode != 0:
    #             print(f"Failed to upload current run metadata (uuid): {result.stderr}")
    #             return

    #         # print(f"Current run uploaded to: {VOLUME_URI}/{CURRENT_RUNS_FOLDER}/latest.* and run_{uid_safe}.*")
    #     except Exception as e:
    #         print(f"Current run upload error: {e}")