package sandbox

import (
	"archive/tar"
	"bytes"
	"context"
	"fmt"
	"io"
	"log/slog"
	"maps"
	"os"
	"path/filepath"

	"github.com/criyle/go-judge/pb"
	"github.com/joint-online-judge/JOJ3/internal/stage"
	"google.golang.org/protobuf/proto"
)

const (
	tarStreamThreshold = 128 * 1024
	streamChunkSize    = 32 * 1024
)

func (e *Sandbox) Run(cmds []stage.Cmd) ([]stage.ExecutorResult, error) {
	var err error
	if e.execClient == nil {
		slog.Debug("create exec client", "server", e.execServer)
		e.execClient, err = createExecClient(e.execServer, e.token)
		if err != nil {
			return nil, err
		}
	}
	// cannot use range loop since we need to change the value
	for i := 0; i < len(cmds); i += 1 {
		cmd := &cmds[i]
		if cmd.CopyIn == nil {
			cmd.CopyIn = make(map[string]stage.CmdFile)
		}
		for k, v := range cmd.CopyInCached {
			if fileID, ok := e.cachedMap[v]; ok {
				cmd.CopyIn[k] = stage.CmdFile{FileID: &fileID}
			}
		}
	}
	if needStream, tarData := prepareTarStream(cmds); needStream {
		return e.runWithTarStream(cmds, tarData)
	}
	return e.runUnary(cmds)
}

func prepareTarStream(cmds []stage.Cmd) (bool, []byte) {
	if len(cmds) == 0 {
		return false, nil
	}
	for i := range cmds {
		if cmds[i].CopyInDir != "" &&
			estimateCopyInSize(&cmds[i]) >= tarStreamThreshold {
			tarData, err := createCopyInTar(&cmds[i])
			if err != nil {
				slog.Error("create copyIn tar", "error", err)
				return false, nil
			}
			for j := range cmds {
				cmds[j].CopyInDir = ""
			}
			return true, tarData
		}
	}
	return false, nil
}

func (e *Sandbox) runUnary(cmds []stage.Cmd) ([]stage.ExecutorResult, error) {
	pbCmds := convertPBCmd(cmds)
	for i, pbCmd := range pbCmds {
		slog.Debug("sandbox execute", "i", i, "pbCmd size", proto.Size(pbCmd))
	}
	pbReq := &pb.Request{}
	pbReq.SetCmd(pbCmds)
	slog.Debug("sandbox execute", "pbReq size", proto.Size(pbReq))
	pbRet, err := e.execClient.Exec(context.TODO(), pbReq)
	if err != nil {
		return nil, err
	}
	if pbRet.GetError() != "" {
		return nil, fmt.Errorf("sandbox execute error: %s", pbRet.GetError())
	}
	results := convertPBResult(pbRet.GetResults())
	for _, result := range results {
		maps.Copy(e.cachedMap, result.FileIDs)
	}
	return results, nil
}

func (e *Sandbox) runWithTarStream(cmds []stage.Cmd, tarData []byte) ([]stage.ExecutorResult, error) {
	tarCmd := stage.Cmd{
		// FIXME: `/w` won't work if go-judge change default working directory
		Args:         []string{"/bin/tar", "xf", "-", "-C", "/w"},
		Env:          cmds[0].Env,
		Stdin:        &stage.CmdFile{StreamIn: true},
		CPULimit:     cmds[0].CPULimit,
		ClockLimit:   cmds[0].ClockLimit,
		MemoryLimit:  cmds[0].MemoryLimit,
		StackLimit:   cmds[0].StackLimit,
		ProcLimit:    cmds[0].ProcLimit,
		CPURateLimit: cmds[0].CPURateLimit,
		CPUSetLimit:  cmds[0].CPUSetLimit,
	}

	allCmds := append([]stage.Cmd{tarCmd}, cmds...)
	pbCmds := convertPBCmd(allCmds)

	pbReq := &pb.Request{}
	pbReq.SetCmd(pbCmds)
	slog.Debug(
		"sandbox stream execute",
		"pbReq size", proto.Size(pbReq),
		"tarSize", len(tarData),
	)

	stream, err := e.execClient.ExecStream(context.TODO())
	if err != nil {
		return nil, fmt.Errorf("exec stream: %w", err)
	}

	sr := &pb.StreamRequest{}
	sr.SetExecRequest(pbReq)
	if err := stream.Send(sr); err != nil {
		return nil, fmt.Errorf("stream send request: %w", err)
	}

	for offset := 0; offset < len(tarData); offset += streamChunkSize {
		end := offset + streamChunkSize
		if end > len(tarData) {
			end = len(tarData)
		}
		input := pb.StreamRequest_builder{
			ExecInput: (&pb.StreamRequest_Input_builder{
				Index:   0,
				Fd:      0,
				Content: tarData[offset:end],
			}).Build(),
		}.Build()
		if err := stream.Send(input); err != nil {
			return nil, fmt.Errorf("stream send input: %w", err)
		}
	}

	if err := stream.CloseSend(); err != nil {
		return nil, fmt.Errorf("stream close send: %w", err)
	}

	var finalResponse *pb.Response
	for {
		resp, err := stream.Recv()
		if err == io.EOF {
			break
		}
		if err != nil {
			return nil, fmt.Errorf("stream recv: %w", err)
		}
		if resp.HasExecResponse() {
			finalResponse = resp.GetExecResponse()
		}
	}

	if finalResponse == nil {
		return nil, fmt.Errorf("stream: no exec response received")
	}
	if finalResponse.GetError() != "" {
		return nil, fmt.Errorf("stream execute error: %s", finalResponse.GetError())
	}

	results := convertPBResult(finalResponse.GetResults())
	if len(results) == 0 {
		return nil, fmt.Errorf("stream: empty results")
	}
	tarResult := results[0]
	if tarResult.Status != stage.StatusAccepted || tarResult.ExitStatus != 0 {
		return nil, fmt.Errorf(
			"tar extraction failed: status=%v, exit=%d, error=%s",
			tarResult.Status, tarResult.ExitStatus, tarResult.Error,
		)
	}

	results = results[1:]
	for _, result := range results {
		maps.Copy(e.cachedMap, result.FileIDs)
	}
	return results, nil
}

