import os
import json
import asyncio
from pathlib import Path
import traceback
from blocks import bash
from starlette.requests import Request
from starlette.responses import JSONResponse
from blocks_control_sdk.control.agent_base import LLM, NotifyCompleteArgs, NotificationsV2
from pydantic import BaseModel, Field
from typing import List, Dict, Optional
from blocks_control_sdk.prompt_builder import get_explore_subagent_prompt
from blocks_control_sdk.constants.core import WORKSPACE_DIR
from blocks_control_sdk.tools.github import (
    get_new_github_token,
    download_images_from_github_comment_id
)
from blocks_control_sdk.tools.hybrid import (
    clone_repository_into_folder,
    register_pull_request,
    send_internal_message,
    create_pull_request,
    create_pr_review_comment_on_file,
    create_pr_review_comment_on_lines,
)
from blocks_control_sdk.control.agent_codex import CodexAgentCLI, CodexAgentConfig
from blocks_control_sdk.tools.slack import download_image_from_slack, upload_file__slack, join_channel__slack
from blocks_control_sdk.tools.linear import download_linear_image
from blocks_control_sdk.clients.api import BlocksMessageCreate, Chat
from blocks_control_sdk.tools.blocks import send_message__blocks
from blocks_control_sdk.constants.openai import OpenAIModels
from blocks_control_sdk.utils import get_blocks_runtime_config, BlocksRuntimeConfigKeys
import tempfile
import time

SLACK_TOKEN = os.getenv("SLACK_TOKEN")
BLOCKS_TASK_ID = os.getenv("BLOCKS_TASK_ID")
LINEAR_TOKEN = os.getenv("LINEAR_TOKEN")

BLOCKS_TRIGGER_ALIAS = os.getenv("BLOCKS_TRIGGER_ALIAS")

# Configuration for the number of IMPORTANT clone tools to register
NUM_IMPORTANT_CLONE_TOOLS = 1

EXLORE_SUBAGENT_PROMPT = get_explore_subagent_prompt()



# Pydantic models for AskUserQuestion tool
class QuestionOption(BaseModel):
    label: str = Field(..., description="The display text for this option that the user will see and select. Should be concise (1-5 words) UNDER 30 CHARACTERS and clearly describe the choice.")
    description: str = Field(..., description="Explanation of what this option means or what will happen if chosen. Useful for providing context about trade-offs or implications.")


class Question(BaseModel):
    question: str = Field(..., description='The complete question to ask the user. Should be clear, specific, and end with a question mark. Example: "Which library should we use for date formatting?" If multiSelect is true, phrase it accordingly, e.g. "Which features do you want to enable?"')
    header: str = Field(..., description="Very short label displayed as a chip/tag (max 12 chars). Examples: 'Auth method', 'Library', 'Approach'.")
    options: List[QuestionOption] = Field(..., min_length=2, max_length=4, description="The available choices for this question. Must have 2-4 options. Each option should be a distinct, mutually exclusive choice (unless multiSelect is enabled). There should be no 'Other' option, that will be provided automatically.")
    multiSelect: bool = Field(..., description="Set to true to allow the user to select multiple options instead of just one. Use when choices are not mutually exclusive.")


class AskUserQuestionInput(BaseModel):
    questions: List[Question] = Field(..., min_length=1, max_length=4, description="Questions to ask the user (1-4 questions)")
    answers: Optional[Dict[str, str]] = Field(None, description="User answers collected by the permission component")

