from __future__ import annotations

import json
import logging
from io import BytesIO
from types import TracebackType
from typing import TYPE_CHECKING

from core.session.cli_api import CliApiSessionManager
from core.virtual_environment.__base.helpers import execute
from core.virtual_environment.__base.virtual_environment import VirtualEnvironment

from .bash.dify_cli import DifyCliConfig
from .constants import DIFY_CLI_CONFIG_PATH, DIFY_CLI_PATH
from .manager import SandboxManager
from .utils.debug import sandbox_debug

if TYPE_CHECKING:
    from core.tools.__base.tool import Tool

    from .bash.bash_tool import SandboxBashTool

logger = logging.getLogger(__name__)


class SandboxSession:
    def __init__(
        self,
        *,
        workflow_execution_id: str,
        tenant_id: str,
        user_id: str,
        tools: list[Tool],
    ) -> None:
        self._workflow_execution_id = workflow_execution_id
        self._tenant_id = tenant_id
        self._user_id = user_id
        self._tools = tools

        self._sandbox: VirtualEnvironment | None = None
        self._bash_tool: SandboxBashTool | None = None
        self._session_id: str | None = None

    def __enter__(self) -> SandboxSession:
        sandbox = SandboxManager.get(self._workflow_execution_id)
        if sandbox is None:
            raise RuntimeError(f"Sandbox not found for workflow_execution_id={self._workflow_execution_id}")

        session = CliApiSessionManager().create(tenant_id=self._tenant_id, user_id=self._user_id)
        self._session_id = session.id

        try:
            config = DifyCliConfig.create(session, self._tools)
            config_json = json.dumps(config.model_dump(mode="json"), ensure_ascii=False)

            sandbox_debug("sandbox", "config_json", config_json)
            sandbox.upload_file(DIFY_CLI_CONFIG_PATH, BytesIO(config_json.encode("utf-8")))

            execute(
                sandbox,
                [DIFY_CLI_PATH, "init"],
                timeout=30,
                error_message="Failed to initialize Dify CLI in sandbox",
            )

        except Exception:
            CliApiSessionManager().delete(session.id)
            self._session_id = None
            raise

        from .bash.bash_tool import SandboxBashTool

        self._sandbox = sandbox
        self._bash_tool = SandboxBashTool(sandbox=sandbox, tenant_id=self._tenant_id)
        return self

    def __exit__(
        self,
        exc_type: type[BaseException] | None,
        exc: BaseException | None,
        tb: TracebackType | None,
    ) -> bool:
        try:
            self.cleanup()
        except Exception:
            logger.exception("Failed to cleanup SandboxSession")
        return False

    @property
    def bash_tool(self) -> SandboxBashTool:
        if self._bash_tool is None:
            raise RuntimeError("SandboxSession is not initialized")
        return self._bash_tool

    def cleanup(self) -> None:
        if self._session_id is None:
            return

        CliApiSessionManager().delete(self._session_id)
        logger.debug("Cleaned up SandboxSession session_id=%s", self._session_id)
        self._session_id = None