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.blocks_events.context import get_trigger_provider_from_trigger_alias
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 (
    get_file_path,
    poll_for_port_open,
    patch_blocks_runtime_config,
    BlocksRuntimeConfigKeys,
    get_blocks_runtime_config,
    has_system_command,
)
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
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, get_plan_prompt
from blocks_control_sdk.sandbox.monitor import SandboxMonitor

BLOCKS_INITIAL_CHAT_THREAD_ID = os.getenv("BLOCKS_INITIAL_CHAT_THREAD_ID")
BLOCKS_TRIGGER_ALIAS = os.getenv("BLOCKS_TRIGGER_ALIAS")

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)

        PATH__MESSAGE_HISTORY = get_file_path("BLOCKS_MESSAGE_HISTORY.xml")
        is_resumable_session = bool(os.getenv("BLOCKS_SNAPSHOT_DOWNLOAD_URL"))

        if is_resumable_session:
            agent.is_session_active = True

        _user_messages = [message for message in messages if message.role == "user" and message.type == "message"]
        latest_user_message_concat = ". ".join([user_message.message for user_message in _user_messages])
        # Plan mode specific prompt
        has_plan_mode = any([has_system_command(message.message, "plan") for message in _user_messages])

        if has_plan_mode:
            new_prompt = get_plan_prompt(options.llm_provider)
            latest_user_message_concat = new_prompt + "\n\n" + latest_user_message_concat
            agent.enter_plan_mode()

        latest_user_message = _user_messages[-1]

        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)

        resolved_prompt = prompt.initial(latest_user_message_concat)

        resolved_source_provider = latest_user_message.source_provider if (is_resumable_session and latest_user_message_concat) else get_trigger_provider_from_trigger_alias(BLOCKS_TRIGGER_ALIAS)


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

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

        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()

            # Check if we can immediately snapshot or not
            user_messages = [message for message in messages if message.role == "user" and message.type == "message"]

            is_any_message_interrupt = any([message for message in messages if message.type == "interrupt"])
            is_any_message_continue = any([has_system_command(message.message, "continue") for message in messages if message.role == "user"])
            is_any_message_plan = any([has_system_command(message.message, "plan") for message in user_messages])


            if is_any_message_continue or is_any_message_interrupt or user_messages:
                message_history = "\n" + chat.format_message_history(user_messages) 
                # Append to the message history file
                with open(PATH__MESSAGE_HISTORY, "a") as f:
                    f.write(message_history)
      
            if is_any_message_interrupt:
                print("Last message is interrupt, interrupting agent...")
                agent.interrupt()
                continue

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

            if user_messages:
                chat_thread_id = user_messages[-1].chat_thread_id
                latest_user_message_source_provider = user_messages[-1].source_provider
                concat_message = ". ".join([user_message.message for user_message in user_messages])

                agent.interrupt()

                agent.new_chat_thread(chat_thread_id=chat_thread_id)

                patch_blocks_runtime_config({
                    BlocksRuntimeConfigKeys.LATEST_USER_MESSAGE: concat_message,
                    BlocksRuntimeConfigKeys.LATEST_USER_MESSAGE_SOURCE_PROVIDER: latest_user_message_source_provider
                })

                if is_any_message_plan:
                    _plan_prompt = get_plan_prompt(options.llm_provider)
                    concat_message = _plan_prompt + "\n\n" + concat_message
                    agent.enter_plan_mode()
                    new_prompt = prompt.followup(concat_message)
                    agent.query(new_prompt)

                elif is_any_message_continue:
                    print("Last message is continue, restarting agent session...")
                    agent.exit_plan_mode()
                    agent.interrupt()
                    chat_thread_id = user_messages[-1].chat_thread_id if user_messages else BLOCKS_INITIAL_CHAT_THREAD_ID
                    agent.new_chat_thread(chat_thread_id=chat_thread_id, new_session=True)

                    message_history = PATH__MESSAGE_HISTORY.read_text()
                    new_message = f"""Your task is to implement the plan, described in detail in the latest file in the ~/.config/blocks/plans directory. Read my messages leading up to the plan creation in {PATH__MESSAGE_HISTORY.absolute()} for additional context of what I want."""

                    resolved_prompt = prompt.followup(new_message)
                    agent.query(resolved_prompt)
                else:
                    new_prompt = prompt.followup(concat_message)
                    agent.query(new_prompt)

            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
