package pythonrt

import (
	"context"
	"crypto/md5"
	"errors"
	"fmt"
	"os"
	"sync"

	"go.jetify.com/typeid"
	"go.uber.org/zap"

	"go.autokitteh.dev/autokitteh/sdk/sdktypes"
)

type DockerRuntimeConfig struct {
	WorkerAddressProvider  func() string
	LogRunnerCode          bool
	LogBuildCode           bool
	MaxMemoryPerWorkflowMB int64
	MaxCPUsPerWorkflow     float32
}

type dockerRunnerManager struct {
	logger                *zap.Logger
	client                *dockerClient
	runnerIDToContainerID map[string]string
	mu                    *sync.Mutex
	workerAddressProvider func() string
}

const baseImage = "baseimage:latest"

func configureDockerRunnerManager(log *zap.Logger, cfg DockerRuntimeConfig) error {
	dc, err := NewDockerClient(log, cfg)
	if err != nil {
		return err
	}

	_, err = dc.ensureNetwork()
	if err != nil {
		return err
	}

	log.Info(fmt.Sprintf("docker connected and synced succesffully, there are %d active runners on network %s", dc.ActiveRunnersCount(), networkName))

	// we don't reconnect to existing runners, we start new ones
	// so in case server started and there are some runners running
	// we stop them
	if len(dc.activeRunnerIDs) > 0 {
		log.Info("Stopping orphand runners")
		for rid := range dc.activeRunnerIDs {
			if err := dc.StopRunner(context.Background(), rid); err != nil {
				log.Warn(fmt.Sprintf("failed stopping runner %s: %s", rid, err.Error()))
				continue
			}

			log.Debug("stopped runner: " + rid)
		}
	}

	baseImagePath, err := prepareBaseImageCode("")
	if err != nil {
		return fmt.Errorf("prepare base image code: %w", err)
	}

	log.Info("building base image", zap.String("image", baseImage))
	if err := dc.BuildImage(context.Background(), baseImage, baseImagePath); err != nil {
		return fmt.Errorf("build base image: %w", err)
	}

	drm := &dockerRunnerManager{
		logger:                log,
		client:                dc,
		runnerIDToContainerID: map[string]string{},
		mu:                    new(sync.Mutex),
		workerAddressProvider: cfg.WorkerAddressProvider,
	}

	configuredRunnerType = runnerTypeDocker
	runnerManager = drm
	drm.logger.Info("configured")
	return nil
}

func createStartCommand(entrypoint, workerAddress, runnerID string) []string {
	// ["main.py", "--code-dir", "workflow", "--worker-address", "host.docker.internal:9980"]
	return []string{
		entrypoint,
		"--code-dir", "/workflow",
		"--worker-address", workerAddress,
		"--runner-id", runnerID,
	}
}

func (rm *dockerRunnerManager) prepareCustomReqImage(ctx context.Context, sessionID sdktypes.SessionID, details userCodeDetails) (string, error) {

	reqFileBytes, err := os.ReadFile(details.requirementsFilePath)
	if err != nil {
		return "", fmt.Errorf("read requirements file: %w", err)
	}
	hash := md5.Sum(reqFileBytes)
	version := fmt.Sprintf("u%x", hash)
	selectedImage := "usercode:" + version

	exists, err := rm.client.ImageExists(ctx, selectedImage)
	if err != nil {
		return "", fmt.Errorf("check image exists: %w", err)
	}
	if !exists {
		rm.logger.Info("building new user code image", zap.String("image", selectedImage), zap.String("session", sessionID.String()))
		tmpBaseImageCode, err := prepareBaseImageCode(details.requirementsFilePath)
		if err != nil {
			return "", fmt.Errorf("prepare base image code: %w", err)
		}

		if err := rm.client.BuildImage(ctx, selectedImage, tmpBaseImageCode); err != nil {
			return "", fmt.Errorf("build image: %w", err)
		}
	}

	return selectedImage, nil
}

func (rm *dockerRunnerManager) Start(ctx context.Context, sessionID sdktypes.SessionID, buildArtifacts []byte, vars map[string]string) (string, *RunnerClient, error) {
	if len(buildArtifacts) == 0 {
		return "", nil, errors.New("no build artifacts")
	}

	details, err := prepareUserCode(buildArtifacts, false)
	if err != nil {
		return "", nil, fmt.Errorf("prepare user code: %w", err)
	}

	selectedImage := baseImage
	if details.hasCustomRequirements {
		selectedImage, err = rm.prepareCustomReqImage(ctx, sessionID, details)
		if err != nil {
			return "", nil, err
		}
	}

	rid, err := typeid.WithPrefix("runner")
	if err != nil {
		return "", nil, err
	}
	runnerID := rid.String()
	cmd := createStartCommand("/runner/main.py", rm.workerAddressProvider(), runnerID)

	cid, port, err := rm.client.StartRunner(ctx, selectedImage, details.codePath, sessionID, cmd, vars)
	if err != nil {
		return "", nil, fmt.Errorf("start runner: %w", err)
	}

	runnerAddr := "127.0.0.1:" + port
	client, err := dialRunner(ctx, runnerAddr)
	if err != nil {
		defer func() {
			if err := rm.client.StopRunner(ctx, cid); err != nil {
				rm.logger.Warn("StopRunner", zap.Error(err))
			}
		}()
		exitCode, exitCodeErr := rm.client.getContainerExitCode(ctx, cid)
		// We try to get an exitCode to better understand on dial error
		// we return to the user, but if we can't we return what we have (the dialRunner error)
		// To be clear, exitCodeErr is error getting the exit code, not exitCode itself
		if exitCodeErr != nil {
			rm.logger.Warn("getContainerExitCode", zap.Error(exitCodeErr))
			return "", nil, err
		}

		if exitCode != 0 {
			rm.logger.Warn("runner exited with non-zero code", zap.Int("exit_code", exitCode))
			logs, err := rm.client.getContainerLogs(ctx, cid)
			if err != nil {
				return "", nil, errors.New("container exit code != 0, but we could not read logs")
			}
			return "", nil, errors.New(logs)
		}

		return "", nil, err
	}

	rm.mu.Lock()
	rm.runnerIDToContainerID[runnerID] = cid
	rm.mu.Unlock()
	return runnerID, client, nil
}

func (rm *dockerRunnerManager) RunnerHealth(ctx context.Context, runnerID string) error {
	rm.mu.Lock()
	cid, ok := rm.runnerIDToContainerID[runnerID]
	rm.mu.Unlock()

	if !ok {
		return errors.New("runner not found")
	}

	isRunning, err := rm.client.IsRunning(cid)
	if err != nil {
		return err
	}
	if !isRunning {
		return errors.New("runner not running")
	}

	return nil
}

func (rm *dockerRunnerManager) Stop(ctx context.Context, runnerID string) error {
	rm.mu.Lock()
	cid, ok := rm.runnerIDToContainerID[runnerID]
	rm.mu.Unlock()

	if !ok {
		return errors.New("runner not found")
	}

	return rm.client.StopRunner(ctx, cid)
}
func (*dockerRunnerManager) Health(ctx context.Context) error { return nil }
