"""
Generic SSE task monitoring for background operations.

This module provides a reusable SSE streaming interface for monitoring
any long-running background task (ingestion, backup, restore, etc.).
"""

import asyncio
import json
import time
from collections.abc import AsyncGenerator
from datetime import UTC, datetime
from typing import Any

from robosystems.graph_api.models.tasks import TaskType
from robosystems.logger import logger


async def generate_task_sse_events(
  task_manager,
  task_id: str,
  task_type: TaskType = TaskType.INGESTION,
  heartbeat_interval: int = 30,
) -> AsyncGenerator[dict[str, Any]]:
  """
  Generate SSE events for any background task with progress monitoring.

  This is a generic implementation that can be used for:
  - Database backups
  - Database restores
  - Any other long-running task

  Args:
      task_manager: Task manager instance with get_task method
      task_id: Unique task identifier
      task_type: Type of task being monitored
      heartbeat_interval: Seconds between heartbeat events

  Yields:
      SSE event dictionaries with event type and data
  """
  last_heartbeat = time.time()
  last_progress = -1

  # Send initial connection event
  yield {
    "event": "connected",
    "data": json.dumps(
      {
        "task_id": task_id,
        "task_type": task_type.value,
        "timestamp": datetime.now(UTC).isoformat(),
        "message": f"Connected to {task_type.value} task monitor",
      }
    ),
  }

  while True:
    try:
      # Get current task status
      task = await task_manager.get_task(task_id)

      if not task:
        yield {
          "event": "error",
          "data": json.dumps(
            {"error": f"Task {task_id} not found", "task_type": task_type.value}
          ),
        }
        break

      # Send heartbeat every interval to prevent timeout
      current_time = time.time()
      if current_time - last_heartbeat > heartbeat_interval:
        yield {
          "event": "heartbeat",
          "data": json.dumps(
            {
              "task_id": task_id,
              "task_type": task_type.value,
              "status": task["status"],
              "timestamp": datetime.now(UTC).isoformat(),
              "message": "Task is still running...",
            }
          ),
        }
        last_heartbeat = current_time
        logger.debug(f"[SSE] Sent heartbeat for {task_type.value} task {task_id}")

      # Send progress updates
      current_progress = task.get("progress_percent", 0)
      if current_progress != last_progress:
        yield {
          "event": "progress",
          "data": json.dumps(
            {
              "task_id": task_id,
              "task_type": task_type.value,
              "status": task["status"],
              "progress_percent": current_progress,
              "records_processed": task.get("records_processed", 0),
              "estimated_records": task.get("estimated_records", 0),
              "started_at": task.get("started_at"),
              "message": _get_progress_message(task_type, task),
              "metadata": task.get("metadata", {}),
            }
          ),
        }
        last_progress = current_progress

      # Check for completion
      if task["status"] == "completed":
        yield {
          "event": "completed",
          "data": json.dumps(
            {
              "task_id": task_id,
              "task_type": task_type.value,
              "status": "completed",
              "result": task.get("result"),
              "duration_seconds": _calculate_duration(task),
              "message": _get_completion_message(task_type, task),
              "metadata": task.get("metadata", {}),
            }
          ),
        }
        break

      # Check for failure
      if task["status"] == "failed":
        yield {
          "event": "failed",
          "data": json.dumps(
            {
              "task_id": task_id,
              "task_type": task_type.value,
              "status": "failed",
              "error": task.get("error"),
              "message": _get_failure_message(task_type, task),
              "metadata": task.get("metadata", {}),
            }
          ),
        }
        break

      # Wait a bit before checking again
      await asyncio.sleep(2)

    except Exception as e:
      logger.error(
        f"[SSE] Error generating events for {task_type.value} task {task_id}: {e}"
      )
      yield {
        "event": "error",
        "data": json.dumps({"error": str(e), "task_type": task_type.value}),
      }
      break


