package runner

import (
	"fmt"
	"git.0x7f.app/WOJ/woj-server/internal/e"
	"git.0x7f.app/WOJ/woj-server/pkg/file"
	"git.0x7f.app/WOJ/woj-server/pkg/utils"
	"go.uber.org/zap"
	"os"
	"path/filepath"
	"strings"
	"time"
)

type ProblemRunResult struct {
	QueueId uint64
	Status  RuntimeStatus
}

type ProblemRunResults map[int]*ProblemRunResult

func (s *service) SandboxArgsBuilder(meta *JudgeMeta, id int) string {
	var args []string

	args = append(args, fmt.Sprintf("--memory_limit=%d", meta.Cfg.Lang.Runtime.Run.MemoryLimit))
	args = append(args, fmt.Sprintf("--nproc_limit=%d", meta.Cfg.Lang.Runtime.Run.NProcLimit))
	args = append(args, fmt.Sprintf("--time_limit=%d", meta.Cfg.Lang.Runtime.Run.TimeLimit))
	args = append(args, fmt.Sprintf("--sandbox_template=%s", meta.Cfg.Lang.Lang))
	args = append(args, fmt.Sprintf("--sandbox_action=%s", "ret"))
	args = append(args, fmt.Sprintf("--uid=%d", 1000))
	args = append(args, fmt.Sprintf("--gid=%d", 1000))
	args = append(args, fmt.Sprintf("--file_input=/woj/problem/data/input/%d.input", id))
	args = append(args, fmt.Sprintf("--file_output=/woj/user/%d.out.usr", id))
	args = append(args, fmt.Sprintf("--file_info=/woj/user/%d.info", id))

	if meta.Cfg.Lang.JudgeInterpreter() != "" {
		sourceFile := fmt.Sprintf("%s.%s", meta.Run.User, meta.Run.Lang)

		args = append(args, fmt.Sprintf("--program=%s", meta.Cfg.Lang.JudgeInterpreter()))
		args = append(args, fmt.Sprintf("--program_arg=/woj/user/%s", sourceFile))
	} else {
		args = append(args, fmt.Sprintf("--program=/woj/user/%s.out", meta.Run.User))
	}

	s.log.Debug("[run] sandbox args", zap.Strings("args", args))

	return strings.Join(args, " ")
}

func (s *service) ProblemRun(meta *JudgeMeta) ProblemRunResults {
	workDir := filepath.Join(UserDir, meta.Run.User)
	dataDir := filepath.Join(ProblemDir, fmt.Sprintf("%d", meta.Run.Version), "data", "input")

	runtimeArgs := RuntimeArgs{
		Rootfs: RootfsRunDir,
		// sh, woj_launcher:program, woj_launcher:killer, woj_launcher:stat
		Pid:    int64(meta.Cfg.Lang.Runtime.Run.NProcLimit + 4),
		Memory: uint64(meta.Cfg.Lang.Runtime.Run.MemoryLimit * 1024 * 1024),
		// woj-sandbox killer will add 1 more second, here we add 1 also
		Timeout: time.Duration((meta.Cfg.Lang.Runtime.Run.TimeLimit+1000)/1000+1+1) * time.Second,
	}

	result := make(ProblemRunResults)
	for _, task := range meta.Cfg.All.Tasks {
		f := func(id int) func() (interface{}, error) {
			return func() (interface{}, error) {
				testCase := filepath.Join(dataDir, fmt.Sprintf("%d.input", id))
				ansFile := filepath.Join(workDir, fmt.Sprintf("%d.out.usr", id))
				ifoFile := filepath.Join(workDir, fmt.Sprintf("%d.info", id))

				targetFile := fmt.Sprintf("%s.out", meta.Run.User)
				targetPath := filepath.Join(workDir, targetFile)
				sourceFile := fmt.Sprintf("%s.%s", meta.Run.User, meta.Run.Lang)
				sourcePath := filepath.Join(workDir, sourceFile)

				args := &RunArgs{
					Program: ProgramArgs{
						Args: []string{
							"/bin/sh", "-c",
							"cd /woj/user && /woj/framework/scripts/woj_launcher " +
								s.SandboxArgsBuilder(meta, id),
						},
					},
					Runtime: runtimeArgs,
				}
				args.Runtime.Mount = []MountInfo{
					{
						Source:      testCase,
						Destination: fmt.Sprintf("/woj/problem/data/input/%d.input", id),
						Readonly:    true,
					},
					{
						Source:      utils.If(meta.Cfg.Lang.JudgeInterpreter() != "", sourcePath, targetPath),
						Destination: fmt.Sprintf("/woj/user/%s", utils.If(meta.Cfg.Lang.JudgeInterpreter() != "", sourceFile, targetFile)),
						Readonly:    true,
					},
					{
						Source:      ansFile,
						Destination: fmt.Sprintf("/woj/user/%d.out.usr", id),
						Readonly:    false,
					},
					{
						Source:      ifoFile,
						Destination: fmt.Sprintf("/woj/user/%d.info", id),
						Readonly:    false,
					},
				}

				err := utils.NewMust().
					DoAny(func() error { return os.Remove(ansFile) }).
					DoAny(func() error { return os.Remove(ifoFile) }).
					Do(func() error { return file.TouchErr(ansFile) }).
					Do(func() error { return file.TouchErr(ifoFile) }).
					Done()

				if err != nil {
					s.log.Info("[run] prepare failed", zap.Error(err), zap.Any("meta", *meta))
					return nil, err
				}

				return s.JailRun(args)
			}
		}(task.Id)

		queueId := s.pool.AddTask(f)
		result[task.Id] = &ProblemRunResult{
			QueueId: queueId,
		}
	}

	for i := range result {
		waitBuf := s.pool.WaitForTask(result[i].QueueId)
		if waitBuf.Error != nil {
			s.log.Error(
				"[run] wait for problem run failed",
				zap.Error(waitBuf.Error),
				zap.Any("meta", *meta))
			continue
		}
		val, ok := waitBuf.Value.(RuntimeStatus)
		if !ok {
			s.log.Error(
				"[run] container run is not returning RuntimeStatus",
				zap.Any("waitBuf", waitBuf),
				zap.Any("meta", *meta))
			continue
		}

		result[i].Status = val
	}

	return result
}

