/*
 * Copyright 2021 PingCAP, Inc.
 *
 * Licensed under the Apache License, Version 2.0 (the "License");
 * you may not use this file except in compliance with the License.
 * You may obtain a copy of the License at
 *
 *     http://www.apache.org/licenses/LICENSE-2.0
 *
 * Unless required by applicable law or agreed to in writing, software
 * distributed under the License is distributed on an "AS IS" BASIS,
 * See the License for the specific language governing permissions and
 * limitations under the License.
 */

package main

import (
	"crypto/tls"
	"fmt"
	"io/fs"
	"io/ioutil"
	"net/http"
	"net/http/httputil"
	"net/url"
	"os"
	"os/exec"
	"path"
	"path/filepath"
	"strings"
	"syscall"
	"time"

	"github.com/pingcap/log"
	"go.uber.org/zap"

	"github.com/pingcap/mimir/pkg/cert"
	"github.com/pingcap/mimir/pkg/constant"
)

type flinkCluster struct {
	masterPID      int
	taskmanagerPID int
	flinkDir       string
	logsDir        string
}

func newFlinkCluster(flinkDir string) (*flinkCluster, error) {
	logsDir, err := ioutil.TempDir("/tmp", "mim-logs-*")
	if err != nil {
		return nil, err
	}
	return &flinkCluster{
		flinkDir: flinkDir,
		logsDir:  logsDir,
	}, nil
}

func (c *flinkCluster) bootstrap(endpoint string) (err error) {
	defer func() {
		if err != nil {
			c.terminate()
		}
	}()

	if endpoint != "" {
		url, err := url.Parse(endpoint + "/flink-api")
		if err != nil {
			return err
		}
		var tlsCfg *tls.Config
		if url.Scheme == "https" {
			tlsCfg, err = cert.LoadClientCert(path.Join(constant.HomeDir(), "tls"))
			if err != nil {
				return err
			}
		}
		go proxy(url, tlsCfg)
		return nil
	}

	if err := c.bootMaster(); err != nil {
		return err
	}
	if err := c.bootTaskmanager(); err != nil {
		return err
	}
	return nil
}

func (c *flinkCluster) bootSQLClient() error {
	jars := []string{}
	jars = append(jars, searchJars(filepath.Join(c.flinkDir, "lib"))...)
	jars = append(jars, searchJars(filepath.Join(c.flinkDir, "opt"))...)

	args := []string{
		"-XX:+UseG1GC",
		fmt.Sprintf("-Dlog.file=%s/flink-root-sql-client.log", c.logsDir),
		fmt.Sprintf("-Dlog4j.configuration=file:%s/conf/log4j-cli.properties", c.flinkDir),
		fmt.Sprintf("-Dlog4j.configurationFile=file:%s/conf/log4j-cli.properties", c.flinkDir),
		fmt.Sprintf("-Dlogback.configurationFile=file:%s/conf/logback.xml", c.flinkDir),
		"-classpath", strings.Join(jars, ":"),
		"org.apache.flink.table.client.SqlClient",
		"embedded",
	}

	cmd := exec.Command("java", args...)
	cmd.Env = []string{
		fmt.Sprintf("FLINK_CONF_DIR=%s", filepath.Join(c.flinkDir, "conf")),
	}
	cmd.Stdin = os.Stdin
	cmd.Stdout = os.Stdout
	cmd.Stderr = os.Stderr
	return cmd.Run()
}

func (c *flinkCluster) terminate() {
	if c.masterPID != 0 {
		if err := syscall.Kill(c.masterPID, syscall.SIGKILL); err != nil {
			log.Error("kill flink master failed", zap.Int("PID", c.masterPID))
		}
	}
	if c.taskmanagerPID != 0 {
		if err := syscall.Kill(c.taskmanagerPID, syscall.SIGKILL); err != nil {
			log.Error("kill flink taskmanager failed", zap.Int("PID", c.taskmanagerPID))
		}
	}
}

