debug_env/server/database/managers/task_manager.py

"""
Task Manager — Manages task metadata, episodes, and watch channels.

Follows the ACLManager pattern: class-based CRUD with user scoping, pagination,
sync tokens, and ETags for cache validation.
"""

import base64
import json
import logging
from datetime import datetime, timedelta
from typing import Optional, Dict, Any, List
from uuid import uuid4

from sqlalchemy import and_, desc, or_
from sqlalchemy.orm import Session

from debug_env.server.database.models.task import TaskRecord, EpisodeRecord, WatchChannel

logger = logging.getLogger(__name__)


class TaskManager:
    """
    Manages task operations with user-scoped access control.

    Args:
        db: SQLAlchemy Session
        user_id: User identifier for scoping operations
    """

    def __init__(self, db: Session, user_id: str):
        self.db = db
        self.user_id = user_id

    def list_tasks(
        self,
        max_results: int = 50,
        page_token: Optional[str] = None,
        sync_token: Optional[str] = None,
    ) -> "TaskListResponse":
        """
        List tasks with pagination and sync token support.

        Args:
            max_results: Maximum number of tasks per page (1-250, default 50)
            page_token: Pagination token from previous response
            sync_token: Sync token for incremental updates

        Returns:
            TaskListResponse with items, nextPageToken, nextSyncToken
        """
        try:
            # Handle sync token for incremental synchronization
            if sync_token:
                return self._handle_sync_request(sync_token, max_results)

            # Base query: all tasks (not user-scoped)
            query = self.db.query(TaskRecord).order_by(TaskRecord.id)

            # Handle pagination with base64 offset token
            offset = 0
            if page_token:
                try:
                    offset = int(base64.b64decode(page_token).decode("utf-8"))
                except (ValueError, TypeError):
                    raise ValueError("Invalid pageToken")

            # Get one extra to determine if there's a next page
            items = query.offset(offset).limit(max_results + 1).all()

            has_next_page = len(items) > max_results
            if has_next_page:
                items = items[:max_results]
                next_page_token = base64.b64encode(
                    str(offset + max_results).encode("utf-8")
                ).decode("utf-8")
            else:
                next_page_token = None

            # Generate next sync token based on latest update
            latest_updated = (
                self.db.query(TaskRecord.updated_at)
                .order_by(desc(TaskRecord.updated_at))
                .first()
            )

            next_sync_token = None
            if latest_updated and latest_updated[0]:
                sync_data = {
                    "timestamp": latest_updated[0].isoformat(),
                    "user_id": self.user_id,
                }
                next_sync_token = base64.b64encode(
                    json.dumps(sync_data).encode("utf-8")
                ).decode("utf-8")

            # Generate collection etag
            etag = f'"{uuid4()}"'

            return TaskListResponse(
                etag=etag,
                items=items,
                nextPageToken=next_page_token,
                nextSyncToken=next_sync_token,
            )

        except Exception as e:
            logger.error(f"Error listing tasks: {e}")
            raise

    def _handle_sync_request(
        self, sync_token: str, max_results: int
    ) -> "TaskListResponse":
        """
        Handle incremental synchronization request.

        Args:
            sync_token: Sync token from previous response
            max_results: Maximum items to return

        Returns:
            TaskListResponse with changes since sync token
        """
        try:
            # Decode sync token
            sync_data = json.loads(base64.b64decode(sync_token).decode("utf-8"))
            last_sync_time = datetime.fromisoformat(sync_data["timestamp"])

            # Check if sync token is too old (7-day expiration)
            if (datetime.utcnow() - last_sync_time).days > 7:
                raise ValueError("Sync token expired")

            # Query for changes since last sync
            query = (
                self.db.query(TaskRecord)
                .filter(TaskRecord.updated_at > last_sync_time)
                .order_by(TaskRecord.updated_at, TaskRecord.id)
            )

            items = query.limit(max_results + 1).all()

            has_more = len(items) > max_results
            if has_more:
                items = items[:max_results]

            # Generate new sync token
            next_sync_token = None
            if items:
                latest_time = max(item.updated_at for item in items)
                sync_data = {
                    "timestamp": latest_time.isoformat(),
                    "user_id": self.user_id,
                }
                next_sync_token = base64.b64encode(
                    json.dumps(sync_data).encode("utf-8")
                ).decode("utf-8")
            else:
                # No changes, return same sync token
                next_sync_token = sync_token

            # Generate etag
            etag = f'"{uuid4()}"'

            return TaskListResponse(
                etag=etag,
                items=items,
                nextPageToken=None,
                nextSyncToken=next_sync_token,
            )

        except json.JSONDecodeError:
            raise ValueError("Invalid sync token format")
        except Exception as e:
            logger.error(f"Error handling sync request: {e}")
            raise

    def get_task(self, task_id: str) -> Optional[TaskRecord]:
        """
        Retrieve a specific task by ID.

        Args:
            task_id: Task identifier (e.g., "task1")

        Returns:
            TaskRecord or None if not found
        """
        return self.db.query(TaskRecord).filter(TaskRecord.id == task_id).first()

    def create_episode(self, task_id: str, workdir: str) -> EpisodeRecord:
        """
        Create a new episode (task run) for a user.

        Args:
            task_id: Task identifier
            workdir: Path to working directory with task files

        Returns:
            Created EpisodeRecord
        """
        try:
            # Verify task exists
            task = self.get_task(task_id)
            if not task:
                raise ValueError(f"Task '{task_id}' not found")

            # Create episode
            episode_id = str(uuid4())
            episode = EpisodeRecord(
                id=episode_id,
                task_id=task_id,
                user_id=self.user_id,
                workdir=workdir,
                status="active",
                etag=f'"{uuid4()}"',
            )
            self.db.add(episode)
            self.db.commit()
            self.db.refresh(episode)

            logger.info(
                f"Created episode {episode_id} for task {task_id}, user {self.user_id}"
            )
            return episode

        except Exception as e:
            self.db.rollback()
            logger.error(f"Error creating episode: {e}")
            raise

    def get_episode(self, task_id: str, episode_id: str) -> Optional[EpisodeRecord]:
        """
        Retrieve a specific episode.

        Args:
            task_id: Task identifier
            episode_id: Episode identifier

        Returns:
            EpisodeRecord or None if not found
        """
        return (
            self.db.query(EpisodeRecord)
            .filter(
                and_(
                    EpisodeRecord.id == episode_id,
                    EpisodeRecord.task_id == task_id,
                    EpisodeRecord.user_id == self.user_id,
                )
            )
            .first()
        )

    def update_episode(
        self, task_id: str, episode_id: str, status: str, pass_rate: Optional[float] = None
    ) -> Optional[EpisodeRecord]:
        """
        Update an episode's status and pass rate.

        Args:
            task_id: Task identifier
            episode_id: Episode identifier
            status: New status ("active", "passed", "failed", "abandoned")
            pass_rate: Pass rate (0.0-1.0) if completed

        Returns:
            Updated EpisodeRecord or None if not found
        """
        try:
            episode = self.get_episode(task_id, episode_id)
            if not episode:
                return None

            episode.status = status
            if pass_rate is not None:
                episode.pass_rate = pass_rate
            episode.etag = f'"{uuid4()}"'
            episode.updated_at = datetime.utcnow()

            self.db.commit()
            self.db.refresh(episode)

            logger.info(
                f"Updated episode {episode_id} status to {status} (pass_rate={pass_rate})"
            )
            return episode

        except Exception as e:
            self.db.rollback()
            logger.error(f"Error updating episode: {e}")
            raise

    def list_episodes(
        self, task_id: str, max_results: int = 50, page_token: Optional[str] = None
    ) -> "EpisodeListResponse":
        """
        List episodes for a task with pagination.

        Args:
            task_id: Task identifier
            max_results: Maximum items per page
            page_token: Pagination token

        Returns:
            EpisodeListResponse with items and nextPageToken
        """
        try:
            # Base query: user-scoped
            query = (
                self.db.query(EpisodeRecord)
                .filter(
                    and_(EpisodeRecord.task_id == task_id, EpisodeRecord.user_id == self.user_id)
                )
                .order_by(EpisodeRecord.created_at, EpisodeRecord.id)
            )

            # Handle pagination
            offset = 0
            if page_token:
                try:
                    offset = int(base64.b64decode(page_token).decode("utf-8"))
                except (ValueError, TypeError):
                    raise ValueError("Invalid pageToken")

            items = query.offset(offset).limit(max_results + 1).all()

            has_next_page = len(items) > max_results
            if has_next_page:
                items = items[:max_results]
                next_page_token = base64.b64encode(
                    str(offset + max_results).encode("utf-8")
                ).decode("utf-8")
            else:
                next_page_token = None

            return EpisodeListResponse(items=items, nextPageToken=next_page_token)

        except Exception as e:
            logger.error(f"Error listing episodes: {e}")
            raise

    def watch_task(self, task_id: str, watch_request: Dict[str, Any]) -> WatchChannel:
        """
        Set up a watch channel for task state changes.

        Args:
            task_id: Task identifier
            watch_request: Dict with 'id', 'address', 'token' (optional)

        Returns:
            Created WatchChannel
        """
        try:
            # Verify task exists
            task = self.get_task(task_id)
            if not task:
                raise ValueError(f"Task '{task_id}' not found")

            # Check for duplicate watch ID
            existing = self.db.query(WatchChannel).filter(
                WatchChannel.id == watch_request["id"]
            ).first()
            if existing:
                raise ValueError(f"Watch channel with ID {watch_request['id']} already exists")

            # Create watch channel (24 hour default expiration)
            expires_at = datetime.utcnow() + timedelta(hours=24)
            channel = WatchChannel(
                id=watch_request["id"],
                task_id=task_id,
                user_id=self.user_id,
                webhook_address=watch_request["address"],
                webhook_token=watch_request.get("token"),
                expires_at=expires_at,
                is_active=True,
            )
            self.db.add(channel)
            self.db.commit()
            self.db.refresh(channel)

            logger.info(f"Created watch channel {watch_request['id']} for task {task_id}")
            return channel

        except Exception as e:
            self.db.rollback()
            logger.error(f"Error creating watch channel: {e}")
            raise

    def cleanup_expired_channels(self) -> int:
        """
        Clean up expired watch channels for this user.

        Returns:
            Number of channels cleaned up
        """
        try:
            current_time = datetime.utcnow()

            expired_channels = (
                self.db.query(WatchChannel)
                .filter(
                    and_(
                        WatchChannel.user_id == self.user_id,
                        WatchChannel.expires_at < current_time,
                        WatchChannel.is_active == True,
                    )
                )
                .all()
            )

            cleanup_count = 0
            for channel in expired_channels:
                channel.is_active = False
                cleanup_count += 1

            if cleanup_count > 0:
                self.db.commit()
                logger.info(f"Cleaned up {cleanup_count} expired watch channels for user {self.user_id}")

            return cleanup_count

        except Exception as e:
            self.db.rollback()
            logger.error(f"Error cleaning up expired channels: {e}")
            return 0


# Response models (simple classes matching the pattern)
class TaskListResponse:
    def __init__(
        self,
        etag: str,
        items: List[TaskRecord],
        nextPageToken: Optional[str] = None,
        nextSyncToken: Optional[str] = None,
    ):
        self.etag = etag
        self.items = items
        self.nextPageToken = nextPageToken
        self.nextSyncToken = nextSyncToken


class EpisodeListResponse:
    def __init__(self, items: List[EpisodeRecord], nextPageToken: Optional[str] = None):
        self.items = items
        self.nextPageToken = nextPageToken


def get_task_manager(db: Session, user_id: str) -> TaskManager:
    """Factory function to create a TaskManager instance."""
    return TaskManager(db, user_id)