import os
import json
import traceback

from blocks_control_sdk.activity_providers.base import AgentActivityProvider, SandboxPausingError
from blocks_control_sdk.constants.core import is_micro_vm
from blocks_control_sdk.control.agent_base import LLM, NotifyStartArgs, NotifyResumeArgs, NotifyMessageArgs, NotifyCompleteArgs, NotifyToolCallArgs, NotifyInterruptArgs
from blocks_control_sdk.sandbox.monitor import SandboxMonitor, MonitorState
from blocks_control_sdk.clients.api import Orchestrator, BlocksUserTelemetryMessage
from blocks_control_sdk.tools.blocks import send_message__blocks
from blocks_control_sdk.utils import get_blocks_runtime_config, BlocksRuntimeConfigKeys


class SandboxActivityProvider(AgentActivityProvider):
    """Encapsulates sandbox monitor lifecycle wiring as an activity provider."""

    def __init__(self, trigger_alias: str):
        super().__init__(trigger_alias)

    @property
    def current_chat_thread_id(self) -> str:
        return get_blocks_runtime_config().get(BlocksRuntimeConfigKeys.CHAT_THREAD_ID)

    def setup(self, input: dict, llm_provider: LLM) -> None:
        _is_micro_vm = is_micro_vm()
        timeout = 2.5 * 60 if _is_micro_vm else 15 * 60

        def on_pause_requested():
            print("[SandboxActivityProvider] Pause requested, sending idle status...")
            Orchestrator().send_idle_status()

        def on_tick_callback(remaining_seconds: float):
            print(f"TIMER TICKING CALLBACK CALLED WITH REMAINING SECONDS: {remaining_seconds}")
            try:
                system_metrics_dict = {
                    "time_remaining_in_minutes": remaining_seconds / 60,
                }
                send_message__blocks(
                    message=json.dumps(system_metrics_dict),
                    role="assistant",
                    type="telemetry",
                    chat_thread_id=self.current_chat_thread_id
                )
            except Exception as e:
                print(f"ERROR IN TIMER TICKING CALLBACK: {e}")
                traceback.print_exc()

        print("[SandboxActivityProvider] Setting up sandbox monitor")
        monitor = SandboxMonitor()
        monitor.configure(
            default_timeout_seconds=timeout,
            on_pause_requested=on_pause_requested if _is_micro_vm else None,
            on_tick=on_tick_callback,
        )
        # Start with a long soft timeout during initialization
        SandboxMonitor().start_idle(timeout_seconds=60 * 60 * 2)  # 2 hours initial

    def on_start(self, notification: NotifyStartArgs) -> None:
        print("!" * 100)
        print("SANDBOX PROVIDER: NOTIFY START V2 CALLED")
        print("!" * 100)
        SandboxMonitor().start_active()
        chat_thread_id = notification.chat_thread_id
        send_message__blocks(
            message="[Agent is starting a session]",
            role="assistant",
            type="start",
            urgency_level_between_zero_to_10=10,
            chat_thread_id=chat_thread_id,
        )

    def on_resume(self, notification: NotifyResumeArgs) -> None:
        print("!" * 100)
        print("SANDBOX PROVIDER: NOTIFY RESUME V2 CALLED")
        print("!" * 100)
        SandboxMonitor().start_active()

    def on_complete(self, notification: NotifyCompleteArgs) -> None:
        print("!" * 100)
        print("SANDBOX PROVIDER: NOTIFY COMPLETE V2 CALLED")
        print("!" * 100)
        chat_thread_id = notification.chat_thread_id
        SandboxMonitor().start_idle()
        send_message__blocks(
            message="[Agent has finished a session]",
            role="assistant",
            type="end",
            urgency_level_between_zero_to_10=10,
            chat_thread_id=chat_thread_id,
        )

    def check(self, messages: list) -> None:
        sandbox_monitor = SandboxMonitor()

        # Telemetry heartbeat dispatch (ported from entrypoint.py)
        telemetry_messages = [m for m in messages if m.role == "user" and m.type == "telemetry"]
        any_user_messages = [m for m in messages if m.role == "user"]

        if telemetry_messages:
            for msg in telemetry_messages:
                if msg.message == BlocksUserTelemetryMessage.USER_ACTIVITY.value:
                    sandbox_monitor.heartbeat(extend_seconds=2.5 * 60)
                elif msg.message == BlocksUserTelemetryMessage.WAKE_UP.value:
                    sandbox_monitor.heartbeat(extend_seconds=5 * 60)
                elif msg.message == BlocksUserTelemetryMessage.CODE_EDITOR_EXPANDED.value:
                    sandbox_monitor.heartbeat(extend_seconds=10 * 60)
                elif msg.message == BlocksUserTelemetryMessage.CODE_EDITOR_OPENED.value:
                    sandbox_monitor.heartbeat(extend_seconds=10 * 60)
        elif any_user_messages:
            sandbox_monitor.heartbeat(extend_seconds=1 * 60)

        # Monitor state check (ported from entrypoint.py)
        monitor_state = sandbox_monitor.check()

        if monitor_state == MonitorState.IDLE:
            remaining = sandbox_monitor.get_remaining_seconds()
            if remaining > 0:
                print(f"Agent idle, counting down... {remaining / 60:.2f} minutes remaining")

        if monitor_state == MonitorState.PAUSING:
            print("[SandboxActivityProvider] Monitor state is PAUSING")
            if not is_micro_vm():
                print("Max idle time reached, stopping agent...")
                raise SandboxPausingError("Max idle time reached")
            # VM mode: on_pause_requested callback already fired. VM will snapshot.

    def on_interrupt(self, notification: NotifyInterruptArgs) -> None:
        SandboxMonitor().start_idle(timeout_seconds=5 * 60)

    def on_message(self, notification: NotifyMessageArgs) -> None:
        pass

    def on_tool_call(self, notification: NotifyToolCallArgs) -> None:
        pass
