"""
Scheduler handler for scheduler management functionality.

This module handles all scheduler-related operations including status checking,
waiting for idle state, and streaming progress updates.
"""

import json
import time
import traceback

from collections import Counter
from datetime import datetime, timezone
from typing import Any

from fastapi import HTTPException
from fastapi.responses import StreamingResponse

# Imports for new implementation
from memos.api.product_models import (
    AllStatusResponse,
    AllStatusResponseData,
    StatusResponse,
    StatusResponseItem,
    TaskQueueData,
    TaskQueueResponse,
    TaskSummary,
)
from memos.log import get_logger
from memos.mem_scheduler.base_scheduler import BaseScheduler
from memos.mem_scheduler.optimized_scheduler import OptimizedScheduler
from memos.mem_scheduler.utils.status_tracker import TaskStatusTracker


logger = get_logger(__name__)


def handle_scheduler_allstatus(
    mem_scheduler: BaseScheduler,
    status_tracker: TaskStatusTracker,
) -> AllStatusResponse:
    """
    Get aggregated scheduler status metrics (no per-task payload).

    Args:
        mem_scheduler: The BaseScheduler instance.
        status_tracker: The TaskStatusTracker instance.

    Returns:
        AllStatusResponse with aggregated status data.
    """

    def _summarize_tasks(task_details: list[dict[str, Any]]) -> TaskSummary:
        """Aggregate counts by status for the provided task details (tracker data)."""
        counter = Counter()
        for detail in task_details:
            status = detail.get("status")
            if status:
                counter[status] += 1

        total = sum(counter.values())
        return TaskSummary(
            waiting=counter.get("waiting", 0),
            in_progress=counter.get("in_progress", 0),
            completed=counter.get("completed", 0),
            pending=counter.get("pending", counter.get("waiting", 0)),
            failed=counter.get("failed", 0),
            cancelled=counter.get("cancelled", 0),
            total=total,
        )

    def _aggregate_counts_from_redis(
        tracker: TaskStatusTracker, max_age_seconds: float = 86400
    ) -> TaskSummary | None:
        """Stream status counts directly from Redis to avoid loading all task payloads."""
        redis_client = getattr(tracker, "redis", None)
        if not redis_client:
            return None

        counter = Counter()
        now = datetime.now(timezone.utc).timestamp()

        # Scan task_meta keys, then hscan each hash in batches
        cursor: int | str = 0
        while True:
            cursor, keys = redis_client.scan(cursor=cursor, match="memos:task_meta:*", count=200)
            for key in keys:
                h_cursor: int | str = 0
                while True:
                    h_cursor, fields = redis_client.hscan(key, cursor=h_cursor, count=500)
                    for value in fields.values():
                        try:
                            payload = json.loads(
                                value.decode("utf-8") if isinstance(value, bytes) else value
                            )
                            # Skip stale entries to reduce noise and load
                            ts = payload.get("submitted_at") or payload.get("started_at")
                            if ts:
                                try:
                                    ts_dt = datetime.fromisoformat(ts)
                                    ts_seconds = ts_dt.timestamp()
                                except Exception:
                                    ts_seconds = None
                                if ts_seconds and (now - ts_seconds) > max_age_seconds:
                                    continue
                            status = payload.get("status")
                            if status:
                                counter[status] += 1
                        except Exception:
                            continue
                    if h_cursor == 0 or h_cursor == "0":
                        break
            if cursor == 0 or cursor == "0":
                break

        if not counter:
            return TaskSummary()  # Empty summary if nothing found

        total = sum(counter.values())
        return TaskSummary(
            waiting=counter.get("waiting", 0),
            in_progress=counter.get("in_progress", 0),
            completed=counter.get("completed", 0),
            pending=counter.get("pending", counter.get("waiting", 0)),
            failed=counter.get("failed", 0),
            cancelled=counter.get("cancelled", 0),
            total=total,
        )

    try:
        # Prefer streaming aggregation to avoid pulling all task payloads
        all_tasks_summary = _aggregate_counts_from_redis(status_tracker)
        if all_tasks_summary is None:
            # Fallback: load all details then aggregate
            global_tasks = status_tracker.get_all_tasks_global()
            all_task_details: list[dict[str, Any]] = []
            for _, tasks in global_tasks.items():
                all_task_details.extend(tasks.values())
            all_tasks_summary = _summarize_tasks(all_task_details)

        # Scheduler view: assume tracker contains scheduler tasks; overlay queue monitor for live queue depth
        sched_waiting = all_tasks_summary.waiting
        sched_in_progress = all_tasks_summary.in_progress
        sched_pending = all_tasks_summary.pending
        sched_completed = all_tasks_summary.completed
        sched_failed = all_tasks_summary.failed
        sched_cancelled = all_tasks_summary.cancelled

        # If queue monitor is available, prefer its live waiting/in_progress counts.
        #
        # Two queue backends produce different ``queue_status_data`` shapes:
        #   - Redis queue: ``{"running": .., "remaining": .., "pending": ..,
        #     "<stream_key_prefix>:<user>:<cube>:<label>": {"running": ..,
        #     "remaining": .., "pending": ..}, ...}`` — the per-stream entries
        #     are nested dicts and their keys start with the scheduler stream
        #     prefix (default ``"scheduler:messages:stream:v2.0:"``). The loop
        #     below aggregates those per-stream entries.
        #   - Local in-memory queue: only the top-level totals
        #     ``{"running", "remaining", "pending"}`` are present (see
        #     ``TaskScheduleMonitor._get_local_tasks_status``). There are no
        #     per-stream keys to enumerate, so the prefix loop would silently
        #     collapse the summary to all zeros. Detect that case and fall back
        #     to the flat totals so Docker / no-Redis deployments still surface
        #     accurate in_progress / waiting / pending counts (issue #1395).
        if mem_scheduler.task_schedule_monitor:
            queue_status_data = mem_scheduler.task_schedule_monitor.get_tasks_status() or {}
            scheduler_waiting = 0
            scheduler_in_progress = 0
            scheduler_pending = 0
            per_stream_keys = [
                key
                for key, value in queue_status_data.items()
                if isinstance(key, str) and key.startswith("scheduler:") and isinstance(value, dict)
            ]

            if per_stream_keys:
                for key in per_stream_keys:
                    value = queue_status_data[key]
                    scheduler_in_progress += int(value.get("running", 0) or 0)
                    scheduler_pending += int(value.get("pending", value.get("remaining", 0)) or 0)
                    scheduler_waiting += int(value.get("remaining", 0) or 0)
            elif not getattr(mem_scheduler, "use_redis_queue", True):
                # Local-queue mode: there are no per-stream keys; use the flat
                # totals from ``_get_local_tasks_status``. ``remaining`` reflects
                # queued depth (waiting / pending) and ``running`` reflects the
                # dispatcher's in-flight count.
                scheduler_in_progress = int(queue_status_data.get("running", 0) or 0)
                scheduler_waiting = int(queue_status_data.get("remaining", 0) or 0)
                scheduler_pending = int(
                    queue_status_data.get("pending", queue_status_data.get("remaining", 0)) or 0
                )

            sched_waiting = scheduler_waiting
            sched_in_progress = scheduler_in_progress
            sched_pending = scheduler_pending

        scheduler_summary = TaskSummary(
            waiting=sched_waiting,
            in_progress=sched_in_progress,
            pending=sched_pending,
            completed=sched_completed,
            failed=sched_failed,
            cancelled=sched_cancelled,
            total=sched_waiting
            + sched_in_progress
            + sched_completed
            + sched_failed
            + sched_cancelled,
        )

        return AllStatusResponse(
            data=AllStatusResponseData(
                scheduler_summary=scheduler_summary,
                all_tasks_summary=all_tasks_summary,
            )
        )
    except Exception as err:
        logger.error(f"Failed to get full scheduler status: {traceback.format_exc()}")
        raise HTTPException(status_code=500, detail="Failed to get full scheduler status") from err


