{#- Context variables:
   - databricks_profile: str | None
-#}
"""Fetch error details for a failed Databricks job run.

Retrieves task-level errors, skipped (upstream-failed) tasks, and SQL query
failures from the Query History API — giving a complete picture of what went
wrong and which tasks were affected downstream.

Usage:
    python scripts/get_job_errors.py <JOB_RUN_ID> [--profile PROFILE] [--json]

The argument is the **job run ID** (top-level run), NOT a task run ID.
The script inspects every task within that run automatically.

Examples:
    python scripts/get_job_errors.py 35295526214877
    python scripts/get_job_errors.py 35295526214877 --profile {{ databricks_profile or 'DEFAULT' }}
    python scripts/get_job_errors.py 35295526214877 --json
"""

from __future__ import annotations

import argparse
import json
import sys
import time
from dataclasses import dataclass, field

from databricks.sdk import WorkspaceClient
from databricks.sdk.service.jobs import RunResultState
from databricks.sdk.service.sql import QueryFilter, QueryStatus, TimeRange


# ---------------------------------------------------------------------------
# Data structures
# ---------------------------------------------------------------------------


@dataclass
class TaskError:
    """A task that directly failed."""

    task_key: str
    run_id: int
    state: str
    error_message: str
    run_page_url: str = ""


@dataclass
class SkippedTask:
    """A task skipped because an upstream task failed."""

    task_key: str
    run_id: int


@dataclass
class QueryError:
    """A failed SQL query from query history, linked to a task."""

    task_run_id: str
    task_key: str
    query_text: str
    error_message: str


@dataclass
class JobRunErrors:
    """All errors collected from a job run."""

    run_id: int
    job_id: int
    state: str
    result_state: str
    state_message: str
    task_errors: list[TaskError] = field(default_factory=list)
    skipped_tasks: list[SkippedTask] = field(default_factory=list)
    query_errors: list[QueryError] = field(default_factory=list)

    @property
    def has_errors(self) -> bool:
        return bool(self.task_errors or self.query_errors)

    def to_dict(self) -> dict[str, object]:
        return {
            "run_id": self.run_id,
            "job_id": self.job_id,
            "state": self.state,
            "result_state": self.result_state,
            "state_message": self.state_message,
            "task_errors": [
                {
                    "task_key": e.task_key,
                    "run_id": e.run_id,
                    "state": e.state,
                    "error_message": e.error_message,
                    "run_page_url": e.run_page_url,
                }
                for e in self.task_errors
            ],
            "skipped_tasks": [
                {"task_key": s.task_key, "run_id": s.run_id}
                for s in self.skipped_tasks
            ],
            "query_errors": [
                {
                    "task_run_id": e.task_run_id,
                    "task_key": e.task_key,
                    "query_text": e.query_text,
                    "error_message": e.error_message,
                }
                for e in self.query_errors
            ],
        }


# ---------------------------------------------------------------------------
# Core logic
# ---------------------------------------------------------------------------


def get_job_errors(
    run_id: int, profile: str = "{{ databricks_profile or 'DEFAULT' }}"
) -> JobRunErrors:
    """Fetch all errors for a given job run.

    Args:
        run_id: The Databricks **job** run ID (top-level, not a task run ID).
        profile: Databricks CLI profile name.

    Returns:
        A JobRunErrors object with task errors, skipped tasks, and query errors.
    """
    w = WorkspaceClient(profile=profile)
    run = w.jobs.get_run(run_id)

    state = run.state
    errors = JobRunErrors(
        run_id=run_id,
        job_id=run.job_id or 0,
        state=(
            state.life_cycle_state.value
            if state and state.life_cycle_state
            else "UNKNOWN"
        ),
        result_state=(
            state.result_state.value
            if state and state.result_state
            else "UNKNOWN"
        ),
        state_message=state.state_message or "" if state else "",
    )

    # Build a mapping from task run_id -> task_key for cross-referencing
    task_run_id_to_key: dict[str, str] = {}

    # --- Task-level errors + skipped tasks ---
    for task in run.tasks or []:
        task_state = task.state
        if not task_state or not task_state.result_state:
            continue

        if task.run_id:
            task_run_id_to_key[str(task.run_id)] = task.task_key or "unknown"

        if task_state.result_state in (
            RunResultState.FAILED,
            RunResultState.TIMEDOUT,
            RunResultState.CANCELED,
        ):
            # Fetch the task run URL for direct navigation
            page_url = ""
            if task.run_id:
                try:
                    task_run = w.jobs.get_run(task.run_id)
                    page_url = task_run.run_page_url or ""
                except Exception:  # noqa: BLE001
                    pass

            errors.task_errors.append(
                TaskError(
                    task_key=task.task_key or "unknown",
                    run_id=task.run_id or 0,
                    state=task_state.result_state.value,
                    error_message=(
                        task_state.state_message
                        or "No error message available"
                    ),
                    run_page_url=page_url,
                )
            )
        elif task_state.result_state == RunResultState.UPSTREAM_FAILED:
            errors.skipped_tasks.append(
                SkippedTask(
                    task_key=task.task_key or "unknown",
                    run_id=task.run_id or 0,
                )
            )

    # --- SQL query errors from Query History ---
    task_run_ids = set(task_run_id_to_key.keys())

    start_time = run.start_time or 0
    end_time = run.end_time or int(time.time() * 1000)

    if start_time and task_run_ids:
        try:
            queries = w.query_history.list(
                filter_by=QueryFilter(
                    statuses=[QueryStatus.FAILED],
                    query_start_time_range=TimeRange(
                        start_time_ms=start_time,
                        end_time_ms=end_time,
                    ),
                ),
                max_results=50,
            )

            for q in queries.res or []:
                if q.query_source and q.query_source.job_info:
                    tid = q.query_source.job_info.job_task_run_id
                    if tid and tid in task_run_ids:
                        errors.query_errors.append(
                            QueryError(
                                task_run_id=tid,
                                task_key=task_run_id_to_key.get(
                                    tid, "unknown"
                                ),
                                query_text=q.query_text or "",
                                error_message=(
                                    q.error_message or "No error details"
                                ),
                            )
                        )
        except Exception as exc:  # noqa: BLE001
            print(
                f"Warning: Could not fetch query history: {exc}",
                file=sys.stderr,
            )

    return errors


# ---------------------------------------------------------------------------
# Display helpers
# ---------------------------------------------------------------------------


def print_errors(errors: JobRunErrors) -> None:
    """Print errors in a human-readable format."""
    print(f"Job Run: {errors.run_id}  |  Job ID: {errors.job_id}")
    print(f"State:   {errors.state}  |  Result: {errors.result_state}")
    if errors.state_message:
        print(f"Message: {errors.state_message}")
    print("=" * 70)

    # Build set of task run IDs that have query errors
    tasks_with_query_errors = {qe.task_run_id for qe in errors.query_errors}

    # --- Direct task failures ---
    if errors.task_errors:
        print(f"\nFailed Tasks ({len(errors.task_errors)}):")
        print("-" * 70)
        for te in errors.task_errors:
            print(f"  Task:    {te.task_key}")
            print(f"  Run ID:  {te.run_id}")
            print(f"  State:   {te.state}")
            print(f"  Error:   {te.error_message}")
            if te.run_page_url:
                print(f"  URL:     {te.run_page_url}")
            if str(te.run_id) not in tasks_with_query_errors:
                print(
                    "  Note:    No SQL errors in query history."
                    " Check notebook output at the URL above."
                )
            print()

    # --- Query errors (the real tracebacks) ---
    if errors.query_errors:
        print(f"SQL / Notebook Errors ({len(errors.query_errors)}):")
        print("-" * 70)
        for qe in errors.query_errors:
            print(f"  Task:    {qe.task_key} (run_id: {qe.task_run_id})")
            print(f"  Query:   {qe.query_text[:300]}")
            print(f"  Error:   {qe.error_message}")
            print()

    # --- Skipped (downstream cascade) ---
    if errors.skipped_tasks:
        print(
            f"Skipped Tasks — upstream failed ({len(errors.skipped_tasks)}):"
        )
        print("-" * 70)
        for st in errors.skipped_tasks:
            print(f"  - {st.task_key}")
        print()

    if not errors.has_errors:
        print("\nNo errors found for this run.")
        print(
            "The run may still be in progress, or errors may not be"
            " in query history."
        )


# ---------------------------------------------------------------------------
# CLI
# ---------------------------------------------------------------------------


def main() -> None:
    """CLI entry point."""
    parser = argparse.ArgumentParser(
        description=(
            "Fetch error details for a failed Databricks job run. "
            "Pass the JOB run ID (top-level), not a task run ID."
        ),
    )
    parser.add_argument(
        "run_id",
        type=int,
        help="The job run ID (top-level run, not a task run ID)",
    )
    parser.add_argument(
        "--profile",
        default="{{ databricks_profile or 'DEFAULT' }}",
        help=(
            "Databricks CLI profile"
            " (default: {{ databricks_profile or 'DEFAULT' }})"
        ),
    )
    parser.add_argument(
        "--json",
        action="store_true",
        dest="as_json",
        help="Output errors as JSON (useful for programmatic consumption)",
    )
    args = parser.parse_args()

    errors = get_job_errors(args.run_id, profile=args.profile)

    if args.as_json:
        print(json.dumps(errors.to_dict(), indent=2))
    else:
        print_errors(errors)

    sys.exit(1 if errors.has_errors else 0)


if __name__ == "__main__":
    main()
