from blocks_control_sdk.activity_providers.base import AgentActivityProvider
from blocks_control_sdk.blocks_events.context import BLOCKS_EVENT__PROVIDER, BLOCKS_EVENT__GITLAB__TYPE
from blocks_control_sdk.control.agent_base import LLM, NotifyMessageArgs, NotifyCompleteArgs, NotifyToolCallArgs
from blocks_control_sdk.tools.blocks import get_task_header
from blocks_control_sdk.tools.gitlab import (
    create_gitlab_issue_comment,
    create_gitlab_issue_discussion_reply,
    create_gitlab_merge_request_comment,
    create_gitlab_merge_request_discussion_reply,
    update_gitlab_issue_comment_with_header,
    update_gitlab_issue_comment,
    update_gitlab_merge_request_comment_with_header,
    update_gitlab_merge_request_comment,
)
from blocks_control_sdk.utils import patch_blocks_runtime_config, get_blocks_runtime_config

class GitLabActivityProvider(AgentActivityProvider):
    def _is_issue(self) -> bool:
        return self.trigger_alias == "gitlab.issue_comment"

    def setup(self, input: dict, llm_provider: LLM) -> None:
        super().setup(input, llm_provider)
        if self._is_issue():
            self._setup_issue(input, llm_provider)
        else:
            self._setup_merge_request(input, llm_provider)

    def _setup_issue(self, input: dict, llm_provider: LLM) -> None:
        issue = input.get("issue", {})
        issue_id = issue.get("id")
        issue_iid = issue.get("iid")

        project = input.get("project", {})
        project_path = project.get("path_with_namespace", "")

        issue_header = get_task_header()
        initial_gitlab_comment_id = input.get("$blocks.provider.comment_id")

        gitlab_comment_id = initial_gitlab_comment_id
        if not initial_gitlab_comment_id:
            comment_result = create_gitlab_issue_comment(project_path, issue_iid, issue_header)
            gitlab_comment_id = comment_result.get("id") if isinstance(comment_result, dict) else None
        else:
            comment_result = update_gitlab_issue_comment(project_path, issue_iid, initial_gitlab_comment_id, issue_header)
            gitlab_comment_id = comment_result.get("id") if isinstance(comment_result, dict) else initial_gitlab_comment_id

        patch_blocks_runtime_config({
            "event_provider": BLOCKS_EVENT__PROVIDER.GITLAB.value,
            "event_type": BLOCKS_EVENT__GITLAB__TYPE.ISSUE.value,
            "llm_provider": llm_provider.value,
            "metadata": {
                "comment_id": gitlab_comment_id,
                "issue_id": issue_id,
                "issue_iid": issue_iid,
                "project_path": project_path
            }
        })

    def _setup_merge_request(self, input: dict, llm_provider: LLM) -> None:
        merge_request = input.get("merge_request", {})
        mr_id = merge_request.get("id")
        mr_iid = merge_request.get("iid")

        project = input.get("project", {})
        project_path = project.get("path_with_namespace", "")

        issue_header = get_task_header()
        initial_gitlab_comment_id = input.get("$blocks.provider.comment_id")

        gitlab_comment_id = initial_gitlab_comment_id
        if not initial_gitlab_comment_id:
            comment_result = create_gitlab_merge_request_comment(project_path, mr_iid, issue_header)
            gitlab_comment_id = comment_result.get("id") if isinstance(comment_result, dict) else None
        else:
            comment_result = update_gitlab_merge_request_comment(project_path, mr_iid, initial_gitlab_comment_id, issue_header)
            gitlab_comment_id = comment_result.get("id") if isinstance(comment_result, dict) else initial_gitlab_comment_id

        patch_blocks_runtime_config({
            "event_provider": BLOCKS_EVENT__PROVIDER.GITLAB.value,
            "event_type": BLOCKS_EVENT__GITLAB__TYPE.MERGE_REQUEST.value,
            "llm_provider": llm_provider.value,
            "metadata": {
                "comment_id": gitlab_comment_id,
                "mr_id": mr_id,
                "mr_iid": mr_iid,
                "project_path": project_path
            }
        })

    def on_message(self, notification: NotifyMessageArgs) -> None:
        summary = self._prepare_progress_summary(notification)
        if summary:
            blocks_runtime_config = get_blocks_runtime_config()
            metadata = blocks_runtime_config.get("metadata", {})
            comment_id = metadata.get("comment_id")
            project_path = metadata.get("project_path")

            if comment_id and project_path:
                if self._is_issue():
                    issue_iid = metadata.get("issue_iid")
                    if issue_iid:
                        update_gitlab_issue_comment_with_header(project_path, issue_iid, comment_id, summary)
                else:
                    mr_iid = metadata.get("mr_iid")
                    if mr_iid:
                        update_gitlab_merge_request_comment_with_header(project_path, mr_iid, comment_id, summary)

    def on_complete(self, notification: NotifyCompleteArgs) -> None:
        message = notification.last_message
        if message:
            blocks_runtime_config = get_blocks_runtime_config()
            metadata = blocks_runtime_config.get("metadata", {})
            comment_id = metadata.get("comment_id")
            project_path = metadata.get("project_path")

            if comment_id and project_path:
                if self._is_issue():
                    issue_iid = metadata.get("issue_iid")
                    if issue_iid:
                        update_gitlab_issue_comment(project_path, issue_iid, comment_id, message)
                else:
                    mr_iid = metadata.get("mr_iid")
                    if mr_iid:
                        update_gitlab_merge_request_comment(project_path, mr_iid, comment_id, message)

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

    def check(self, messages: list) -> None:
        user_messages = self._get_user_messages(messages)
        if not user_messages:
            return

        new_comment = self.input.get("new_comment", {})
        discussion_id = new_comment.get("discussion_id")
        issue_header = get_task_header()

        comment_result = None

        if self._is_issue():
            issue = self.input.get("issue", {})
            issue_iid = issue.get("iid")
            project = self.input.get("project", {})
            project_path = project.get("path_with_namespace", "")

            if discussion_id:
                comment_result = create_gitlab_issue_discussion_reply(project_path, issue_iid, discussion_id, issue_header)
            else:
                comment_result = create_gitlab_issue_comment(project_path, issue_iid, issue_header)
        else:
            merge_request = self.input.get("merge_request", {})
            mr_iid = merge_request.get("iid")
            project = self.input.get("project", {})
            project_path = project.get("path_with_namespace", "")

            if discussion_id:
                comment_result = create_gitlab_merge_request_discussion_reply(project_path, mr_iid, discussion_id, issue_header)
            else:
                comment_result = create_gitlab_merge_request_comment(project_path, mr_iid, issue_header)

        new_comment_id = None
        if comment_result and isinstance(comment_result, dict):
            new_comment_id = comment_result.get("id")

        patch_blocks_runtime_config({"metadata": {"comment_id": new_comment_id}}, deep_patch=True)