def handle_scheduler_status(
    user_id: str, status_tracker: TaskStatusTracker, task_id: str | None = None
) -> StatusResponse:
    """
    Get scheduler running status for one or all tasks of a user.

    Retrieves task statuses from the persistent TaskStatusTracker.

    Args:
        user_id: User ID to query for.
        status_tracker: The TaskStatusTracker instance.
        task_id: Optional Task ID to query. Can be either:
                 - business_task_id (will aggregate all related item statuses)
                 - item_id (will return single item status)

    Returns:
        StatusResponse with a list of task statuses.

    Raises:
        HTTPException: If a specific task is not found.
    """
    response_data: list[StatusResponseItem] = []

    try:
        if task_id:
            # First try as business_task_id (aggregated query)
            business_task_data = status_tracker.get_task_status_by_business_id(task_id, user_id)
            if business_task_data:
                response_data.append(
                    StatusResponseItem(task_id=task_id, status=business_task_data["status"])
                )
            else:
                # Fallback: try as item_id (single item query)
                item_task_data = status_tracker.get_task_status(task_id, user_id)
                if not item_task_data:
                    raise HTTPException(
                        status_code=404, detail=f"Task {task_id} not found for user {user_id}"
                    )
                response_data.append(
                    StatusResponseItem(task_id=task_id, status=item_task_data["status"])
                )
        else:
            all_tasks = status_tracker.get_all_tasks_for_user(user_id)
            # The plan returns an empty list, which is good.
            # No need to check "if not all_tasks" explicitly before the list comprehension
            response_data = [
                StatusResponseItem(task_id=tid, status=t_data["status"])
                for tid, t_data in all_tasks.items()
            ]

        return StatusResponse(data=response_data)
    except HTTPException:
        # Re-raise HTTPException directly to preserve its status code (e.g., 404)
        raise
    except Exception as err:
        logger.error(f"Failed to get scheduler status for user {user_id}: {traceback.format_exc()}")
        raise HTTPException(status_code=500, detail="Failed to get scheduler status") from err


