from __future__ import annotations

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

from relationalai.services._shared.sql_rows import row_to_mapping


@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


def coerce_datetime(value: Any) -> datetime | None:
    """Best-effort datetime coercion across SQL/DA timestamp shapes.

    Supported inputs include:
    - `datetime`
    - SQL-style strings like `2026-03-06 13:39:54`
    - ISO 8601 strings, including trailing `Z`
    - Unix epoch seconds / milliseconds as `int`, `float`, `Decimal`, or numeric strings
    """
    if value is None:
        return None
    if isinstance(value, datetime):
        return value
    if isinstance(value, bool):
        return None
    if isinstance(value, Decimal):
        value = int(value)
    if isinstance(value, (int, float)):
        raw = float(value)
        candidates: list[datetime] = []
        raw_candidates = [raw / 1000.0, raw] if abs(raw) >= 1e10 else [raw]
        for candidate in raw_candidates:
            try:
                candidates.append(datetime.fromtimestamp(candidate, tz=timezone.utc))
            except Exception:
                continue
        if not candidates:
            return None
        for candidate in candidates:
            if 1900 <= candidate.year <= 3000:
                return candidate
        return candidates[0]
    if isinstance(value, str):
        s = value.strip()
        if not s:
            return None
        if s.isdigit():
            return coerce_datetime(int(s))
        try:
            return datetime.fromisoformat(s.replace("Z", "+00:00"))
        except Exception:
            return None
    return None


def datetime_sort_key(value: Any) -> float:
    """Return a stable comparable sort key for mixed datetime formats.

    Jobs data can come from multiple backends with mixed timezone conventions.
    For ordering purposes, coerce supported timestamp shapes first, then treat
    naive datetimes as UTC and normalize aware ones to UTC before comparing.
    """
    dt = coerce_datetime(value)
    if not isinstance(dt, datetime):
        return float("-inf")
    dt = dt.replace(tzinfo=timezone.utc) if dt.tzinfo is None else dt.astimezone(timezone.utc)
    return dt.timestamp()


@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_to_mapping(row)
    return JobRecord(
        id=str(d.get("ID") or d.get("id") or ""),
        reasoner_type=str(d.get("REASONER_TYPE") or d.get("reasoner_type") or 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=coerce_datetime(d.get("CREATED_ON") if "CREATED_ON" in d else d.get("created_on")),
        finished_at=coerce_datetime(d.get("FINISHED_AT") if "FINISHED_AT" in d else 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("REASONER_NAME")
            or d.get("reasoner_name")
            or 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_to_mapping(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",
    "datetime_sort_key",
    "job_from_row",
    "logic_txn_job_from_row",
    "job_events_from_value",
    "record_to_job",
    "record_to_job_with_include",
    "record_to_raw_dict",
]

