"""
Git-based model downloading utilities for Helium
"""
import os
import shutil
import tempfile
import subprocess
from typing import Optional

def git_clone_model(repo_id: str, branch: str = "main", cache_dir: Optional[str] = None) -> str:
    """Clone a model repository using Git
    
    Args:
        repo_id: The model repository ID (e.g. 'openai/gpt-oss-20b')
        branch: Git branch to clone (default: main)
        cache_dir: Optional cache directory for storing models
        
    Returns:
        Path to the cloned repository
    """
    if cache_dir is None:
        cache_dir = os.path.join(os.path.expanduser("~"), ".cache", "helium", "models")
    os.makedirs(cache_dir, exist_ok=True)
    
    # Convert repo_id to valid directory name
    safe_name = repo_id.replace('/', '_')
    repo_dir = os.path.join(cache_dir, safe_name)
    
    # Clear existing directory if it exists
    if os.path.exists(repo_dir):
        shutil.rmtree(repo_dir)
        
    # Construct git URL
    git_url = f"https://huggingface.co/{repo_id}"
    
    try:
        # Clone the repository
        subprocess.run(
            ["git", "clone", "--depth", "1", "-b", branch, git_url, repo_dir],
            check=True,
            capture_output=True
        )
        
        print(f"Successfully cloned model repository to {repo_dir}")
        return repo_dir
        
    except subprocess.CalledProcessError as e:
        print(f"Git clone failed: {e.stderr.decode()}")
        # Clean up on failure
        if os.path.exists(repo_dir):
            shutil.rmtree(repo_dir)
        raise RuntimeError(f"Failed to clone model repository: {e.stderr.decode()}")

def download_lfs_files(repo_dir: str) -> None:
    """Download LFS files for a cloned repository
    
    Args:
        repo_dir: Path to the cloned repository
    """
    try:
        # Initialize Git LFS
        subprocess.run(
            ["git", "lfs", "install"],
            cwd=repo_dir,
            check=True,
            capture_output=True
        )
        
        # Pull LFS files
        subprocess.run(
            ["git", "lfs", "pull"],
            cwd=repo_dir,
            check=True,
            capture_output=True
        )
        
        print("Successfully downloaded LFS files")
        
    except subprocess.CalledProcessError as e:
        print(f"Git LFS operation failed: {e.stderr.decode()}")
        raise RuntimeError(f"Failed to download LFS files: {e.stderr.decode()}")

def validate_model_files(repo_dir: str) -> bool:
    """Validate that all required model files are present
    
    Args:
        repo_dir: Path to the model repository
        
    Returns:
        True if validation passes
    """
    required_files = [
        "config.json",
        "model.safetensors",
        "pytorch_model.bin",
        "tokenizer.json"
    ]
    
    # Check for at least one weight file
    has_weights = False
    for file in os.listdir(repo_dir):
        if file.endswith(('.bin', '.safetensors')):
            has_weights = True
            break
            
    if not has_weights:
        print("No model weight files found")
        return False
        
    # Check for config file
    if not os.path.exists(os.path.join(repo_dir, "config.json")):
        print("No config.json found")
        return False
        
    return True