def handle_task_queue_status(
    user_id: str, mem_scheduler: OptimizedScheduler, task_id: str | None = None
) -> TaskQueueResponse:
    try:
        queue_wrapper = getattr(mem_scheduler, "memos_message_queue", None)
        if queue_wrapper is None:
            raise HTTPException(status_code=503, detail="Scheduler queue is not available")

        # Unwrap to the underlying queue if wrapped by ScheduleTaskQueue
        queue = getattr(queue_wrapper, "memos_message_queue", queue_wrapper)

        # Only support Redis-backed queue for now; try lazy init if not connected
        redis_conn = getattr(queue, "_redis_conn", None)
        if redis_conn is None:
            try:
                if hasattr(queue, "auto_initialize_redis"):
                    queue.auto_initialize_redis()
                    redis_conn = getattr(queue, "_redis_conn", None)
                if redis_conn and hasattr(queue, "connect"):
                    queue.connect()
            except Exception:
                redis_conn = None

        if redis_conn is None:
            raise HTTPException(status_code=503, detail="Scheduler queue not connected to Redis")

        # Use wrapper to list stream keys so it can adapt to local/redis queue
        stream_keys = queue_wrapper.get_stream_keys()
        # Filter by user_id; stream key format: {prefix}:{user_id}:{mem_cube_id}:{task_label}
        user_stream_keys = [sk for sk in stream_keys if f":{user_id}:" in sk]

        if not user_stream_keys:
            raise HTTPException(
                status_code=404, detail=f"No scheduler streams found for user {user_id}"
            )

        def _parse_user_id_from_stream(stream_key: str) -> str | None:
            try:
                parts = stream_key.split(":")
                if len(parts) < 3:
                    return None
                # prefix may contain multiple segments; user_id is the 2nd segment from the end - 1
                return parts[-3]
            except Exception:
                return None

        user_ids_present = {
            uid for uid in (_parse_user_id_from_stream(sk) for sk in stream_keys) if uid
        }

        pending_total = 0
        pending_detail: list[str] = []
        remaining_total = 0
        remaining_detail: list[str] = []

        consumer_group = getattr(queue, "consumer_group", None) or "scheduler_group"
        for sk in user_stream_keys:
            try:
                pending_info = redis_conn.xpending(sk, consumer_group)
                pending_count = pending_info[0] if pending_info else 0
            except Exception:
                pending_count = 0
            pending_total += pending_count
            pending_detail.append(f"{sk}:{pending_count}")

            try:
                remaining_count = redis_conn.xlen(sk)
            except Exception:
                remaining_count = 0
            remaining_total += remaining_count
            remaining_detail.append(f"{sk}:{remaining_count}")

        data = TaskQueueData(
            user_id=user_id,
            user_name=None,
            mem_cube_id=None,
            stream_keys=user_stream_keys,
            users_count=len(user_ids_present),
            pending_tasks_count=pending_total,
            remaining_tasks_count=remaining_total,
            pending_tasks_detail=pending_detail,
            remaining_tasks_detail=remaining_detail,
        )
        return TaskQueueResponse(data=data)
    except HTTPException:
        # Re-raise HTTPException directly to preserve its status code (e.g., 404)
        raise
    except Exception as err:
        logger.error(
            f"Failed to get task queue status for user {user_id}: {traceback.format_exc()}"
        )
        raise HTTPException(status_code=500, detail="Failed to get scheduler status") from err


