from pydantic import BaseModel, Field
from crewai_tools import BaseTool
import git
import os
import shutil
from typing import Type


class GitCloneInput(BaseModel):
    """Input schema for GitTool."""
    repo_url: str = Field(..., description="URL of the Git repository to clone")
    branch: str = Field(None, description="Branch to checkout (optional)")
    commit: str = Field(None, description="Commit hash to checkout (optional)")
    target_dir: str = Field(None, description="Target directory to clone into (optional)")


class GitTool(BaseTool):
    name: str = "GitTool"
    description: str = "Clones a Git repository and checks out a specific branch or commit"
    args_schema: Type[BaseModel] = GitCloneInput

    def _run(self, repo_url: str, branch: str = None, commit: str = None, target_dir: str = None) -> str:
        """
        Clone a Git repository and checkout a specific branch or commit.

        Args:
            repo_url: URL of the Git repository to clone
            branch: Branch to checkout (optional)
            commit: Commit hash to checkout (optional)
            target_dir: Target directory to clone into (optional)

        Returns:
            A message indicating the result of the operation.
        """
        # If target_dir is not provided, extract the repository name from the URL
        if target_dir is None:
            target_dir = repo_url.split("/")[-1]
            if target_dir.endswith(".git"):
                target_dir = target_dir[:-4]

        # Remove the target directory if it already exists to avoid conflicts
        if os.path.exists(target_dir):
            shutil.rmtree(target_dir)

        try:
            # Clone the repository
            repo = git.Repo.clone_from(repo_url, target_dir)

            # Checkout the specified branch or commit
            if branch:
                repo.git.checkout(branch)
            elif commit:
                repo.git.checkout(commit)

            return f"Successfully cloned {repo_url} into {target_dir} and checked out {'branch: ' + branch if branch else 'commit: ' + commit if commit else 'default branch'}"
        except Exception as e:
            return f"Error cloning repository: {str(e)}"