# SPDX-FileCopyrightText: 2025 Weibo, Inc.
#
# SPDX-License-Identifier: Apache-2.0

"""Database handler for Chat Shell.

Note: This module is responsible ONLY for database operations.
WebSocket notifications should be handled by the emitter chain (WebSocketResultEmitter),
not by database operations. This follows the Single Responsibility Principle.

This module also provides:
- Thread pool executor for running sync DB operations without blocking event loop
- Context manager for database sessions with auto-rollback
- Utility function for running sync functions in executor

Other modules can use run_sync_in_executor() to wrap their sync DB functions.
"""

import asyncio
import logging
from concurrent.futures import ThreadPoolExecutor
from contextlib import contextmanager
from datetime import datetime
from typing import Any, Callable, Dict, Generator, Optional, TypeVar

from sqlalchemy.orm import Session
from sqlalchemy.orm.attributes import flag_modified

from app.models.subtask import Subtask
from app.schemas.kind import TaskStatus
from app.services.adapters.collaboration_strategy import CollaborationStrategy

logger = logging.getLogger(__name__)

# Thread pool for database operations
# Increased max_workers to handle concurrent DB operations across the application
_db_executor = ThreadPoolExecutor(max_workers=20)

# Terminal statuses that mark completion
_TERMINAL_STATUSES = frozenset(["COMPLETED", "FAILED", "CANCELLED"])

T = TypeVar("T")


@contextmanager
def _db_session() -> Generator[Session, None, None]:
    """Context manager for database session with auto-rollback on error."""
    from app.db.session import SessionLocal

    db = SessionLocal()
    try:
        yield db
        db.commit()
    except Exception:
        db.rollback()
        raise
    finally:
        db.close()