def handle_scheduler_wait(
    user_name: str,
    status_tracker: TaskStatusTracker,
    timeout_seconds: float = 120.0,
    poll_interval: float = 0.5,
) -> dict[str, Any]:
    """
    Wait until the scheduler is idle for a specific user.

    Blocks and polls the new /scheduler/status endpoint until no tasks are in
    'waiting' or 'in_progress' state, or until a timeout is reached.

    Args:
        user_name: User name to wait for.
        status_tracker: The TaskStatusTracker instance.
        timeout_seconds: Maximum wait time in seconds.
        poll_interval: Polling interval in seconds.

    Returns:
        Dictionary with wait result and statistics.

    Raises:
        HTTPException: If wait operation fails.
    """
    start_time = time.time()
    try:
        while time.time() - start_time < timeout_seconds:
            # Directly call the new, reliable status logic
            status_response = handle_scheduler_status(
                user_id=user_name, status_tracker=status_tracker
            )

            # System is idle if the data list is empty or no tasks are active
            is_idle = not status_response.data or all(
                task.status in ["completed", "failed", "cancelled"] for task in status_response.data
            )

            if is_idle:
                return {
                    "message": "idle",
                    "data": {
                        "running_tasks": 0,  # Kept for compatibility
                        "waited_seconds": round(time.time() - start_time, 3),
                        "timed_out": False,
                        "user_name": user_name,
                    },
                }

            time.sleep(poll_interval)

        # Timeout occurred
        final_status = handle_scheduler_status(user_id=user_name, status_tracker=status_tracker)
        active_tasks = [t for t in final_status.data if t.status in ["waiting", "in_progress"]]

        return {
            "message": "timeout",
            "data": {
                "running_tasks": len(active_tasks),  # A more accurate count of active tasks
                "waited_seconds": round(time.time() - start_time, 3),
                "timed_out": True,
                "user_name": user_name,
            },
        }
    except HTTPException:
        # Re-raise HTTPException directly to preserve its status code
        raise
    except Exception as err:
        logger.error(
            f"Failed while waiting for scheduler for user {user_name}: {traceback.format_exc()}"
        )
        raise HTTPException(status_code=500, detail="Failed while waiting for scheduler") from err


def handle_scheduler_wait_stream(
    user_name: str,
    status_tracker: TaskStatusTracker,
    timeout_seconds: float = 120.0,
    poll_interval: float = 0.5,
    instance_id: str = "",
) -> StreamingResponse:
    """
    Stream scheduler progress via Server-Sent Events (SSE) using the new status endpoint.

    Emits periodic heartbeat frames while tasks are active, then a final
    status frame indicating idle or timeout.

    Args:
        user_name: User name to monitor.
        status_tracker: The TaskStatusTracker instance.
        timeout_seconds: Maximum stream duration in seconds.
        poll_interval: Polling interval between updates.
        instance_id: Instance ID for response.

    Returns:
        StreamingResponse with SSE formatted progress updates.
    """

    def event_generator():
        start_time = time.time()
        try:
            while True:
                elapsed = time.time() - start_time
                if elapsed > timeout_seconds:
                    # Send timeout message and break
                    final_status = handle_scheduler_status(
                        user_id=user_name, status_tracker=status_tracker
                    )
                    active_tasks = [
                        t for t in final_status.data if t.status in ["waiting", "in_progress"]
                    ]
                    payload = {
                        "user_name": user_name,
                        "active_tasks": len(active_tasks),
                        "elapsed_seconds": round(elapsed, 3),
                        "status": "timeout",
                        "timed_out": True,
                        "instance_id": instance_id,
                    }
                    yield "data: " + json.dumps(payload, ensure_ascii=False) + "\n\n"
                    break

                # Get status
                status_response = handle_scheduler_status(
                    user_id=user_name, status_tracker=status_tracker
                )
                active_tasks = [
                    t for t in status_response.data if t.status in ["waiting", "in_progress"]
                ]
                num_active = len(active_tasks)

                payload = {
                    "user_name": user_name,
                    "active_tasks": num_active,
                    "elapsed_seconds": round(elapsed, 3),
                    "status": "running" if num_active > 0 else "idle",
                    "instance_id": instance_id,
                }
                yield "data: " + json.dumps(payload, ensure_ascii=False) + "\n\n"

                if num_active == 0:
                    break  # Exit loop if idle

                time.sleep(poll_interval)

        except Exception as e:
            err_payload = {
                "status": "error",
                "detail": "stream_failed",
                "exception": str(e),
                "user_name": user_name,
            }
            logger.error(f"Scheduler stream error for {user_name}: {traceback.format_exc()}")
            yield "data: " + json.dumps(err_payload, ensure_ascii=False) + "\n\n"

    return StreamingResponse(event_generator(), media_type="text/event-stream")