func estimateCopyInSize(cmd *stage.Cmd) int {
	total := 0
	if cmd.CopyInDir != "" {
		_ = filepath.Walk(cmd.CopyInDir,
			func(path string, info os.FileInfo, err error) error {
				if err != nil || info.IsDir() {
					return nil
				}
				relPath, err := filepath.Rel(cmd.CopyInDir, path)
				if err != nil {
					return nil
				}
				if _, exists := cmd.CopyIn[relPath]; !exists {
					total += int(info.Size())
				}
				return nil
			})
	}
	for _, f := range cmd.CopyIn {
		if f.Symlink != nil {
			continue
		}
		if f.Src != nil {
			if fi, err := os.Stat(*f.Src); err == nil {
				total += int(fi.Size())
			}
		} else if f.Content != nil {
			total += len(*f.Content)
		}
	}
	return total
}

func createCopyInTar(cmd *stage.Cmd) ([]byte, error) {
	var buf bytes.Buffer
	tw := tar.NewWriter(&buf)

	if cmd.CopyInDir != "" {
		err := filepath.Walk(cmd.CopyInDir,
			func(path string, info os.FileInfo, err error) error {
				if err != nil {
					return err
				}
				relPath, err := filepath.Rel(cmd.CopyInDir, path)
				if err != nil {
					return err
				}
				if relPath == "." {
					return nil
				}
				if _, exists := cmd.CopyIn[relPath]; exists {
					if info.IsDir() {
						return filepath.SkipDir
					}
					return nil
				}
				if info.Mode()&os.ModeSymlink != 0 {
					link, err := os.Readlink(path)
					if err != nil {
						return err
					}
					hdr, err := tar.FileInfoHeader(info, link)
					if err != nil {
						return err
					}
					hdr.Name = relPath
					return tw.WriteHeader(hdr)
				}
				hdr, err := tar.FileInfoHeader(info, "")
				if err != nil {
					return err
				}
				hdr.Name = relPath
				if info.IsDir() {
					hdr.Name += "/"
				}
				if err := tw.WriteHeader(hdr); err != nil {
					return err
				}
				if info.IsDir() {
					return nil
				}
				f, err := os.Open(path)
				if err != nil {
					return err
				}
				_, err = io.Copy(tw, f)
				f.Close()
				return err
			})
		if err != nil {
			return nil, fmt.Errorf("walk copyInDir: %w", err)
		}
	}

	if err := tw.Close(); err != nil {
		return nil, fmt.Errorf("close tar: %w", err)
	}
	return buf.Bytes(), nil
}

func (e *Sandbox) Cleanup() error {
	for k, fileID := range e.cachedMap {
		req := &pb.FileID{}
		req.SetFileID(fileID)
		_, err := e.execClient.FileDelete(context.TODO(), req)
		if err != nil {
			slog.Error("sandbox cleanup", "error", err)
		}
		delete(e.cachedMap, k)
	}
	return nil
}