from __future__ import annotations

from dataclasses import asdict, dataclass
from datetime import datetime
from decimal import Decimal
from typing import Any
import json


@dataclass(frozen=True)
class JobEvents:
    """Job event stream page (e.g. progress/log events)."""

    events: list[dict[str, Any]]
    continuation_token: str | None


def _to_int(v: Any) -> int | None:
    if v is None:
        return None
    if isinstance(v, bool):
        return int(v)
    if isinstance(v, int):
        return v
    if isinstance(v, Decimal):
        return int(v)
    # Snowpark may return numpy scalars depending on environment; fall back to int(...)
    try:
        return int(v)
    except Exception:
        return None


@dataclass(frozen=True)
class JobRecord:
    """Backend job/transaction record as returned by SQL or Direct Access.

    This is an internal/backend-facing model (rows/payloads). The stable public
    model returned by `client.jobs.get/list` lives in `views.py`.
    """

    id: str
    reasoner_type: str
    # Transaction-specific fields (LOGIC transactions), also populated when available for jobs.
    database_name: str | None = None
    read_only: bool | None = None
    state: str | None = None
    created_by: str | None = None
    created_on: datetime | None = None
    finished_at: datetime | None = None
    duration: int | None = None
    payload: str | None = None
    reasoner_name: str | None = None
    abort_reason: str | None = None
    timeout_ms: int | None = None


def job_from_row(row: Any) -> JobRecord:
    """Parse a Snowpark row (or row-like dict) into a `JobRecord`."""
    d = row.asDict() if hasattr(row, "asDict") else dict(row)
    return JobRecord(
        id=str(d.get("ID") or d.get("id") or ""),
        reasoner_type=str(d.get("TYPE") or d.get("type") or ""),
        database_name=(d.get("DATABASE_NAME") or d.get("database_name") or d.get("DBNAME") or d.get("dbname")),
        read_only=(
            d.get("READ_ONLY")
            if "READ_ONLY" in d
            else d.get("read_only") if "read_only" in d else d.get("READONLY") if "READONLY" in d else d.get("readonly")
        ),
        state=(d.get("STATE") or d.get("state")),
        created_by=(d.get("CREATED_BY") or d.get("created_by")),
        created_on=(d.get("CREATED_ON") or d.get("created_on")),
        finished_at=(d.get("FINISHED_AT") or d.get("finished_at")),
        duration=_to_int(d.get("DURATION") if "DURATION" in d else d.get("duration")),
        payload=(d.get("PAYLOAD") or d.get("payload")),
        reasoner_name=(d.get("ENGINE_NAME") or d.get("engine_name") or d.get("WORKER_NAME") or d.get("worker_name")),
        abort_reason=(d.get("ABORT_REASON") or d.get("abort_reason")),
        timeout_ms=_to_int(d.get("TIMEOUT_MS") if "TIMEOUT_MS" in d else d.get("timeout_ms")),
    )


def logic_txn_job_from_row(row: Any) -> JobRecord:
    """Parse a transaction row/object into a JobRecord with reasoner_type='LOGIC'."""
    d = row.asDict() if hasattr(row, "asDict") else dict(row)
    # Ensure expected fields exist for `job_from_row`
    d2 = dict(d)
    d2.setdefault("type", "LOGIC")
    # Transactions use engine_name; normalize to the jobs parser expectations.
    if "engine_name" in d2 and "ENGINE_NAME" not in d2:
        d2["ENGINE_NAME"] = d2.get("engine_name")
    if "ENGINE_NAME" in d2 and "engine_name" not in d2:
        d2["engine_name"] = d2.get("ENGINE_NAME")
    return job_from_row(d2)


def job_events_from_value(value: Any) -> JobEvents:
    """Parse a job-events payload into a `JobEvents` object.

    Snowflake UDFs often return JSON as a string; Direct Access may return dicts.
    """
    if value is None or value == "":
        return JobEvents(events=[], continuation_token=None)

    payload: Any = value
    if isinstance(value, str):
        try:
            payload = json.loads(value)
        except Exception:
            return JobEvents(events=[], continuation_token=None)

    if not isinstance(payload, dict):
        return JobEvents(events=[], continuation_token=None)

    events = payload.get("events", [])
    token = payload.get("continuation_token", None)
    if not isinstance(events, list):
        events = []
    if token is not None and not isinstance(token, str):
        token = str(token)

    # Ensure events are dicts (best-effort; don't throw on unexpected shapes).
    out_events: list[dict[str, Any]] = []
    for e in events:
        if isinstance(e, dict):
            out_events.append(e)
    return JobEvents(events=out_events, continuation_token=token)


