package main

import (
	"archive/zip"
	"flag"
	"fmt"
	"io"
	"os"
	"os/exec"
	"path/filepath"
	"strings"
	"time"
)

var preservePaths = map[string]bool{
	"config/settings.json": true,
	"data":                 true,
}

func shouldPreserve(rel string) bool {
	rel = filepath.ToSlash(rel)
	for p := range preservePaths {
		if rel == p || strings.HasPrefix(rel, p+"/") {
			return true
		}
	}
	return false
}

func isProcessRunning(pid int) bool {
	cmd := exec.Command("tasklist", "/FI", fmt.Sprintf("PID eq %d", pid), "/NH")
	out, err := cmd.Output()
	if err != nil {
		return false
	}
	return strings.Contains(string(out), fmt.Sprintf("%d", pid))
}

func main() {
	pid := flag.Int("pid", 0, "PID of main app to wait for")
	zipPath := flag.String("zip", "", "Path to update zip")
	target := flag.String("target", "", "Install directory")
	launch := flag.String("launch", "", "Executable to launch after update")
	flag.Parse()

	if *zipPath == "" || *target == "" {
		fmt.Println("Usage: updater --pid PID --zip FILE --target DIR --launch EXE")
		os.Exit(1)
	}

	logsDir := filepath.Join(*target, "logs")
	os.MkdirAll(logsDir, 0755)
	logFile, _ := os.Create(filepath.Join(logsDir, "updater.log"))
	log := func(msg string, args ...interface{}) {
		line := fmt.Sprintf("[%s] %s\n", time.Now().Format("15:04:05"), fmt.Sprintf(msg, args...))
		fmt.Print(line)
		if logFile != nil {
			logFile.WriteString(line)
		}
	}

	if *pid > 0 {
		log("Waiting for process %d to exit...", *pid)
		for i := 0; i < 30; i++ {
			if !isProcessRunning(*pid) {
				log("Process exited")
				break
			}
			time.Sleep(500 * time.Millisecond)
		}
		time.Sleep(500 * time.Millisecond)
	}

	log("Extracting %s", *zipPath)
	tmpDir := *zipPath + "_extract"
	os.RemoveAll(tmpDir)
	if err := extractZip(*zipPath, tmpDir); err != nil {
		log("Extract failed: %v", err)
		os.Exit(1)
	}

	// Auto-detect nested top-level directory and strip it
	entries, _ := os.ReadDir(tmpDir)
	if len(entries) == 1 && entries[0].IsDir() {
		nested := filepath.Join(tmpDir, entries[0].Name())
		log("Detected nested directory: %s, stripping", entries[0].Name())
		tmpDir = nested
	}

	log("Installing to %s", *target)
	count := 0
	skipped := 0
	err := filepath.Walk(tmpDir, func(path string, info os.FileInfo, err error) error {
		if err != nil || info.IsDir() {
			return err
		}
		rel, _ := filepath.Rel(tmpDir, path)
		if shouldPreserve(rel) {
			log("  Skip: %s", rel)
			skipped++
			return nil
		}
		dst := filepath.Join(*target, rel)
		os.MkdirAll(filepath.Dir(dst), 0755)

		src, err := os.Open(path)
		if err != nil {
			return err
		}
		defer src.Close()

		out, err := os.Create(dst)
		if err != nil {
			log("  Failed to write: %s (%v)", rel, err)
			return nil
		}
		defer out.Close()
		io.Copy(out, src)
		count++
		return nil
	})
	if err != nil {
		log("Install error: %v", err)
	}
	log("Updated %d files, skipped %d", count, skipped)

	os.RemoveAll(tmpDir)
	os.Remove(*zipPath)
	log("Cleanup done")

	if *launch != "" {
		exe := filepath.Join(*target, *launch)
		log("Launching %s", exe)
		cmd := exec.Command(exe)
		cmd.Dir = *target
		cmd.Start()
	}

	log("Update complete")
	if logFile != nil {
		logFile.Close()
	}
}

func extractZip(src, dst string) error {
	r, err := zip.OpenReader(src)
	if err != nil {
		return err
	}
	defer r.Close()

	for _, f := range r.File {
		path := filepath.Join(dst, f.Name)
		if f.FileInfo().IsDir() {
			os.MkdirAll(path, 0755)
			continue
		}
		os.MkdirAll(filepath.Dir(path), 0755)
		rc, err := f.Open()
		if err != nil {
			return err
		}
		out, err := os.Create(path)
		if err != nil {
			rc.Close()
			return err
		}
		io.Copy(out, rc)
		out.Close()
		rc.Close()
	}
	return nil
}