#!/usr/bin/env python3
import os
import subprocess
import platform
import sys
import tempfile
import shutil
from pathlib import Path
import argparse

def run_command(command, shell=True, check=True):
    """Run a shell command and print output"""
    print(f"Running: {command}")
    try:
        result = subprocess.run(
            command, 
            shell=shell, 
            check=check,
            stdout=subprocess.PIPE, 
            stderr=subprocess.PIPE,
            text=True
        )
        print(result.stdout)
        if result.stderr:
            print(f"Error: {result.stderr}")
        return True
    except subprocess.CalledProcessError as e:
        print(f"Command failed with return code {e.returncode}")
        print(e.stderr)
        return False

def install_ubuntu_dependencies():
    """Install dependencies for Ubuntu systems"""
    # Get the distro and codename of the Ubuntu distribution
    try:
        # Try with explicit path first
        paths_to_try = ["/usr/bin/lsb_release", "lsb_release"]
        
        for lsb_cmd in paths_to_try:
            try:
                distro = subprocess.check_output(
                    [lsb_cmd, "--id", "--short", "|", "tr", 'A-Z', 'a-z'], 
                    text=True,
                    stderr=subprocess.DEVNULL
                ).strip()
                
                codename = subprocess.check_output(
                    [lsb_cmd, "-cs"], 
                    text=True,
                    stderr=subprocess.DEVNULL
                ).strip()
                
                # If we get here, the command succeeded
                break
            except (subprocess.SubprocessError, FileNotFoundError):
                continue
        else:  # This else belongs to the for loop, executes if no break occurred
            raise FileNotFoundError("Could not run lsb_release successfully")
            
    except Exception as e:
        print(f"Could not determine Ubuntu distro/codename: {e}")
        print("Defaulting to 'ubuntu/jammy'")
        distro = "ubuntu"
        codename = "jammy"
    
    print(f"Detected Ubuntu distro: {distro}, codename: {codename}")
    
    # Create a temporary directory for downloads
    with tempfile.TemporaryDirectory() as tmp_dir:
        # Download the Arrow APT source
        arrow_deb = f"apache-arrow-apt-source-latest-{codename}.deb"
        arrow_url = f"https://apache.jfrog.io/artifactory/arrow/{distro}/{arrow_deb}"
        
        # Download the .deb file
        download_cmd = f"wget {arrow_url} -P {tmp_dir}"
        if not run_command(download_cmd):
            print("Failed to download Apache Arrow APT source")
            return False
        
        # Install the .deb file
        deb_path = os.path.join(tmp_dir, arrow_deb)
        install_cmd = f"sudo apt install -y {deb_path}"
        if not run_command(install_cmd):
            print("Failed to install Apache Arrow APT source")
            return False
        
        # Update package lists
        if not run_command("sudo apt update"):
            print("Failed to update package lists")
            return False
        
        # Install Arrow C++ libraries and development files
        arrow_cmd = "sudo apt install -y libarrow-dev libarrow-dataset-dev libparquet-dev"
        if not run_command(arrow_cmd):
            print("Failed to install Arrow C++ libraries")
            return False
        
        # Install additional dependencies
        deps_cmd = "sudo apt install -y liblz4-dev libzstd-dev libre2-dev libthrift-dev"
        if not run_command(deps_cmd):
            print("Failed to install additional dependencies")
            return False
    
    print("Successfully installed all Ubuntu dependencies")
    return True

def install_mac_dependencies():
    """Install dependencies for macOS systems using Homebrew"""
    # Check if Homebrew is installed
    if shutil.which("brew") is None:
        print("Homebrew is not installed. Please install Homebrew first:")
        print("/bin/bash -c \"$(curl -fsSL https://raw.githubusercontent.com/Homebrew/install/HEAD/install.sh)\"")
        return False
    
    # Install Apache Arrow and dependencies
    commands = [
        "brew install apache-arrow",
        "brew install lz4",
        "brew install zstd",
        "brew install re2",
        "brew install thrift",
        "brew install libomp"  # For OpenMP support on macOS
    ]
    
    for cmd in commands:
        if not run_command(cmd):
            print(f"Failed to run: {cmd}")
            print("Continuing with remaining installations...")
    
    print("Successfully installed all macOS dependencies")
    return True

def parse_args():
    parser = argparse.ArgumentParser(description='Install dependencies for jsonl-parser')
    parser.add_argument('--force-platform', choices=['ubuntu', 'mac'], 
                       help='Force installation for a specific platform')
    return parser.parse_args()

def main():
    args = parse_args()
    
    if args.force_platform:
        platform_name = args.force_platform
    else:
        # Determine the operating system
        system = platform.system()
        if system == "Linux":
            # Check if it's Ubuntu
            if os.path.exists("/etc/lsb-release"):
                with open("/etc/lsb-release") as f:
                    if "Ubuntu" in f.read():
                        platform_name = "ubuntu"
                    else:
                        platform_name = "linux"
            else:
                platform_name = "linux"
        elif system == "Darwin":
            platform_name = "mac"
        else:
            platform_name = system.lower()
    
    # Install dependencies based on the platform
    if platform_name == "ubuntu":
        print("Installing dependencies for Ubuntu...")
        if install_ubuntu_dependencies():
            print("All Ubuntu dependencies installed successfully!")
        else:
            sys.exit(1)
    elif platform_name == "mac":
        print("Installing dependencies for macOS...")
        if install_mac_dependencies():
            print("All macOS dependencies installed successfully!")
        else:
            sys.exit(1)
    else:
        print(f"Unsupported platform: {platform_name}")
        print("This script currently supports Ubuntu and macOS only.")
        sys.exit(1)
    
    print("\nInstallation complete!")
    print("You can now build the parser with the build_cpp_parser.py script.")

if __name__ == "__main__":
    main()