"""
Timestamp collection state management for orchestrator initial mode.

Timestamps are stored in a CogniteFile DMS node (cdf_cdm:CogniteFile:v1)
referenced via a DirectRelation on the OrchestrationState node. This avoids
the CDF DMS 40KB JSON property limit that makes it impossible to store large
timestamp lists directly on the state node.

State stored on OrchestrationState (DMS):
- timestamp_collection_phase: 'collecting' | 'complete' | null
- timestamp_collection_count: Number of timestamps (for visibility)
- timestamp_collection_started: When collection started
- last_updated: Used for concurrency/staleness detection
- timestamps_file: DirectRelation → CogniteFile node

State stored in the CogniteFile (file content, JSON):
- timestamps: list[int]  — milliseconds since epoch
- cursor: str | null     — DMS sync cursor for resuming
"""

from __future__ import annotations

import json
from typing import TYPE_CHECKING

if TYPE_CHECKING:
    from cognite.client import CogniteClient


def _file_ext_id(orchestration_id: str) -> str:
    return f"ts_collect_{orchestration_id}"


def load_timestamp_collection_state(
    client: CogniteClient, orchestration_id: str, data_quality_space: str = "dataQuality"
) -> dict | None:
    """
    Load saved timestamp collection state for this orchestration.

    Returns None if no collecting state exists (first run or already complete).
    Returns dict with keys: cursor, timestamps, count, last_updated_ms
    """
    try:
        from handlers.orchestrator import load_orchestration_state

        state = load_orchestration_state(client, orchestration_id, data_quality_space)
        if not state:
            return None

        phase = state.get("timestamp_collection_phase")
        if phase not in ("collecting", "checkpoint"):
            return None

        count = state.get("timestamp_collection_count", 0)
        last_updated = state.get("last_updated")
        last_updated_ms = 0
        if last_updated:
            if hasattr(last_updated, "timestamp"):
                last_updated_ms = int(last_updated.timestamp() * 1000)
            elif isinstance(last_updated, str):
                try:
                    import datetime as _dt

                    lu_dt = _dt.datetime.fromisoformat(last_updated.replace("Z", "+00:00"))
                    last_updated_ms = int(lu_dt.timestamp() * 1000)
                except Exception:
                    pass

        # Load timestamps + cursor from the CogniteFile
        timestamps: list[int] = []
        cursor: str | None = None
        file_ref = state.get("timestamps_file")
        if file_ref:
            try:
                from cognite.client.data_classes.data_modeling import NodeId

                file_space = file_ref.get("space", data_quality_space)
                file_xid = file_ref.get("externalId", _file_ext_id(orchestration_id))
                file_bytes = client.files.download_bytes(instance_id=NodeId(file_space, file_xid))
                data = json.loads(file_bytes)
                timestamps = data.get("timestamps", [])
                cursor = data.get("cursor")
                print(f"  Loaded {len(timestamps):,} timestamps from CogniteFile ({file_xid})")
                print(f"    Cursor: {str(cursor)[:50] if cursor else 'None'}...")
            except Exception as e:
                print(f"  ⚠ Could not load CogniteFile: {e} — will restart from scratch")
        else:
            print("  No timestamps_file relation found — will restart from scratch")

        print("  Loaded timestamp collection state:")
        print(f"    Timestamps: {len(timestamps):,}  Count in DMS: {count:,}")
        print(f"    Last updated: {last_updated}")

        return {
            "cursor": cursor,
            "timestamps": timestamps,
            "count": count,
            "last_updated_ms": last_updated_ms,
        }

    except Exception as e:
        print(f"Failed to load timestamp collection state: {e}")
        import traceback

        traceback.print_exc()
        return None


