#!/usr/bin/env python3

import argparse
import json
import os
import shutil
import subprocess
import sys
from pathlib import Path
import jinja2

def get_platform():
    """Get the platform using the platform.sh script"""
    try:
        script_path = os.path.join(os.path.dirname(os.path.realpath(__file__)), "home/scripts/platform.sh")
        result = subprocess.run([script_path], capture_output=True, text=True, check=True)
        return result.stdout.strip()
    except subprocess.CalledProcessError as e:
        print(f"Error getting platform: {e}")
        sys.exit(1)

def get_device_name():
    """Get the device name using the system_name.sh script"""
    try:
        script_path = os.path.join(os.path.dirname(os.path.realpath(__file__)), "home/scripts/system_name.sh")
        result = subprocess.run([script_path], capture_output=True, text=True, check=True)
        return result.stdout.strip()
    except subprocess.CalledProcessError as e:
        print(f"Error getting device name: {e}")
        sys.exit(1)

def load_context(platform, device_name):
    """Load the appropriate context from contexts.json"""
    try:
        with open('contexts.json', 'r') as f:
            contexts = json.load(f)

        if platform not in contexts:
            print(f"Warning: Platform '{platform}' not found in contexts.json")
            return {}

        if device_name not in contexts[platform]:
            print(f"Warning: Device '{device_name}' not found for platform '{platform}' in contexts.json")
            # Try to use 'default' for the platform if device not found
            if 'default' in contexts[platform]:
                print(f"Using default context for platform '{platform}'")
                return contexts[platform]['default']
            return {}

        return contexts[platform][device_name]
    except FileNotFoundError:
        print("Warning: contexts.json not found")
        return {}
    except json.JSONDecodeError as e:
        print(f"Error parsing contexts.json: {e}")
        sys.exit(1)

def compile_dotfiles(source_dir, target_dir, context):
    """Compile dotfiles, processing Jinja templates with context"""
    # Create jinja environment
    env = jinja2.Environment(
        undefined=jinja2.StrictUndefined,
        trim_blocks=True,
        lstrip_blocks=True
    )

    # Ensure target directory exists
    target_path = Path(target_dir)
    target_path.mkdir(exist_ok=True, parents=True)

    # Process each file in source directory
    for root, dirs, files in os.walk(source_dir):
        # Skip .git directories
        if '.git' in dirs:
            dirs.remove('.git')

        # Create relative path from source_dir
        rel_path = os.path.relpath(root, source_dir)
        if rel_path == '.':
            rel_path = ''

        # Create target directory
        if rel_path:
            target_subdir = target_path / rel_path
            target_subdir.mkdir(exist_ok=True, parents=True)
        else:
            target_subdir = target_path

        for file in files:
            source_file = os.path.join(root, file)

            # Determine target filename (remove .j2 extension for templates)
            target_file_name = file[:-3] if file.endswith('.j2') else file
            target_file = target_subdir / target_file_name

            print(f"Processing: {source_file} -> {target_file}")

            if file.endswith('.j2'):
                # Render Jinja2 template
                try:
                    with open(source_file, 'r') as f:
                        template_content = f.read()

                    # Use the environment to create templates
                    template = env.from_string(template_content)
                    rendered_content = template.render(**context)

                    # Write rendered content to target file
                    with open(target_file, 'w') as f:
                        f.write(rendered_content)

                    # Make executable if source is executable
                    if os.access(source_file, os.X_OK):
                        os.chmod(target_file, 0o755)

                except Exception as e:
                    print(f"Error rendering template {source_file}: {e}")
            else:
                # Copy file as-is
                shutil.copy2(source_file, target_file)

def stow_dotfiles(dotfiles_dir, target_dir=None):
    """Use GNU Stow to symlink the compiled dotfiles"""
    if target_dir is None:
        target_dir = os.path.expanduser("~")

    # Check if stow is installed
    try:
        subprocess.run(["stow", "--version"], capture_output=True, check=True)
    except (subprocess.CalledProcessError, FileNotFoundError):
        print("Error: GNU Stow not found. Please install it before using --stow.")
        sys.exit(1)

    # Get list of directories in dotfiles_dir (each is a stow package)
    packages = [d for d in os.listdir(dotfiles_dir) if os.path.isdir(os.path.join(dotfiles_dir, d))]

    for package in packages:
        print(f"Stowing package: {package}")
        try:
            # Use --adopt to replace existing files
            # Use --no-folding to enable leaf mode (each file individually linked)
            subprocess.run([
                "stow",
                "--dir=" + dotfiles_dir,
                "--target=" + target_dir,
                "--adopt",
                "--no-folding",
                package
            ], check=True)
            print(f"Successfully stowed {package}")
        except subprocess.CalledProcessError as e:
            print(f"Error stowing {package}: {e}")

def main():
    parser = argparse.ArgumentParser(description="Dotfiles manager")
    parser.add_argument("--compile", action="store_true", help="Compile dotfiles")
    parser.add_argument("--stow", action="store_true", help="Stow compiled dotfiles")
    parser.add_argument("--target", help="Target directory for stow (default: $HOME)")
    args = parser.parse_args()

    if not args.compile and not args.stow:
        parser.print_help()
        sys.exit(1)

    platform = get_platform()
    device_name = get_device_name()
    print(f"Platform: {platform}, Device: {device_name}")

    context = load_context(platform, device_name)
    print(f"Loaded context: {context}")

    # Add platform and device_name to context
    context['platform'] = platform
    context['device_name'] = device_name

    compiled_dir = ".compiled_dotfiles"

    if args.compile:
        print(f"Compiling dotfiles from 'dotfiles' to '{compiled_dir}'...")
        compile_dotfiles("dotfiles", compiled_dir, context)
        print("Compilation complete.")

    if args.stow:
        target_dir = args.target if args.target else None
        print(f"Stowing dotfiles from '{compiled_dir}'...")
        stow_dotfiles(compiled_dir, target_dir)
        print("Stowing complete.")

if __name__ == "__main__":
    main()