from optparse import Option
import os
import contextlib
from pathlib import Path
import traceback
from blocks import bash, config
from blocks_control_sdk.constants.core import WORKSPACE_DIR
import shutil
from blocks_control_sdk.activity_providers import AgentActivity
from blocks_control_sdk.activity_providers.base import SandboxPausingError
from blocks_control_sdk.control.agent_gemini_exp import GeminiAgentConfig
from blocks_control_sdk.control.agent_opencode import OpenCodeAgentConfig, SisyphusAgentConfig
from blocks_control_sdk.control.agent_cursor import CursorAgentConfig
from blocks_control_sdk.control.agent_kimi import KimiAgentConfig
from blocks_control_sdk.tools.github import get_new_github_token
from blocks_control_sdk.tools.gitlab import get_new_gitlab_token
from blocks_control_sdk.utils import (
    poll_for_port_open,
    get_blocks_runtime_config,
    patch_blocks_runtime_config,
    BlocksRuntimeConfigKeys,
)
from blocks_control_sdk.blocks_events.context import get_trigger_provider_from_trigger_alias
from blocks_control_sdk.tools.blocks import send_message__blocks
from blocks_control_sdk.control.agent_base import LLM, CodingAgentBaseCLI
from blocks_control_sdk.clients.api import Chat, UserQueryMode
from blocks_control_sdk.control.agent_codex import CodexAgentConfig
from blocks_control_sdk.control.agent_claude_exp import ClaudeAgentConfig
from pydantic import BaseModel, ConfigDict
import time
from typing import Optional, Union, List
from blocks_control_sdk.clients.api import BlocksMessage
from blocks_control_sdk.prompt_builder import PromptBuilder
from blocks_control_sdk.sandbox.monitor import SandboxMonitor

BLOCKS_INITIAL_CHAT_THREAD_ID = os.getenv("BLOCKS_INITIAL_CHAT_THREAD_ID")

config.set_token_resolver(lambda provider: (
    get_new_github_token() if provider == "github" else
    get_new_gitlab_token() if provider == "gitlab" else
    None
))

class AgentOptions(BaseModel):
    agent: CodingAgentBaseCLI
    llm_provider: LLM
    use_exp: bool = False
    mcp_server_enabled: bool = True
    runner: str = "fargate.medium"
    task_input: dict

    chat_messages: List[BlocksMessage] = []
    chat: Chat

    agent_config: Optional[Union[CodexAgentConfig, ClaudeAgentConfig, GeminiAgentConfig, OpenCodeAgentConfig, CursorAgentConfig, SisyphusAgentConfig, KimiAgentConfig]] = None

    activity: AgentActivity

    mcp_server_options: Optional[dict] = {
        "host": "127.0.0.1",
        "port": 8000
    }

    # allow arbitrary pydantic fields
    model_config = ConfigDict(arbitrary_types_allowed=True)