def save_timestamp_collection_state(
    client: CogniteClient,
    orchestration_id: str,
    cursor: str | None,
    timestamps: list[int],
    data_quality_space: str = "dataQuality",
    phase: str = "collecting",
):
    """
    Save timestamp collection progress.

    phase: 'collecting' (mid-run intermediate save) or 'checkpoint' (controlled
           shutdown — function exited cleanly, monitor should re-trigger immediately).

    1. Creates/updates a CogniteFile node and uploads timestamps + cursor as JSON.
    2. Saves lightweight DMS state (phase, count, last_updated, relation to file).

    Raises RuntimeError if the file upload fails — resumption depends on it.
    """
    try:
        import datetime

        from cognite.client.data_classes.data_modeling import NodeId
        from cognite.client.data_classes.data_modeling.cdm.v1 import CogniteFileApply
        from handlers.orchestrator import load_orchestration_state, save_orchestration_state

        file_xid = _file_ext_id(orchestration_id)

        # --- 1. Upload timestamps + cursor to CogniteFile ---
        payload = json.dumps({"timestamps": timestamps, "cursor": cursor}).encode()

        # Ensure the CogniteFile node exists
        client.data_modeling.instances.apply(
            CogniteFileApply(
                name=f"timestamp_collection_{orchestration_id}.json",
                space=data_quality_space,
                external_id=file_xid,
            )
        )

        # Upload file content
        client.files.upload_content_bytes(
            payload,
            instance_id=NodeId(data_quality_space, file_xid),
        )

        print("  Saving timestamp collection state...")
        print(f"    Timestamps: {len(timestamps):,} → CogniteFile {file_xid}")
        print(f"    Cursor: {str(cursor)[:50] if cursor else 'None'}...")

        # --- 2. Save lightweight DMS state ---
        state = load_orchestration_state(client, orchestration_id, data_quality_space)
        if not state:
            state = {
                "orchestration_id": orchestration_id,
                "timestamp_collection_started": datetime.datetime.now(tz=datetime.timezone.utc),
            }

        state["timestamp_collection_count"] = len(timestamps)
        state["timestamp_collection_phase"] = phase
        state["timestamps_file"] = {"space": data_quality_space, "externalId": file_xid}
        state["last_updated"] = datetime.datetime.now(tz=datetime.timezone.utc)

        if "timestamp_collection_started" not in state:
            state["timestamp_collection_started"] = state["last_updated"]

        save_orchestration_state(client, state, data_quality_space)
        print("  ✓ Timestamp collection state saved")

    except Exception as e:
        error_msg = f"CRITICAL ERROR saving timestamp collection state: {e}"
        print(error_msg)
        print(f"  Error type: {type(e).__name__}")
        import traceback

        traceback.print_exc()
        raise RuntimeError(error_msg) from e


def mark_timestamp_collection_complete(
    client: CogniteClient,
    orchestration_id: str,
    timestamps: list[int],
    data_quality_space: str = "dataQuality",
):
    """
    Mark timestamp collection as complete and delete the CogniteFile.

    Timestamps are already in memory at this point — the file is no longer needed.
    """
    try:
        import datetime

        from cognite.client.data_classes.data_modeling import NodeId
        from handlers.orchestrator import load_orchestration_state, save_orchestration_state

        state = load_orchestration_state(client, orchestration_id, data_quality_space)
        if not state:
            print("  ⚠ Warning: No orchestration state found when marking collection complete")
            return

        # Mark complete
        state["timestamp_collection_phase"] = "complete"
        state["timestamp_collection_count"] = len(timestamps)
        state["last_updated"] = datetime.datetime.now(tz=datetime.timezone.utc)

        # Delete the CogniteFile node (and its file content)
        file_ref = state.get("timestamps_file")
        if file_ref:
            try:
                file_space = file_ref.get("space", data_quality_space)
                file_xid = file_ref.get("externalId", _file_ext_id(orchestration_id))
                client.data_modeling.instances.delete(nodes=[NodeId(file_space, file_xid)])
                print(f"  ✓ Deleted CogniteFile {file_xid}")
            except Exception as del_err:
                print(f"  ⚠ Could not delete CogniteFile: {del_err}")

        state["timestamps_file"] = None

        print("  Marking timestamp collection complete...")
        print(f"    Total timestamps: {len(timestamps):,}")

        save_orchestration_state(client, state, data_quality_space)
        print("  ✓ Timestamp collection marked as complete")

    except Exception as e:
        print(f"Warning: Failed to mark timestamp collection complete: {e}")
        # Don't raise — collection is done, this is cleanup only
