import os
import shutil
import subprocess
import sys
import tempfile
from pathlib import Path

# Load environment variables manually
def load_env():
    env_path = Path(".env")
    if env_path.exists():
        with open(env_path, "r") as f:
            for line in f:
                line = line.strip()
                if not line or line.startswith("#"):
                    continue
                if "=" in line:
                    key, value = line.split("=", 1)
                    # Strip quotes if present
                    if (value.startswith('"') and value.endswith('"')) or \
                       (value.startswith("'") and value.endswith("'")):
                        value = value[1:-1]
                    os.environ[key] = value

load_env()

HF_TOKEN = os.getenv("HF_TOKEN")
HF_ORG = "AmaniQuery"

if not HF_TOKEN:
    print("Error: HF_TOKEN not found in .env file or environment variables.")
    sys.exit(1)

def run_command(command, cwd=None, env=None):
    try:
        updated_env = os.environ.copy()
        if env:
            updated_env.update(env)
            
        result = subprocess.run(
            command,
            cwd=cwd,
            env=updated_env,
            check=True,
            shell=True,
            capture_output=True,
            text=True
        )
        return result.stdout.strip()
    except subprocess.CalledProcessError as e:
        print(f"Error running command: {command}")
        print(f"Output: {e.stdout}")
        print(f"Error: {e.stderr}")
        raise

def deploy_service(service_type):
    # Configuration based on service type
    if service_type == "agent":
        repo_name = "amaniquery-agent"
        dockerfile_src = "deployments/huggingface/Dockerfile.hf"
        readme_src = "deployments/huggingface/README.md"
        context_path = "."
    elif service_type == "memory":
        repo_name = "amaniquery-memory"
        dockerfile_src = "deployments/huggingface/Dockerfile.rust.hf"
        readme_src = "deployments/huggingface/README.md" # We might need a specific one or just use default and edit
        context_path = "." # Rust service needs root context if using workspace
    else:
        print(f"Unknown service type: {service_type}")
        return

    print(f"Deploying {service_type} to https://huggingface.co/{HF_ORG}/{repo_name}...")

    # Create temp directory for cloning
    with tempfile.TemporaryDirectory() as temp_dir:
        repo_url = f"https://oauth2:{HF_TOKEN}@huggingface.co/spaces/{HF_ORG}/{repo_name}"
        
        print(f"Cloning {repo_name}...")
        try:
            run_command(f"git clone {repo_url} .", cwd=temp_dir)
        except Exception as e:
            print("Failed to clone. Make sure the Space exists first!")
            return

        # Configure git user
        run_command("git config user.email 'deploy-script@amaniquery.com'", cwd=temp_dir)
        run_command("git config user.name 'Deployment Script'", cwd=temp_dir)

        # Copy files respecting gitignore (using git ls-files to get list of tracked files)
        print("Copying files...")
        project_root = Path.cwd()
        
        # Get list of files tracked by git in current repo
        tracked_files = run_command("git ls-files", cwd=project_root).splitlines()
        
        for file_path in tracked_files:
            src = project_root / file_path
            dst = Path(temp_dir) / file_path
            
            # Skip deployments folder (handle separately to avoid overwriting custom setup)
            if file_path.startswith("deployments"):
                continue

            # Skip architectural diagrams and images
            if file_path.endswith(".png") or file_path.endswith(".jpg") or file_path.endswith(".jpeg"):
                continue
                
            dst.parent.mkdir(parents=True, exist_ok=True)
            if src.exists():
                shutil.copy2(src, dst)

        # Copy specific deployment artifacts
        print("Configuring deployment artifacts...")
        shutil.copy2(project_root / dockerfile_src, Path(temp_dir) / "Dockerfile")
        
        # Always overwrite README to ensure correct configuration
        shutil.copy2(project_root / readme_src, Path(temp_dir) / "README.md")

        # Explicitly copy start.sh (needed for runtime, might be untracked)
        start_sh_src = project_root / "deployments/huggingface/start.sh"
        start_sh_dst = Path(temp_dir) / "deployments/huggingface/start.sh"
        if start_sh_src.exists():
             start_sh_dst.parent.mkdir(parents=True, exist_ok=True)
             shutil.copy2(start_sh_src, start_sh_dst)

        for wfile in ["go.work", "go.work.sum"]:
            src = project_root / wfile
            if src.exists():
                shutil.copy2(src, Path(temp_dir) / wfile)
                # Force add because it might be in .gitignore
                run_command(f"git add -f {wfile}", cwd=temp_dir)

        # Explicitly copy rust-memory-service including Cargo.lock (which might be ignored or missed)
        rust_src = project_root / "rust-memory-service"
        rust_dst = Path(temp_dir) / "rust-memory-service"
        if rust_src.exists():
            # Remove destination if exists (from git ls-files copy) to ensure clean full copy
            if rust_dst.exists():
                shutil.rmtree(rust_dst)
            shutil.copytree(rust_src, rust_dst, ignore=shutil.ignore_patterns("target", ".git"))
            # Force add to ensure ignored files (like Cargo.lock if ignored) are included
            run_command("git add -f rust-memory-service", cwd=temp_dir)

        # Explicitly copy generated protobuf files (not tracked in git)
        proto_gen_src = project_root / "pkg/proto/gen"
        proto_gen_dst = Path(temp_dir) / "pkg/proto/gen"
        if proto_gen_src.exists():
            if proto_gen_dst.exists():
                shutil.rmtree(proto_gen_dst)
            proto_gen_dst.parent.mkdir(parents=True, exist_ok=True)
            shutil.copytree(proto_gen_src, proto_gen_dst)
            run_command("git add -f pkg/proto/gen", cwd=temp_dir)
        
        # Commit and push
        print("Pushing changes...")
        run_command("git add .", cwd=temp_dir)
        status = run_command("git status --porcelain", cwd=temp_dir)
        
        if status:
            run_command('git commit -m "Automated deployment update"', cwd=temp_dir)
            run_command("git push", cwd=temp_dir)
            print(f"Successfully deployed {service_type}!")
        else:
            print("No changes to deploy.")

if __name__ == "__main__":
    if len(sys.argv) < 2:
        print("Usage: python deploy_hf.py [agent|memory|all]")
        sys.exit(1)
        
    target = sys.argv[1]
    
    if target == "all":
        deploy_service("agent")
        deploy_service("memory")
    else:
        deploy_service(target)