func (c *flinkCluster) bootMaster() error {
	jars := []string{}
	jars = append(jars, searchJars(filepath.Join(c.flinkDir, "lib"))...)
	jars = append(jars, searchJars(filepath.Join(c.flinkDir, "opt"))...)

	args := []string{
		"-XX:+UseG1GC",
		fmt.Sprintf("-Dlog.file=%s/flink-root-standalonesession.log", c.logsDir),
		fmt.Sprintf("-Dlog4j.configuration=file:%s/conf/log4j-cli.properties", c.flinkDir),
		fmt.Sprintf("-Dlog4j.configurationFile=file:%s/conf/log4j-cli.properties", c.flinkDir),
		fmt.Sprintf("-Dlogback.configurationFile=file:%s/conf/logback.xml", c.flinkDir),
		"-classpath", strings.Join(jars, ":"),
		"org.apache.flink.runtime.entrypoint.StandaloneSessionClusterEntrypoint",
		"--configDir", filepath.Join(c.flinkDir, "conf"),
		"--executionMode", "cluster",
	}

	cmd := exec.Command("java", args...)
	if err := cmd.Start(); err != nil {
		return err
	}
	c.masterPID = cmd.Process.Pid
	return nil
}

func (c *flinkCluster) bootTaskmanager() error {
	jars := []string{}
	jars = append(jars, searchJars(filepath.Join(c.flinkDir, "lib"))...)
	jars = append(jars, searchJars(filepath.Join(c.flinkDir, "opt"))...)

	args := []string{
		"-XX:+UseG1GC",
		fmt.Sprintf("-Dlog.file=%s/flink-root-taskexecutor.log", c.logsDir),
		fmt.Sprintf("-Dlog4j.configuration=file:%s/conf/log4j-cli.properties", c.flinkDir),
		fmt.Sprintf("-Dlog4j.configurationFile=file:%s/conf/log4j-cli.properties", c.flinkDir),
		fmt.Sprintf("-Dlogback.configurationFile=file:%s/conf/logback.xml", c.flinkDir),
		"-classpath", strings.Join(jars, ":"),
		"org.apache.flink.runtime.taskexecutor.TaskManagerRunner",
		"--configDir", filepath.Join(c.flinkDir, "conf"),
		"-D", "taskmanager.memory.framework.off-heap.size=134217728b",
		"-D", "taskmanager.memory.network.max=275146345b",
		"-D", "taskmanager.memory.network.min=275146345b",
		"-D", "taskmanager.memory.framework.heap.size=134217728b",
		"-D", "taskmanager.memory.managed.size=1100585384b",
		"-D", "taskmanager.cpu.cores=1.0",
		"-D", "taskmanager.memory.task.heap.size=1107296234b",
		"-D", "taskmanager.memory.task.off-heap.size=0b",
	}

	cmd := exec.Command("java", args...)
	if err := cmd.Start(); err != nil {
		return err
	}
	c.taskmanagerPID = cmd.Process.Pid
	return nil
}

func searchJars(dir string) []string {
	jars := []string{}
	if err := filepath.WalkDir(dir, func(path string, info fs.DirEntry, err error) error {
		if err != nil {
			log.Warn("retrievin flink lib", zap.Error(err))
			return nil
		}
		if info.IsDir() {
			return nil
		}
		if strings.HasSuffix(info.Name(), ".jar") {
			jars = append(jars, path)
		}
		return nil
	}); err != nil {
		log.Error("search jars", zap.Error(err))
	}
	return jars
}

func proxy(url *url.URL, tlsCfg *tls.Config) {
	proxy := httputil.NewSingleHostReverseProxy(url)
	if tlsCfg != nil {
		proxy.Transport = makeTransport(tlsCfg)
	}

	s := &http.Server{
		Addr:           ":8081",
		Handler:        proxy,
		ReadTimeout:    10 * time.Second,
		WriteTimeout:   10 * time.Second,
		MaxHeaderBytes: 1 << 20,
	}

	log.Fatal("bootstrap proxy", zap.Error(s.ListenAndServe()))
}

func makeTransport(tlsCfg *tls.Config) *http.Transport {
	// Start with the DefaultTransport for sane defaults.
	transport := http.DefaultTransport.(*http.Transport).Clone()
	// Conservatively disable HTTP keep-alives as this program will only
	// ever need a single HTTP request.
	transport.DisableKeepAlives = true
	// Timeout early if the server doesn't even return the headers.
	transport.ResponseHeaderTimeout = time.Minute
	// We should clone a tlsCfg because we use it across goroutine
	if tlsCfg != nil {
		transport.TLSClientConfig = tlsCfg.Clone()
	}

	transport.Proxy = http.ProxyFromEnvironment
	return transport
}