func (s *service) ProblemJudge(meta *JudgeMeta) {
	workDir := filepath.Join(UserDir, meta.Run.User)
	dataDir := filepath.Join(ProblemDir, fmt.Sprintf("%d", meta.Run.Version), "data")
	judgeDir := filepath.Join(ProblemDir, fmt.Sprintf("%d", meta.Run.Version), "judge")
	script := meta.Cfg.Lang.JudgeScript()

	runtimeArgs := RuntimeArgs{
		Rootfs:  RootfsFullDir,
		Pid:     int64(meta.Cfg.Lang.Runtime.Check.NProcLimit + 2), // bash + make
		Memory:  uint64(meta.Cfg.Lang.Runtime.Check.MemoryLimit * 1024 * 1024),
		Timeout: time.Duration((meta.Cfg.Lang.Runtime.Check.TimeLimit+1000)/1000) * time.Second,
	}

	ids := make([]uint64, 0)
	for _, task := range meta.Cfg.All.Tasks {
		f := func(id int) func() (interface{}, error) {
			return func() (interface{}, error) {
				ansFile := filepath.Join(workDir, fmt.Sprintf("%d.out.usr", id))
				jdgFile := filepath.Join(workDir, fmt.Sprintf("%d.judge", id))

				args := &RunArgs{
					Program: ProgramArgs{
						Args: []string{
							"/bin/sh", "-c",
							fmt.Sprintf("cd /woj/user && make -f %s judge", script),
						},
						Env: []string{
							fmt.Sprintf("TEST_NUM=%d", id),
							fmt.Sprintf("CMP=%s", meta.Cfg.Lang.Judge.Cmp),
						},
					},
					Runtime: runtimeArgs,
				}
				args.Runtime.Mount = []MountInfo{
					{
						Source:      judgeDir,
						Destination: "/woj/problem/judge",
						Readonly:    true,
					},
					{
						Source:      dataDir,
						Destination: "/woj/problem/data",
						Readonly:    true,
					},
					{
						Source:      ansFile,
						Destination: fmt.Sprintf("/woj/user/%d.out.usr", id),
						Readonly:    true,
					},
					{
						Source:      jdgFile,
						Destination: fmt.Sprintf("/woj/user/%d.judge", id),
						Readonly:    false,
					},
				}

				err := utils.NewMust().
					DoAny(func() error { return os.Remove(jdgFile) }).
					Do(func() error { return file.TouchErr(jdgFile) }).
					Done()

				if err != nil {
					s.log.Info("[judge] judge prepare failed", zap.Error(err), zap.Any("meta", *meta))
					return nil, err
				}

				return s.JailRun(args)
			}
		}(task.Id)

		id := s.pool.AddTask(f)
		ids = append(ids, id)
	}

	for _, id := range ids {
		_ = s.pool.WaitForTask(id)
	}
}

func (s *service) RunAndJudge(meta *JudgeMeta) (*JudgeStatus, int32, e.Status) {
	// 1. run user program
	results := s.ProblemRun(meta)

	// 2. run judge
	s.ProblemJudge(meta)

	// 3. final JudgeStatus
	status, pts := s.CheckResults(meta, results)

	return status, pts, e.Success
}