from typing import List
import llm
import os
import pathlib
import subprocess
import tempfile
import shutil


@llm.hookimpl
def register_fragment_loaders(register):
    register("repomix", repomix_loader)


def repomix_loader(argument: str) -> List[llm.Fragment]:
    """
    Load repository contents as fragments using Repomix
    
    Argument is a git repository URL (https:// or ssh://)
    Examples:
        repomix:https://git.sr.ht/~amolith/willow
        repomix:ssh://git.sr.ht:~amolith/willow
    """
    if not argument.startswith(("https://", "ssh://", "git@")):
        raise ValueError(
            f"Repository URL must start with https://, ssh://, or git@ - got: {argument}"
        )
    
    # Check if repomix is available
    if not shutil.which("repomix"):
        raise ValueError(
            "repomix command not found. Please install repomix first: "
            "https://github.com/yamadashy/repomix"
        )
    
    # Create a temporary directory for the cloned repository
    with tempfile.TemporaryDirectory() as temp_dir:
        repo_path = pathlib.Path(temp_dir) / "repo"
        
        try:
            # Clone the repository
            subprocess.run(
                ["git", "clone", "--depth=1", argument, str(repo_path)],
                check=True,
                capture_output=True,
                text=True,
            )
            
            # Run repomix on the cloned repository
            result = subprocess.run(
                ["repomix", "--stdout", str(repo_path)],
                check=True,
                capture_output=True,
                text=True,
            )
            
            # Create a single fragment with the repomix output
            fragments = [
                llm.Fragment(
                    content=result.stdout,
                    source=f"repomix:{argument}"
                )
            ]
            
            return fragments
            
        except subprocess.CalledProcessError as e:
            # Handle Git or repomix errors
            if "git" in str(e.cmd):
                raise ValueError(
                    f"Failed to clone repository {argument}: {e.stderr}"
                )
            elif "repomix" in str(e.cmd):
                raise ValueError(
                    f"Failed to run repomix on {argument}: {e.stderr}"
                )
            else:
                raise ValueError(
                    f"Command failed: {e.stderr}"
                )
        except Exception as e:
            # Handle other errors
            raise ValueError(
                f"Error processing repository {argument}: {str(e)}"
            )