package main

import (
	"bufio"
	"io"
	"log"
	"os"
	"path/filepath"
	"regexp"
	"syscall"
)

// The Python stub template that Bazel generates for every py_binary()
// always starts with '#!/usr/bin/env python'. This means that even if
// build actions provide their own copy of Python as part of the input
// root, the remote execution environment must also provide a copy for
// evaluating the stub template.
//
// This tool is a drop-in replacement for /usr/bin/python that is only
// capable of accepting Python stub templates, executing them with the
// copy of Python that is stored in the input root. It obtains the path
// of the Python interpreter by extracting the PYTHON_BINARY constant
// declared in the stub template.

func extractPythonBinary(r io.Reader) string {
	s := bufio.NewScanner(r)
	pattern := regexp.MustCompile("^PYTHON_BINARY\\s*=\\s*'(.*)'$")
	for s.Scan() {
		if match := pattern.FindStringSubmatch(s.Text()); match != nil {
			return match[1]
		}
	}
	if err := s.Err(); err != nil {
		log.Fatal("Failed to read line from Python stub template: ", err)
	}
	log.Fatal("Python stub template does not contain a valid PYTHON_BINARY declaration")
	return ""
}

func main() {
	if len(os.Args) < 2 {
		log.Fatal("Fake Python can only be used to run Python stub templates generated by Bazel")
	}

	// Extract the "PYTHON_BINARY = '...'" line from the Python stub
	// template that was generated by Bazel.
	f, err := os.Open(os.Args[1])
	if err != nil {
		log.Fatalf("Cannot open Python stub template %#v: %s", os.Args[1], err)
	}
	pythonBinary := extractPythonBinary(f)
	f.Close()

	// Determine the path inside the runfiles directory at which the
	// Python interpreter is placed. Tests have RUNFILES_DIR set,
	// while build actions do not. Look up the runfiles directory
	// manually if needed.
	runfilesDir, ok := os.LookupEnv("RUNFILES_DIR")
	if !ok {
		runfilesDir = os.Args[1] + ".runfiles"
	}
	pythonPath := filepath.Join(runfilesDir, pythonBinary)

	// Execute the Python stub template using the copy of Python
	// that it was going to use to run the actual program.
	log.Fatalf(
		"Failed to execute %#v: %s",
		pythonPath,
		syscall.Exec(
			pythonPath,
			append([]string{pythonPath}, os.Args[1:]...),
			os.Environ()))
}