def agent_entrypoint(prompt: PromptBuilder, options: AgentOptions):
    print("########################### Agent Entrypoint ###########################")
    try:

        chat = options.chat
        messages = options.chat_messages

        agent = options.agent

        # Send SSH key as message
        chat.send_ssh_key_as_message(BLOCKS_INITIAL_CHAT_THREAD_ID)
        chat.send_code_server_password(lambda: agent.chat_thread_id)

        # Access the singleton sandbox monitor (created by SandboxActivityProvider during resolve)
        sandbox_monitor = SandboxMonitor()

        # get intended outcome summary from blocks runtime config
        intended_outcome_summary = get_blocks_runtime_config().get(BlocksRuntimeConfigKeys.INTENDED_OUTCOME_SUMMARY)

        if intended_outcome_summary:
            send_message__blocks(
                message=intended_outcome_summary, 
                urgency_level_between_zero_to_10=10, 
                role="assistant", 
                type="message", 
                chat_thread_id=BLOCKS_INITIAL_CHAT_THREAD_ID
            )
        else:
            send_message__blocks(
                message="On it.", 
                urgency_level_between_zero_to_10=10, 
                role="assistant", 
                type="message", 
                chat_thread_id=BLOCKS_INITIAL_CHAT_THREAD_ID
            )

        # If agent config is provided, use it   
        if options.agent_config:
            agent.config = options.agent_config

        options.activity.register(agent)

        is_resumable_session = bool(os.getenv("BLOCKS_SNAPSHOT_DOWNLOAD_URL"))

        if is_resumable_session:
            agent.is_session_active = True

        user_query = chat.prepare_query(
            messages=messages,
            llm_provider=options.llm_provider,
        )

        # For the first message, derive source_provider from trigger alias (e.g. "slack.mention" → "slack")
        # rather than relying on the message's source_provider which defaults to "webhook".
        resolved_source_provider = user_query.source_provider if (is_resumable_session and user_query.message) else get_trigger_provider_from_trigger_alias(prompt.trigger_alias)

        patch_blocks_runtime_config({
            BlocksRuntimeConfigKeys.LATEST_USER_MESSAGE: user_query.message,
            BlocksRuntimeConfigKeys.LATEST_USER_MESSAGE_SOURCE_PROVIDER: resolved_source_provider,
        })

        if user_query.mode == UserQueryMode.PLAN:
            agent.enter_plan_mode()

        WORKSPACE_DIR.mkdir(parents=True, exist_ok=True)
        print(f">> Changing directory to: {WORKSPACE_DIR.absolute()}")

        os.chdir(WORKSPACE_DIR.absolute())

        if options.mcp_server_enabled:
            write_config_kwargs = {}

            if options.llm_provider == LLM.CLAUDE:
                has_claude_oauth = chat.oauth_credentials_mapping.get("claude", None) is not None
                write_config_kwargs["has_oauth_credentials"] = has_claude_oauth

            agent.write_config(**write_config_kwargs)

        ##################################################################################################################
        # Process first user message
        ##################################################################################################################

        resolved_prompt = prompt.initial(user_query.message)
        agent.query(resolved_prompt)

        print("Agent initialized in background, giving buffer time...")

        if not is_resumable_session:
            time.sleep(10)

        ##################################################################################################################
        # Start chat polling
        ##################################################################################################################

        print("Waiting for messages...")

        while True:
            print("Waiting for messages in loop...")
            chat.re_create_session()
            messages = chat.get_messages()

            user_query = chat.prepare_query(
                messages=messages,
                llm_provider=options.llm_provider,
            )

            # Nothing to process — no user messages and no special mode
            if user_query.mode == UserQueryMode.EDIT and not user_query.has_user_messages:
                time.sleep(1)
                continue

            if user_query.mode == UserQueryMode.INTERRUPT:
                print("Last message is interrupt, interrupting agent...")
                agent.interrupt()
                time.sleep(1)
                continue

            try:
                options.activity.check(messages)
            except SandboxPausingError:
                break

            if user_query.has_user_messages:
                agent.interrupt()
                agent.new_chat_thread(chat_thread_id=user_query.chat_thread_id)

                if user_query.mode == UserQueryMode.PLAN:
                    agent.enter_plan_mode()
                    agent.query(prompt.followup(user_query.message))

                elif user_query.mode == UserQueryMode.CONTINUE:
                    agent.exit_plan_mode()
                    agent.interrupt()
                    agent.new_chat_thread(chat_thread_id=user_query.chat_thread_id, new_session=True)
                    agent.query(prompt.followup(user_query.message))
                else:
                    agent.query(prompt.followup(user_query.message))

            time.sleep(1)


        # Clean up sandbox monitor
        sandbox_monitor.stop()


    except Exception as e:
        # inlcude stacktrace
        print("Error during agent entrypoint: ", e)
        traceback.print_exc()
        raise e