class OperationStatus:
    """Normalized operation status values + helpers (jobs)."""

    PENDING = "PENDING"
    RUNNING = "RUNNING"
    SUCCEEDED = "SUCCEEDED"
    FAILED = "FAILED"
    CANCELLED = "CANCELLED"

    _TERMINAL = {SUCCEEDED, FAILED, CANCELLED}

    @classmethod
    def done(cls, status: str | None) -> bool:
        return (status or "").upper() in cls._TERMINAL


@dataclass(frozen=True)
class Operation:
    """Long-running operation handle derived from a job execution."""

    id: str
    status: str | None = None
    error: str | None = None
    raw: dict[str, Any] | None = None

    def done(self) -> bool:
        return OperationStatus.done(self.status)


# ---------------------------------------------------------------------------
# Public view models (stable SDK return types).
# ---------------------------------------------------------------------------


@dataclass(frozen=True)
class Job:
    """Normalized public model for `client.jobs.get/list`.

    By default, `Job` contains ONLY fields that are common between the Jobs and
    Transactions backends.

    When `include=` is used, the SDK returns a `JobExpanded` instance (a subclass)
    that adds additional, opt-in fields.
    """

    id: str
    reasoner_type: str
    state: str | None = None
    name: str | None = None
    created_by: str | None = None
    created_on: datetime | None = None
    finished_at: datetime | None = None
    duration: int | None = None
    abort_reason: str | None = None


@dataclass(frozen=True)
class JobExpanded(Job):
    """Expanded job returned only when `include=` is provided."""

    details: dict[str, Any] | None = None
    events: JobEvents | None = None
    artifacts: dict[str, Any] | None = None
    problems: dict[str, Any] | None = None
    raw: dict[str, Any] | None = None


@dataclass(frozen=True)
class JobInclude:
    """Optional include fields for a `Job`.

    Instances should only be created when `include=` was requested.
    """

    details: dict[str, Any] | None = None
    events: JobEvents | None = None
    artifacts: dict[str, Any] | None = None
    problems: dict[str, Any] | None = None
    raw: dict[str, Any] | None = None

    def is_empty(self) -> bool:
        return all(v is None for v in (self.details, self.events, self.artifacts, self.problems, self.raw))


def record_to_job(
    job: JobRecord,
    *,
    include: JobInclude | None = None,
) -> Job:
    """Project a backend `JobRecord` into a public `Job` (optionally included)."""
    if include is None or include.is_empty():
        return Job(
            id=str(job.id),
            reasoner_type=str(job.reasoner_type),
            state=job.state,
            name=job.reasoner_name,
            created_by=job.created_by,
            created_on=job.created_on,
            finished_at=job.finished_at,
            duration=job.duration,
            abort_reason=job.abort_reason,
        )
    return JobExpanded(
        id=str(job.id),
        reasoner_type=str(job.reasoner_type),
        state=job.state,
        name=job.reasoner_name,
        created_by=job.created_by,
        created_on=job.created_on,
        finished_at=job.finished_at,
        duration=job.duration,
        abort_reason=job.abort_reason,
        details=include.details,
        events=include.events,
        artifacts=include.artifacts,
        problems=include.problems,
        raw=include.raw,
    )


def record_to_raw_dict(job: JobRecord) -> dict[str, Any]:
    """Serialize a backend `JobRecord` dataclass to a dict (best-effort)."""
    # `asdict` recursively converts dataclasses; datetime stays as datetime which is
    # fine for Python callers. JSON serialization (if needed) is the caller's job.
    return asdict(job)


def record_to_job_with_include(
    job: JobRecord,
    *,
    include: set[str],
    events: JobEvents | None = None,
    artifacts: dict[str, Any] | None = None,
    problems: dict[str, Any] | None = None,
) -> Job:
    """Build a public `Job` from a backend record plus optionally-fetched include fields."""
    details: dict[str, Any] | None = None
    if "details" in include:
        details = {
            "database_name": job.database_name,
            "read_only": job.read_only,
            "payload": job.payload,
            "timeout_ms": job.timeout_ms,
        }

    raw: dict[str, Any] | None = None
    if "raw" in include:
        raw = record_to_raw_dict(job)

    inc: JobInclude | None = None
    if include:
        inc = JobInclude(
            details=details,
            events=events if "events" in include else None,
            artifacts=artifacts if "artifacts" in include else None,
            problems=problems if "problems" in include else None,
            raw=raw,
        )

    return record_to_job(job, include=inc)


__all__ = [
    "JobRecord",
    "JobEvents",
    "Job",
    "JobExpanded",
    "JobInclude",
    "Operation",
    "OperationStatus",
    "job_from_row",
    "logic_txn_job_from_row",
    "job_events_from_value",
    "record_to_job",
    "record_to_job_with_include",
    "record_to_raw_dict",
]

