ForcePilot/backend/package/yuxi/channel/extensions/gitlab/api.py

from __future__ import annotations

import asyncio
import logging

import gitlab

from yuxi.channel.extensions.gitlab.rate_limiter import gitlab_rate_limiter

logger = logging.getLogger(__name__)


class AsyncGitLabClient:
    """python-gitlab 的异步包装器，所有同步调用通过 asyncio.to_thread 包装"""

    def __init__(self, base_url: str, private_token: str, ssl_verify: bool | str = True):
        self._gl = gitlab.Gitlab(
            url=base_url,
            private_token=private_token,
            ssl_verify=ssl_verify,
            retry_transient_errors=True,
            timeout=30,
        )
        self._project_cache: dict[int, gitlab.v4.objects.Project] = {}

    async def _run(self, func, *args, **kwargs):
        await gitlab_rate_limiter.acquire()
        return await asyncio.to_thread(func, *args, **kwargs)

    async def get_project(self, project_id: int):
        if project_id not in self._project_cache:
            self._project_cache[project_id] = await self._run(self._gl.projects.get, project_id)
        return self._project_cache[project_id]

    async def probe(self) -> dict | None:
        try:
            version = await self._run(self._gl.version)
            return version
        except Exception:
            logger.exception("GitLab API probe failed")
            return None

    async def get_current_user(self) -> dict:
        user = await self._run(self._gl.user)
        return user.as_dict() if hasattr(user, "as_dict") else dict(user.attributes)

    async def list_issues(self, project_id: int, state: str = "opened", **filters):
        project = await self.get_project(project_id)
        return await self._run(project.issues.list, state=state, per_page=100, **filters)

    async def get_issue(self, project_id: int, issue_iid: int):
        project = await self.get_project(project_id)
        return await self._run(project.issues.get, issue_iid)

    async def create_issue(self, project_id: int, data: dict):
        project = await self.get_project(project_id)
        return await self._run(project.issues.create, data)

    async def update_issue(self, project_id: int, issue_iid: int, data: dict):
        issue = await self.get_issue(project_id, issue_iid)
        for key, value in data.items():
            setattr(issue, key, value)
        return await self._run(issue.save)

    async def list_issue_notes(self, project_id: int, issue_iid: int, **filters):
        issue = await self.get_issue(project_id, issue_iid)
        return await self._run(issue.notes.list, per_page=100, **filters)

    async def create_issue_note(self, project_id: int, issue_iid: int, body: str) -> dict:
        issue = await self.get_issue(project_id, issue_iid)
        note = await self._run(issue.notes.create, {"body": body})
        return note.attributes if hasattr(note, "attributes") else dict(note)

    async def edit_note(
        self,
        project_id: int,
        noteable_type: str,
        noteable_iid: int,
        note_id: int,
        body: str,
    ) -> dict:
        project = await self.get_project(project_id)
        if noteable_type == "Issue":
            noteable = await self._run(project.issues.get, noteable_iid)
        elif noteable_type == "MergeRequest":
            noteable = await self._run(project.mergerequests.get, noteable_iid)
        else:
            raise ValueError(f"unsupported noteable_type: {noteable_type}")
        note = await self._run(noteable.notes.get, note_id)
        note.body = body
        await self._run(note.save)
        return note.attributes if hasattr(note, "attributes") else dict(note)

    async def create_mr_note(self, project_id: int, mr_iid: int, body: str) -> dict:
        project = await self.get_project(project_id)
        mr = await self._run(project.mergerequests.get, mr_iid)
        note = await self._run(mr.notes.create, {"body": body})
        return note.attributes if hasattr(note, "attributes") else dict(note)

    async def create_diff_note(
        self,
        project_id: int,
        mr_iid: int,
        body: str,
        file_path: str,
        new_line: int,
        old_line: int | None = None,
        base_sha: str = "",
        head_sha: str = "",
        start_sha: str = "",
    ) -> dict:
        project = await self.get_project(project_id)
        mr = await self._run(project.mergerequests.get, mr_iid)

        if not all([base_sha, head_sha]):
            diff_refs = mr.diff_refs
            base_sha = base_sha or diff_refs.get("base_sha", "")
            head_sha = head_sha or diff_refs.get("head_sha", "")
            start_sha = start_sha or diff_refs.get("start_sha", base_sha)

        position = {
            "base_sha": base_sha,
            "start_sha": start_sha,
            "head_sha": head_sha,
            "position_type": "text",
            "new_path": file_path,
            "old_path": file_path,
            "new_line": new_line,
            "old_line": old_line,
        }

        note = await self._run(mr.notes.create, {"body": body, "position": position})
        return note.attributes if hasattr(note, "attributes") else dict(note)

    async def list_mrs(
        self,
        project_id: int,
        state: str = "opened",
        order_by: str = "created_at",
        sort: str = "desc",
        **filters,
    ):
        project = await self.get_project(project_id)
        return await self._run(
            project.mergerequests.list,
            state=state,
            order_by=order_by,
            sort=sort,
            per_page=100,
            **filters,
        )

    async def get_mr(self, project_id: int, mr_iid: int):
        project = await self.get_project(project_id)
        return await self._run(project.mergerequests.get, mr_iid)

    async def get_mr_changes(self, project_id: int, mr_iid: int):
        mr = await self.get_mr(project_id, mr_iid)
        return await self._run(mr.changes)

    async def create_mr(
        self,
        project_id: int,
        source_branch: str,
        target_branch: str,
        title: str,
        description: str = "",
        labels: str = "",
        assignee_id: int | None = None,
        remove_source_branch: bool = False,
        squash: bool = False,
    ) -> dict:
        project = await self.get_project(project_id)
        data = {
            "source_branch": source_branch,
            "target_branch": target_branch,
            "title": title,
            "description": description,
            "remove_source_branch": remove_source_branch,
            "squash": squash,
        }
        if labels:
            data["labels"] = labels
        if assignee_id:
            data["assignee_id"] = assignee_id
        mr = await self._run(project.mergerequests.create, data)
        return mr.attributes if hasattr(mr, "attributes") else dict(mr)

    async def update_mr(self, project_id: int, mr_iid: int, data: dict) -> dict:
        mr = await self.get_mr(project_id, mr_iid)
        for key, value in data.items():
            setattr(mr, key, value)
        result = await self._run(mr.save)
        return result.attributes if hasattr(result, "attributes") else dict(result)

    async def merge_mr(
        self,
        project_id: int,
        mr_iid: int,
        merge_commit_message: str | None = None,
        squash: bool = False,
        should_remove_source_branch: bool = False,
        merge_when_pipeline_succeeds: bool = False,
    ) -> dict:
        mr = await self.get_mr(project_id, mr_iid)
        data = {
            "squash": squash,
            "should_remove_source_branch": should_remove_source_branch,
        }
        if merge_commit_message:
            data["merge_commit_message"] = merge_commit_message
        if merge_when_pipeline_succeeds:
            data["merge_when_pipeline_succeeds"] = True
        result = await self._run(mr.merge, **data)
        if isinstance(result, dict):
            return result
        if hasattr(result, "attributes"):
            return result.attributes
        return dict(result)

    async def approve_mr(self, project_id: int, mr_iid: int) -> dict:
        mr = await self.get_mr(project_id, mr_iid)
        result = await self._run(mr.approve)
        return result if isinstance(result, dict) else {}

    async def unapprove_mr(self, project_id: int, mr_iid: int) -> dict:
        mr = await self.get_mr(project_id, mr_iid)
        result = await self._run(mr.unapprove)
        return result if isinstance(result, dict) else {}

    async def get_pipeline(self, project_id: int, pipeline_id: int):
        project = await self.get_project(project_id)
        return await self._run(project.pipelines.get, pipeline_id)

    async def get_job_trace(self, project_id: int, job_id: int, tail_lines: int = 200) -> str:
        project = await self.get_project(project_id)
        job = await self._run(project.jobs.get, job_id)
        trace_bytes = await self._run(job.trace)
        trace = trace_bytes.decode("utf-8", errors="replace")
        lines = trace.splitlines()
        if len(lines) > tail_lines:
            return "\n".join(lines[-tail_lines:])
        return trace

    async def add_reaction(
        self,
        project_id: int,
        noteable_type: str,
        noteable_iid: int,
        note_id: int,
        emoji_name: str,
    ) -> dict:
        project = await self.get_project(project_id)
        if noteable_type == "Issue":
            noteable = await self._run(project.issues.get, noteable_iid)
        else:
            noteable = await self._run(project.mergerequests.get, noteable_iid)
        note = await self._run(noteable.notes.get, note_id)
        emoji = await self._run(note.awardemojis.create, {"name": emoji_name})
        return emoji.attributes if hasattr(emoji, "attributes") else dict(emoji)

    async def upload_file(self, project_id: int, file_path: str) -> dict:
        project = await self.get_project(project_id)
        uploaded = await self._run(project.upload, file_path)
        return {
            "url": uploaded.get("url", ""),
            "markdown": uploaded.get("markdown", ""),
        }

    async def list_hooks(self, project_id: int):
        project = await self.get_project(project_id)
        return await self._run(project.hooks.list)

    async def create_hook(self, project_id: int, data: dict):
        project = await self.get_project(project_id)
        return await self._run(project.hooks.create, data)

    async def delete_hook(self, project_id: int, hook_id: int):
        project = await self.get_project(project_id)
        hook = await self._run(project.hooks.get, hook_id)
        await self._run(hook.delete)

    async def search(
        self,
        project_id: int,
        scope: str,
        query: str,
        project_ids: list[int] | None = None,
        group_ids: list[int] | None = None,
    ) -> list:
        project = await self.get_project(project_id)
        kwargs = {}
        if project_ids:
            kwargs["project_ids"] = project_ids
        if group_ids:
            kwargs["group_ids"] = group_ids
        return await self._run(project.search, scope, query, **kwargs)

    async def search_code(self, project_id: int, query: str, scope: str = "blobs"):
        project = await self.get_project(project_id)
        return await self._run(project.search, scope, query)

    async def list_pipelines(self, project_id: int, status: str | None = None, **filters):
        project = await self.get_project(project_id)
        kwargs = {"per_page": 100, **filters}
        if status:
            kwargs["status"] = status
        return await self._run(project.pipelines.list, **kwargs)

    async def cancel_pipeline(self, project_id: int, pipeline_id: int) -> dict:
        project = await self.get_project(project_id)
        pipeline = await self._run(project.pipelines.get, pipeline_id)
        result = await self._run(pipeline.cancel)
        return result if isinstance(result, dict) else dict(result)

    async def retry_pipeline(self, project_id: int, pipeline_id: int) -> dict:
        project = await self.get_project(project_id)
        pipeline = await self._run(project.pipelines.get, pipeline_id)
        result = await self._run(pipeline.retry)
        return result if isinstance(result, dict) else dict(result)

    async def list_jobs(self, project_id: int, scope: str | None = None, **filters):
        project = await self.get_project(project_id)
        kwargs = {"per_page": 100, **filters}
        if scope:
            kwargs["scope"] = scope
        return await self._run(project.jobs.list, **kwargs)

    async def get_job(self, project_id: int, job_id: int):
        project = await self.get_project(project_id)
        return await self._run(project.jobs.get, job_id)

    async def cancel_job(self, project_id: int, job_id: int) -> dict:
        project = await self.get_project(project_id)
        job = await self._run(project.jobs.get, job_id)
        result = await self._run(job.cancel)
        return result if isinstance(result, dict) else dict(result)

    async def retry_job(self, project_id: int, job_id: int) -> dict:
        project = await self.get_project(project_id)
        job = await self._run(project.jobs.get, job_id)
        result = await self._run(job.retry)
        return result if isinstance(result, dict) else dict(result)

    async def list_branches(self, project_id: int, **filters):
        project = await self.get_project(project_id)
        return await self._run(project.branches.list, per_page=100, **filters)

    async def get_branch(self, project_id: int, branch_name: str):
        project = await self.get_project(project_id)
        return await self._run(project.branches.get, branch_name)

    async def create_branch(self, project_id: int, branch_name: str, ref: str):
        project = await self.get_project(project_id)
        branch = await self._run(project.branches.create, {"branch": branch_name, "ref": ref})
        return branch.attributes if hasattr(branch, "attributes") else dict(branch)

    async def delete_branch(self, project_id: int, branch_name: str):
        project = await self.get_project(project_id)
        await self._run(project.branches.delete, branch_name)
        return {"deleted": True}

    async def list_tags(self, project_id: int, **filters):
        project = await self.get_project(project_id)
        return await self._run(project.tags.list, per_page=100, **filters)

    async def create_tag(self, project_id: int, tag_name: str, ref: str, message: str = ""):
        project = await self.get_project(project_id)
        data = {"tag_name": tag_name, "ref": ref}
        if message:
            data["message"] = message
        tag = await self._run(project.tags.create, data)
        return tag.attributes if hasattr(tag, "attributes") else dict(tag)

    async def delete_tag(self, project_id: int, tag_name: str):
        project = await self.get_project(project_id)
        await self._run(project.tags.delete, tag_name)
        return {"deleted": True}

    async def list_releases(self, project_id: int, **filters):
        project = await self.get_project(project_id)
        return await self._run(project.releases.list, per_page=100, **filters)

    async def create_release(
        self,
        project_id: int,
        tag_name: str,
        name: str,
        description: str = "",
        ref: str = "",
    ) -> dict:
        project = await self.get_project(project_id)
        data = {"tag_name": tag_name, "name": name, "description": description}
        if ref:
            data["ref"] = ref
        release = await self._run(project.releases.create, data)
        return release.attributes if hasattr(release, "attributes") else dict(release)

    async def get_file(self, project_id: int, file_path: str, ref: str = "main") -> dict:
        project = await self.get_project(project_id)
        f = await self._run(project.files.get, file_path=file_path, ref=ref)
        return f.attributes if hasattr(f, "attributes") else dict(f)

    async def list_labels(self, project_id: int, **filters):
        project = await self.get_project(project_id)
        return await self._run(project.labels.list, per_page=100, **filters)

    async def create_label(self, project_id: int, name: str, color: str, description: str = ""):
        project = await self.get_project(project_id)
        label = await self._run(project.labels.create, {"name": name, "color": color, "description": description})
        return label.attributes if hasattr(label, "attributes") else dict(label)

    async def reopen_issue(self, project_id: int, issue_iid: int) -> dict:
        return await self.update_issue(project_id, issue_iid, {"state_event": "reopen"})

    async def list_discussions(self, project_id: int, noteable_type: str, noteable_iid: int):
        project = await self.get_project(project_id)
        if noteable_type == "Issue":
            noteable = await self._run(project.issues.get, noteable_iid)
        else:
            noteable = await self._run(project.mergerequests.get, noteable_iid)
        return await self._run(noteable.discussions.list, per_page=100)

    async def trigger_pipeline(self, project_id: int, ref: str, variables: dict | None = None):
        project = await self.get_project(project_id)
        pipeline_data = {"ref": ref}
        if variables:
            pipeline_data["variables"] = [{"key": k, "value": v} for k, v in variables.items()]
        return await self._run(project.pipelines.create, pipeline_data)