def _get_progress_message(task_type: TaskType, task: dict[str, Any]) -> str:
  """Generate task-specific progress message."""
  if task_type == TaskType.INGESTION:
    table_name = task.get("metadata", {}).get("table_name", "table")
    return f"Processing {table_name}..."
  elif task_type == TaskType.BACKUP:
    database = task.get("metadata", {}).get("database", "database")
    return f"Backing up {database}..."
  elif task_type == TaskType.RESTORE:
    database = task.get("metadata", {}).get("database", "database")
    return f"Restoring {database}..."
  elif task_type == TaskType.STAGING:
    table_name = task.get("metadata", {}).get("table_name", "table")
    file_count = task.get("file_count", 0)
    if file_count > 0:
      return f"Staging {table_name} from {file_count} files..."
    return f"Staging {table_name}..."
  elif task_type == TaskType.MIGRATION:
    metadata = task.get("metadata", {})
    operation = metadata.get("operation", "migration")
    current_db = metadata.get("current_database", "")
    completed = metadata.get("databases_completed", 0)
    total = metadata.get("databases_total", 0)
    if current_db and total > 0:
      return f"Migration {operation}: {current_db} ({completed}/{total})"
    return f"Migration {operation} in progress..."
  else:
    return f"Processing {task_type.value} task..."


def _get_completion_message(task_type: TaskType, task: dict[str, Any]) -> str:
  """Generate task-specific completion message."""
  metadata = task.get("metadata", {})
  result = task.get("result", {})

  if task_type == TaskType.INGESTION:
    table_name = metadata.get("table_name", "table")
    records = result.get("records_loaded", 0)
    if records > 0:
      return f"Successfully ingested {records:,} records into {table_name}"
    else:
      return f"Successfully completed ingestion for {table_name}"
  elif task_type == TaskType.BACKUP:
    database = metadata.get("database", "database")
    size_mb = result.get("backup_size_mb", 0)
    if size_mb > 0:
      return f"Successfully backed up {database} ({size_mb:.1f} MB)"
    else:
      return f"Successfully backed up {database}"
  elif task_type == TaskType.RESTORE:
    database = metadata.get("database", "database")
    return f"Successfully restored {database}"
  elif task_type == TaskType.STAGING:
    table_name = metadata.get("table_name", "table")
    duration = result.get("duration_seconds", 0)
    if duration > 0:
      return f"Successfully staged {table_name} in {duration:.1f}s"
    return f"Successfully staged {table_name}"
  elif task_type == TaskType.MIGRATION:
    operation = metadata.get("operation", "migration")
    count = result.get("databases_exported") or result.get("databases_imported", 0)
    return f"Migration {operation} complete: {count} databases"
  else:
    return f"Successfully completed {task_type.value} task"


def _get_failure_message(task_type: TaskType, task: dict[str, Any]) -> str:
  """Generate task-specific failure message."""
  metadata = task.get("metadata", {})

  if task_type == TaskType.INGESTION:
    table_name = metadata.get("table_name", "table")
    return f"Failed to ingest data into {table_name}"
  elif task_type == TaskType.BACKUP:
    database = metadata.get("database", "database")
    return f"Failed to backup {database}"
  elif task_type == TaskType.RESTORE:
    database = metadata.get("database", "database")
    return f"Failed to restore {database}"
  elif task_type == TaskType.STAGING:
    table_name = metadata.get("table_name", "table")
    return f"Failed to stage {table_name}"
  elif task_type == TaskType.MIGRATION:
    operation = metadata.get("operation", "migration")
    return f"Migration {operation} failed"
  else:
    return f"Failed to complete {task_type.value} task"


def _calculate_duration(task: dict[str, Any]) -> float:
  """Calculate task duration in seconds."""
  if task.get("completed_at") and task.get("started_at"):
    try:
      completed = datetime.fromisoformat(task["completed_at"])
      started = datetime.fromisoformat(task["started_at"])
      return (completed - started).total_seconds()
    except (ValueError, TypeError):
      pass
  return 0.0
