package autoupdate

import (
	"archive/tar"
	"compress/gzip"
	"fmt"
	"io"
	"io/ioutil"
	"net/http"
	"os"
	"os/exec"
	"path/filepath"
)

func BuildNewBinary(tarballURL string) (buildDir string, newBinPath string, err error) {
	fmt.Printf("Downloading tarball: %s\n", tarballURL)
	resp, err := http.Get(tarballURL)
	if err != nil {
		return "", "", fmt.Errorf("failed to download tarball: %w", err)
	}
	defer resp.Body.Close()

	tmpTarball, err := ioutil.TempFile("", "source_*.tar.gz")
	if err != nil {
		return "", "", fmt.Errorf("failed to create temp file: %w", err)
	}
	_, err = io.Copy(tmpTarball, resp.Body)
	if err != nil {
		tmpTarball.Close()
		return "", "", fmt.Errorf("failed to write tarball: %w", err)
	}
	tmpTarball.Close()
	defer os.Remove(tmpTarball.Name())

	tmpDir, err := ioutil.TempDir("", "build_")
	if err != nil {
		return "", "", fmt.Errorf("failed to create temp directory: %w", err)
	}

	fmt.Printf("Unpacking tarball to: %s\n", tmpDir)
	if err := unpackTarball(tmpTarball.Name(), tmpDir); err != nil {
		os.RemoveAll(tmpDir)
		return "", "", fmt.Errorf("failed to unpack tarball: %w", err)
	}

	goModPath := filepath.Join(tmpDir, "go.mod")
	if _, err := os.Stat(goModPath); os.IsNotExist(err) {
		entries, err := ioutil.ReadDir(tmpDir)
		if err != nil {
			os.RemoveAll(tmpDir)
			return "", "", fmt.Errorf("failed to list unpacked directory: %w", err)
		}
		var subDir string
		count := 0
		for _, entry := range entries {
			if entry.IsDir() {
				subDir = filepath.Join(tmpDir, entry.Name())
				count++
			}
		}
		if count == 1 {
			goModPath = filepath.Join(subDir, "go.mod")
			if _, err := os.Stat(goModPath); err != nil {
				os.RemoveAll(tmpDir)
				return "", "", fmt.Errorf("go.mod not found in subdirectory")
			}
			buildDir = subDir
		} else {
			os.RemoveAll(tmpDir)
			return "", "", fmt.Errorf("go.mod not found in unpacked directory")
		}
	} else {
		buildDir = tmpDir
	}

	newBinPath = filepath.Join(buildDir, "train")
	fmt.Println("go.mod found, building using go build...")
	cmd := exec.Command("go", "build", "-o", "train")
	cmd.Dir = buildDir
	cmd.Stdout = os.Stdout
	cmd.Stderr = os.Stderr
	if err := cmd.Run(); err != nil {
		os.RemoveAll(tmpDir)
		return "", "", fmt.Errorf("go build failed: %w", err)
	}
	fmt.Println("Build completed successfully using go build")
	return buildDir, newBinPath, nil
}

func unpackTarball(archivePath, destDir string) error {
	f, err := os.Open(archivePath)
	if err != nil {
		return err
	}
	defer f.Close()
	gzr, err := gzip.NewReader(f)
	if err != nil {
		return err
	}
	defer gzr.Close()
	tr := tar.NewReader(gzr)
	for {
		header, err := tr.Next()
		if err == io.EOF {
			break
		}
		if err != nil {
			return err
		}
		target := filepath.Join(destDir, header.Name)
		switch header.Typeflag {
		case tar.TypeDir:
			if err := os.MkdirAll(target, os.FileMode(header.Mode)); err != nil {
				return err
			}
		case tar.TypeReg:
			outFile, err := os.Create(target)
			if err != nil {
				return err
			}
			if _, err := io.Copy(outFile, tr); err != nil {
				outFile.Close()
				return err
			}
			outFile.Close()
		}
	}
	return nil
}