class DatabaseHandler:
    """Handles database operations for chat service."""

    async def _run_in_executor(self, func: Callable[..., T], *args: Any) -> T:
        """Run a synchronous function in the thread pool executor."""
        return await asyncio.get_event_loop().run_in_executor(_db_executor, func, *args)

    async def update_subtask_status(
        self,
        subtask_id: int,
        status: str,
        result: dict[str, Any] | None = None,
        error: str | None = None,
        executor_name: str | None = None,
        executor_namespace: str | None = None,
    ) -> None:
        """Update subtask status asynchronously."""
        await self._run_in_executor(
            self._update_subtask_sync,
            subtask_id,
            status,
            result,
            error,
            executor_name,
            executor_namespace,
        )

    def _update_subtask_sync(
        self,
        subtask_id: int,
        status: str,
        result: dict[str, Any] | None = None,
        error: str | None = None,
        executor_name: str | None = None,
        executor_namespace: str | None = None,
    ) -> None:
        """Synchronous subtask update (runs in thread pool)."""
        from app.models.subtask import Subtask, SubtaskStatus
        from app.utils.thinking_sanitizer import sanitize_result_for_storage

        try:
            with _db_session() as db:
                subtask = db.get(Subtask, subtask_id)
                if not subtask:
                    return

                subtask.status = SubtaskStatus(status)
                subtask.updated_at = datetime.now()

                if result is not None:
                    # Preserve silent_exit fields if they were set by MCP tool
                    # MCP silent_exit tool directly updates subtask.result before this
                    if subtask.result and isinstance(subtask.result, dict):
                        existing_silent_exit = subtask.result.get("silent_exit")
                        existing_silent_exit_reason = subtask.result.get(
                            "silent_exit_reason"
                        )
                        if existing_silent_exit:
                            result["silent_exit"] = existing_silent_exit
                            if existing_silent_exit_reason:
                                result["silent_exit_reason"] = (
                                    existing_silent_exit_reason
                                )
                            logger.info(
                                "[DB] Preserved silent_exit flag from MCP tool for subtask %s",
                                subtask_id,
                            )

                    # Log loaded_skills before sanitization for debugging
                    loaded_skills = result.get("loaded_skills", [])
                    if loaded_skills:
                        logger.info(
                            "[DB_HANDLER] Saving loaded_skills to database: "
                            "subtask_id=%d, status=%s, loaded_skills=%s",
                            subtask_id,
                            status,
                            loaded_skills,
                        )

                    # Sanitize thinking data before storing to database
                    sanitized_result = sanitize_result_for_storage(result)

                    # Log loaded_skills after sanitization to verify it's preserved
                    sanitized_loaded_skills = sanitized_result.get("loaded_skills", [])
                    if loaded_skills:
                        logger.info(
                            "[DB_HANDLER] After sanitization: subtask_id=%d, "
                            "loaded_skills_preserved=%s, sanitized_loaded_skills=%s",
                            subtask_id,
                            sanitized_loaded_skills == loaded_skills,
                            sanitized_loaded_skills,
                        )

                    subtask.result = sanitized_result
                if error is not None:
                    subtask.error_message = error
                if status in _TERMINAL_STATUSES:
                    subtask.completed_at = datetime.now()

                # Save executor info for container reuse in follow-up tasks
                # Only update if provided and subtask doesn't already have executor_name
                if executor_name and not subtask.executor_name:
                    subtask.executor_name = executor_name
                    subtask.executor_namespace = executor_namespace or ""
                    logger.info(
                        "[DB] Saved executor info for subtask %s: "
                        "executor_name=%s, executor_namespace=%s",
                        subtask_id,
                        executor_name,
                        executor_namespace,
                    )

                task_id = subtask.task_id
            # Context manager commits here, then update task status
            self._update_task_status_sync(task_id)
        except Exception:
            logger.exception("Error updating subtask %s status", subtask_id)

    def _update_task_status_sync(self, task_id: int) -> None:
        """Update task status based on subtask status using collaboration strategy."""
        from app.models.subtask import Subtask, SubtaskRole
        from app.models.task import TaskResource
        from app.schemas.kind import Task
        from app.services.adapters.collaboration_strategy import (
            CollaborationStrategy,
            CollaborationStrategyFactory,
        )

        should_schedule_dispatch = False
        try:
            with _db_session() as db:
                task = (
                    db.query(TaskResource)
                    .filter(
                        TaskResource.id == task_id,
                        TaskResource.kind == "Task",
                        TaskResource.is_active,
                    )
                    .first()
                )
                if not task:
                    return

                subtasks = (
                    db.query(Subtask)
                    .filter(
                        Subtask.task_id == task_id,
                        Subtask.role == SubtaskRole.ASSISTANT,
                    )
                    .order_by(Subtask.message_id.asc())
                    .all()
                )
                if not subtasks:
                    return

                task_crd = Task.model_validate(task.json)
                last_subtask = subtasks[-1]

                if task_crd.status:
                    # Use collaboration strategy to determine task status
                    strategy: CollaborationStrategy = (
                        CollaborationStrategyFactory.get_strategy_for_task(db, task_id)
                    )
                    self._apply_status_update(
                        db, task_id, task_crd.status, last_subtask, strategy
                    )
                    task_crd.status.updatedAt = datetime.now()

                    # Check if pipeline should auto-advance to next stage
                    advance_info = strategy.get_auto_advance_info(
                        db, task_id, last_subtask.id, last_subtask.status.value
                    )
                    if advance_info:
                        self._auto_advance_pipeline(
                            db, task, task_crd, last_subtask, advance_info
                        )
                        # Keep task RUNNING while next stage executes
                        task_crd.status.status = "PENDING"
                        task_crd.status.updatedAt = datetime.now()
                        should_schedule_dispatch = True

                task.json = task_crd.model_dump(mode="json")
                task.updated_at = datetime.now()
                flag_modified(task, "json")

        except Exception:
            logger.exception("Error updating task %s status", task_id)
            return

        if should_schedule_dispatch:
            from app.services.execution.schedule_helper import schedule_dispatch

            schedule_dispatch(task_id)

    def _auto_advance_pipeline(
        self,
        db: Session,
        task: Any,
        task_crd: Any,
        last_subtask: Subtask,
        advance_info: dict,
    ) -> None:
        """Create next pipeline stage subtask and advance currentStage in task spec.

        Called when a pipeline stage completes with requireConfirmation=False
        and more stages exist. The new subtask is added to the session; the
        caller is responsible for committing.

        Args:
            db: Database session (open, not yet committed)
            task: TaskResource ORM object
            task_crd: Parsed Task CRD (mutated in place)
            last_subtask: The subtask that just completed
            advance_info: Dict from get_auto_advance_info() with next stage details
        """
        from app.models.subtask import Subtask as SubtaskModel
        from app.models.subtask import SubtaskRole, SubtaskStatus

        next_stage_index = advance_info["next_stage_index"]
        next_bot_id = advance_info["next_bot_id"]
        next_bot_name = advance_info["next_bot_name"]

        # Reuse same executor container (pipeline runs all stages in one container)
        executor_name = last_subtask.executor_name or ""
        executor_namespace = last_subtask.executor_namespace or ""

        # Compute next message_id from highest existing message_id in this task
        latest = (
            db.query(SubtaskModel)
            .filter(SubtaskModel.task_id == task.id)
            .order_by(SubtaskModel.message_id.desc())
            .first()
        )
        next_message_id = (latest.message_id + 1) if latest else 1
        parent_id = latest.message_id if latest else 0

        # Advance currentStage in task spec
        task_crd.spec.currentStage = next_stage_index

        new_subtask = SubtaskModel(
            user_id=task.user_id,
            task_id=task.id,
            team_id=last_subtask.team_id,
            title=f"{task_crd.spec.title} - {next_bot_name}",
            bot_ids=[next_bot_id],
            role=SubtaskRole.ASSISTANT,
            prompt="",
            status=SubtaskStatus.PENDING,
            progress=0,
            message_id=next_message_id,
            parent_id=parent_id,
            executor_name=executor_name,
            executor_namespace=executor_namespace,
            error_message="",
            completed_at=datetime.now(),
            result=None,
        )
        db.add(new_subtask)
        logger.info(
            "[DB] Pipeline auto-advance task %s: stage %s, bot_id=%s",
            task.id,
            next_stage_index,
            next_bot_id,
        )

    def _apply_status_update(
        self,
        db: Session,
        task_id: int,
        status_obj: TaskStatus,
        last_subtask: Subtask,
        strategy: CollaborationStrategy,
    ) -> None:
        """Apply status update based on last subtask using collaboration strategy.

        Args:
            db: Database session
            task_id: Task ID
            status_obj: Task status object to update
            last_subtask: The last assistant subtask
            strategy: Collaboration strategy to use for status determination
        """
        from app.models.subtask import SubtaskStatus

        # Get the status string from subtask status enum
        subtask_status_str = last_subtask.status.value

        # Use strategy to determine task status
        new_status, progress = strategy.get_task_status_on_subtask_complete(
            db, task_id, last_subtask.id, subtask_status_str
        )

        status_obj.status = new_status
        status_obj.result = last_subtask.result

        if progress is not None:
            status_obj.progress = progress

        # Set completedAt for terminal statuses
        if new_status in ("COMPLETED", "CANCELLED", "PENDING_CONFIRMATION"):
            status_obj.completedAt = datetime.now()

        if last_subtask.status == SubtaskStatus.FAILED:
            status_obj.errorMessage = last_subtask.error_message

    async def save_partial_response(
        self, subtask_id: int, content: str, is_streaming: bool = True
    ) -> None:
        """Save partial response during streaming."""
        await self._run_in_executor(
            self._save_partial_sync, subtask_id, content, is_streaming
        )

    def _save_partial_sync(
        self, subtask_id: int, content: str, is_streaming: bool
    ) -> None:
        """Synchronous partial response save."""
        from app.models.subtask import Subtask

        try:
            with _db_session() as db:
                if subtask := db.get(Subtask, subtask_id):
                    subtask.result = {"value": content, "streaming": is_streaming}
                    subtask.updated_at = datetime.now()
        except Exception:
            logger.exception("Error saving partial response for subtask %s", subtask_id)

    async def get_subtask_message_id(self, subtask_id: int) -> int | None:
        """Get the message_id for a subtask.

        Args:
            subtask_id: The subtask ID

        Returns:
            The message_id if found, None otherwise
        """
        return await self._run_in_executor(
            self._get_subtask_message_id_sync, subtask_id
        )

    def _get_subtask_message_id_sync(self, subtask_id: int) -> int | None:
        """Synchronous get subtask message_id."""
        from app.models.subtask import Subtask

        try:
            with _db_session() as db:
                if subtask := db.get(Subtask, subtask_id):
                    return subtask.message_id
        except Exception:
            logger.exception("Error getting message_id for subtask %s", subtask_id)
        return None