def register_tools(mcp, llm_provider: LLM):

    @mcp.custom_route("/v1/runtime/artifacts/{artifact_id}", methods=["PUT"])
    async def runtime_artifacts_notification(request: Request) -> JSONResponse:
        print("*"*100)
        body = await request.json()
        print(f"[Runtime Artifacts] Route Body: {body}")
        print("*"*100)

        artifact_id = body.get("artifact_id")
        filepath = body.get("filepath")
        event_type = body.get("event_type")
        file_size = body.get("file_size")
        metadata = body.get("metadata", {}) or {}

        # If none is provided, the send_message__blocks function will choose a sensible default
        chat_thread_id = metadata.get("chat_thread_id") or None
        tags = metadata.get("tags") or []

        if "events" in tags:
            if "prs" in tags:
                try:
                    with open(filepath, "r") as f:
                        file_contents = f.read()
                        file_contents = json.loads(file_contents)
                        url = file_contents.get("url")
                        is_successful = file_contents.get("is_successful")
                        if url and is_successful:
                            register_pull_request(url)
                        else:
                            print("[Blocks config event]: PR registration failed because url was not provided or was not successful")
                except Exception as e:
                    print("[Blocks config event]: PR registration failed")
            return JSONResponse({"status": "ok"})

        artifact_metadata = {
            "artifact_id": artifact_id,
            "file_path": filepath,
            "file_name": filepath.split("/")[-1],
            "event_type": event_type,
            "file_size_in_kb": file_size if file_size is not None else 0,
            "mime_type": "octet-stream"
        }

        message_body = json.dumps(artifact_metadata)

        send_message__blocks(
            message_body, 
            role="assistant", 
            type="file", 
            artifact_id=artifact_id, 
            artifact_metadata=artifact_metadata,
            chat_thread_id=chat_thread_id
        )
        return JSONResponse({"status": "ok"})
    
    @mcp.custom_route("/v1/logs", methods=["POST"])
    async def otel_logs(request: Request) -> JSONResponse:
        print("*"*100)
        print("OTEL LOGS:")
        body = await request.json()
        print(body)
        print("*"*100)
        return JSONResponse({"status": "ok"})

    print("Registering tools...")

    ################################################################################################################
    # GET-NEW-GITHUB-TOKEN TOOL
    ################################################################################################################

    print("Registering get_new_github_token tool...")
    @mcp.tool(name="get_new_github_token", description="IMPORTANT: When you are having Github authentication issues use this tool to get a new GitHub token.")
    def _get_new_github_token() -> str:
        """
        Get a new GitHub token when there is any Github auth issues
        """
        return get_new_github_token()

    print("Registering register_running_server_port tool...")
    @mcp.tool(name="register_running_server_port", description="Register a running server port")
    def _register_running_server_port(port: int, name: str) -> str:
        """
        Register a running server port
        Args:
            port: int - The port to register
            name: str - A short name for the server
        """

        try:

            subdomain = f"{port}-{BLOCKS_TASK_ID}"

            chat_api = Chat()
            chat_thread_id = get_blocks_runtime_config().get(BlocksRuntimeConfigKeys.CHAT_THREAD_ID)
        
            proxy_template = f"""[[proxies]]
name = "{subdomain}"
type = "http"
localIP = "0.0.0.0"
localPort = {port}
subdomain = "{subdomain}"
    """

            with open(f"/home/user/.proxies/port-{port}.toml", "a") as f:
                f.write(proxy_template)

            chat_api.create_message(
                BlocksMessageCreate(
                    message=json.dumps({
                        "__name__": "preview",
                        "name": name,
                        "subdomain": subdomain
                    }),
                    role="assistant",
                    type="telemetry",
                    chat_thread_id=chat_thread_id
                )
            )

            bash(f"frpc reload -c /var/lib/frpc/frpc.toml")

            return f"Registered running server port {port}. IMPORTANT REMINDER: Make sure to run the server as a background process so it persists indefinitely NOT using `run_in_background` function but the bash way of running it in the background with '&' at the end of the command."
        except Exception as e:
            traceback.print_exc()
            return f"Error registering running server port {port}: {e}"

    ##################################################################################################################
    # ASK-USER-QUESTION TOOL
    ##################################################################################################################

    print("Registering AskUserQuestion tool...")
    @mcp.tool(
        name="AskUserQuestion",
        description="""Use this tool when you need to ask the user questions during execution. This allows you to:
1. Gather user preferences or requirements
2. Clarify ambiguous instructions
3. Get decisions on implementation choices as you work
4. Offer choices to the user about what direction to take.

Usage notes:
- Users will always be able to select "Other" to provide custom text input
- Use multiSelect: true to allow multiple answers to be selected for a question
- If you recommend a specific option, make that the first option in the list and add "(Recommended)" at the end of the label
"""
    )
    def _ask_user_question(input_data: AskUserQuestionInput) -> str:
        """
        Ask the user questions during execution

        Args:
            input_data: AskUserQuestionInput - The questions to ask the user
        """
        # # Emit elicitation activity if in agent session
        # from tools.linear import is_agent_session_active, create_elicitation_activity

        # if is_agent_session_active():
        #     # Format questions for display in Linear
        #     questions_text = []
        #     for q in input_data.questions:
        #         questions_text.append(f"**{q.question}**")
        #         for opt in q.options:
        #             questions_text.append(f"  • {opt.label}: {opt.description}")

        #     elicitation_body = "I need your input:\n\n" + "\n\n".join(questions_text)
        #     # create_elicitation_activity(body=elicitation_body)

        # For now, return a simple JSON representation
        # This will be handled by the MCP server framework
        return "Choose the recommended options on my behalf, ensuring that get additional context if it will help make more informed decisions like recalling my original request, looking up relevant documentation, etc."

    ##################################################################################################################
    # EXIT-PLAN-MODE TOOL
    ##################################################################################################################

    print("Registering exit_plan_mode tool...")
    @mcp.tool(name="ExitPlanMode", description="Exit plan mode")
    def _exit_plan_mode() -> str:
        return "CRITICAL: Ensure you end your turn now. Simply present a final message with a concise overview of the plan and await for further instructions. "


    if llm_provider == LLM.CODEX or llm_provider == LLM.GEMINI or llm_provider == LLM.CURSOR or llm_provider == LLM.KIMI:
        print("Registering ExploreSubAgent tool...")
        @mcp.tool(name="ExploreSubAgent", description=EXLORE_SUBAGENT_PROMPT)
        def _general_sub_agent_prompt(description: str, prompt: str, subagent_type: str) -> str:
            print("[Subagent] Starting Task tool...")

            prompt = "CRITICAL: You are in read-only mode. DO NOT make any code changes. Here is my request: " + prompt
            
            temp_final_message_file = tempfile.mktemp()
            prompt_temp_file = tempfile.mktemp()
            with open(prompt_temp_file, 'w') as f:
                f.write(f"{prompt}")
                f.flush()

            if llm_provider == LLM.CODEX:
                command = f'cd {WORKSPACE_DIR.absolute()} && codex exec --profile readonly --output-last-message {temp_final_message_file} --model gpt-5.1-codex-mini --skip-git-repo-check --dangerously-bypass-approvals-and-sandbox "$(cat {prompt_temp_file})"'

            elif llm_provider == LLM.GEMINI:
                command = f'cd {WORKSPACE_DIR.absolute()} && gemini -m gemini-flash-latest --output-format json -p "$(cat {prompt_temp_file})" | jq -r ".response" > {temp_final_message_file}'

            elif llm_provider == LLM.CURSOR:
                command = f'cd {WORKSPACE_DIR.absolute()} && agent --print --output-format json --mode ask "$(cat {prompt_temp_file})" | jq -r ".result" > {temp_final_message_file}'   

            elif llm_provider == LLM.KIMI:
                explore_config = os.path.join(str(Path.home()), ".config", "blocks", "kimi_explore.config.yaml")
                command = f'cd {WORKSPACE_DIR.absolute()} && kimi --quiet --agent-file {explore_config} -p "$(cat {prompt_temp_file})" > {temp_final_message_file}'

            print("-"*100)
            print(f"SUBAGENT: Executing command: {command}")
            print("-"*100)
            out = bash(command, background=True)

            pid = out.pid


            print("[Subagent] Registering child PID...")
            send_internal_message(
                payload={"action": "register_child_pid", "pid": pid},
                src="Task"
            )

            print("[Subagent] Waiting for final message...")

            out.wait()

            final_message = "No final message received from the subagent."
            with open(temp_final_message_file, 'r') as f:
                final_message = f.read()

            print("[Subagent] Final message received: ", final_message)

            return final_message

    ##################################################################################################################
    # CREATE-PULL-REQUEST TOOL
    ##################################################################################################################

    if llm_provider != LLM.CLAUDE:
        for i in range(1, NUM_IMPORTANT_CLONE_TOOLS + 1):
            print("Registering create_pull_request tool...")
            @mcp.tool(name="create_pull_request", description="Create a pull request in GitHub. IMPORTANT: Always use this tool to create a pull request in GitHub instead of using the Github CLI.")
            def _create_pull_request(source_branch: str, target_branch: str, title: str, body: str, repo_owner: str, repo_name: str) -> str:
                return create_pull_request(source_branch, target_branch, title, body, repo_owner, repo_name)
        
    ##################################################################################################################
    # CREATE-PR-REVIEW-COMMENT-ON-FILE TOOL
    ##################################################################################################################

    print("Registering create_pr_review_comment_on_file tool...")
    @mcp.tool(name="create_pr_review_comment_on_file", description="Create a pull request review comment on a file")
    def _create_pr_review_comment_on_file(commit_id: str, file_path: str, pull_request_number: str, body: str, line: int, subject_type: str, side: str):
        return create_pr_review_comment_on_file(commit_id, file_path, pull_request_number, body, line, subject_type, side)

    ##################################################################################################################
    # CREATE-PR-REVIEW-COMMENT-ON-LINES TOOL
    ##################################################################################################################

    print("Registering create_pr_review_comment_on_lines tool...")
    @mcp.tool(name="create_pr_review_comment_on_lines", description="Create a pull request review comment on lines")
    def _create_pr_review_comment_on_lines(commit_id: str, file_path: str, pull_request_number: str, body: str, end_line: int, start_line: int, side: str):
        return create_pr_review_comment_on_lines(commit_id, file_path, pull_request_number, body, end_line, start_line, side)

    ##################################################################################################################
    # CLONE-REPOSITORY-INTO-FOLDER TOOL
    ##################################################################################################################

    print("Registering clone_repository_into_folder tool...")
    @mcp.tool(
        name="clone_repository_into_folder", 
        description="IMPORTANT: Clone a repository into a folder. Always use this tool to clone repositories belonging to me instead of using the Github CLI."
    )
    async def _clone_repository_into_folder_dynamic(url: str, folder_name: str, ref: str = "main") -> str:
          return await asyncio.to_thread(clone_repository_into_folder, url, folder_name, ref)

    print("Registering blocks_restart_agent tool...")
    @mcp.tool(name="blocks_restart_agent", description="This must be called after cloning one or more repositories into a folder when I specifically ask you to do so.")
    def _blocks_restart_agent() -> str:
        """
        Restart the agent.
        """
        send_internal_message(
            payload={"action": "hot_reload"},
            src="blocks_restart_agent"
        )
        return "Agent restarting..."

    ##################################################################################################################
    # CUSTOM SLACK TOOLS
    ##################################################################################################################

    if SLACK_TOKEN:

        print("Registering download_image_from_slack tool...")
        @mcp.tool(name="download_image_from_slack", description="Download an image from Slack")
        def _download_image_from_slack(file_id: str) -> str:
            """Download an image from Slack. This tool is used to download images from Slack.

            IMPORTANT: If you only have an image URL, you can extract the file ID from the URL:
            Example:
                - URL: https://files.slack.com/files-tmb/T086XLXR6VB-F09PUR6LFJL-d7c9117e60/screen_shot_2024-11-27_at_11.08.29_pm_480.png
                - File ID: F09PUR6LFJL

            Args:
                file_id: str - The file ID to download
            """
            return download_image_from_slack(file_id)
        
        print("Registering join_channel__slack tool...")
        @mcp.tool(name="join_channel_slack", description="Join a channel in Slack")
        def _join_channel_slack(channel_id: str) -> str:
            """
            Join a channel in Slack
            Args:
                channel_id: str - The channel ID to join
            """
            return join_channel__slack(channel_id)

        print("Registering upload_file__slack tool...")
        @mcp.tool(name="upload_file_slack", description="Upload a file to Slack channel")
        def _upload_file_slack(file_path: str, title: str, thread_ts: str, channel_id: str, initial_comment: str = None) -> str:
            """Upload a file to Slack channel for slack mentions
            Args:
                file_path: str - The path to the file to upload
                title: str - The title of the file
                thread_ts: str - The thread timestamp to add the file to (if not provided it will use the slack thread we are currently in)
                channel_id: str - The channel ID to upload the file to (if not provided it will use the slack channel we are currently in)
                initial_comment: str - The initial comment to add to the file
            """
            return upload_file__slack(file_path, title, thread_ts, channel_id, initial_comment)

    if LINEAR_TOKEN:

        print("Registering download_image_from_linear tool...")
        @mcp.tool(name="download_image_from_linear", description="Download an image from Linear")
        def _download_image_from_linear(image_url: str) -> str:
            """This tool is used to download images from Linear.
            Args:
                image_url: str - The URL of the image to download
            """
            return download_linear_image(image_url)

    ##################################################################################################################
    # CUSTOM GITHUB TOOLS
    ##################################################################################################################

    print("Registering download_images_from_github_comment_id tool...")
    @mcp.tool(name="download_images_from_github_comment_id", description="Download all images from a GitHub issue/PR comment")
    def _download_images_from_github_comment_id(owner: str, repo: str, comment_id: str) -> str:
        """Download all images from a GitHub issue or pull request comment. This tool fetches the comment's
        HTML representation through the GitHub API, which automatically converts user-attachment URLs to
        temporary signed URLs that can be downloaded.

        Args:
            owner: str - Repository owner (e.g., "BlocksOrg")
            repo: str - Repository name (e.g., "agents")
            comment_id: str - The comment ID (not the issue/PR number). This is the numeric ID of the comment,
                             which can be found in the comment's metadata or URL.

        Returns:
            Success message with paths to all downloaded images, or error message if no images found.
        """
        return download_images_from_github_comment_id(owner, repo, comment_id)

    ##################################################################################################################
    # SEND-MESSAGE TOOL
    ##################################################################################################################


    # if llm_provider == LLM.CLAUDE or llm_provider == LLM.CODEX:
    #     pass
    # else:
    #     if BLOCKS_TRIGGER_ALIAS == BLOCKS_EVENT__TRIGGER_ALIAS.SLACK_MENTION.value:
    #         print("Registering send_message__slack tool...")
    #         @mcp.tool(name="send_message", description="Send a message to the user")
    #         def _send_message(message: str, urgency_level_between_zero_to_10: int = 5) -> str:
    #             return send_message__slack(message, urgency_level_between_zero_to_10)

    #     elif BLOCKS_TRIGGER_ALIAS == BLOCKS_EVENT__TRIGGER_ALIAS.WEBHOOK.value:
    #         #no-op for now
    #         pass

    #     else:
    #         print("Registering send_message__blocks tool...")
    #         @mcp.tool(name="send_message", description="Send a message to the user")
    #         def _send_message(message: str, urgency_level_between_zero_to_10: int = 5) -> str:
    #             return send_message__blocks(message, urgency_level_between_zero_to_10)