# Global database handler instance
db_handler = DatabaseHandler()


# ============================================================
# Public utility functions for other modules
# ============================================================


async def run_sync_in_executor(func: Callable[..., T], *args: Any) -> T:
    """
    Run a synchronous function in the shared thread pool executor.

    This utility allows other modules to run their sync database operations
    without blocking the async event loop. The function runs in a shared
    thread pool with other database operations.

    Usage:
        from app.services.chat.storage.db import run_sync_in_executor

        async def my_async_function():
            result = await run_sync_in_executor(my_sync_db_function, arg1, arg2)
            return result

    Args:
        func: Synchronous function to execute
        *args: Arguments to pass to the function

    Returns:
        The return value of the synchronous function
    """
    return await asyncio.get_event_loop().run_in_executor(_db_executor, func, *args)


def get_db_session() -> Generator[Session, None, None]:
    """
    Get a database session context manager.

    This is a public alias for _db_session() for use by other modules.

    Usage:
        from app.services.chat.storage.db import get_db_session

        def my_sync_function():
            with get_db_session() as db:
                # do database operations
                pass

    Returns:
        Context manager yielding a database session
    """
    return _db_session()


def with_session_in_executor(func: Callable[..., T]) -> Callable[..., T]:
    """
    Decorator that wraps a sync function (db: Session, ...) -> T into an async function.

    The decorated function will:
    1. Create a new database session
    2. Run the original function in a thread pool executor
    3. Handle commit/rollback automatically
    4. Close the session when done

    The original sync function MUST have `db: Session` as its first parameter.

    Usage:
        from app.services.chat.storage.db import with_session_in_executor

        @with_session_in_executor
        def can_access_task_sync(db: Session, user_id: int, task_id: int) -> bool:
            task = db.query(Task).filter(Task.id == task_id).first()
            return task is not None and task.user_id == user_id

        # Now can_access_task_sync can be called as an async function:
        result = await can_access_task_sync(user_id, task_id)

    Args:
        func: Synchronous function with signature (db: Session, *args, **kwargs) -> T

    Returns:
        Async function with signature (*args, **kwargs) -> T
    """
    import functools

    @functools.wraps(func)
    async def wrapper(*args: Any, **kwargs: Any) -> T:
        def run_with_session() -> T:
            with _db_session() as db:
                return func(db, *args, **kwargs)

        return await asyncio.get_event_loop().run_in_executor(
            _db_executor, run_with_session
        )

    return wrapper
