/*
** Job Arranger for ZABBIX
** Copyright (C) 2025 Daiwa Institute of Research Ltd. All Rights Reserved.
**
** This program is free software; you can redistribute it and/or modify
** it under the terms of the GNU General Public License as published by
** the Free Software Foundation; either version 2 of the License, or
** (at your option) any later version.
**
** This program is distributed in the hope that it will be useful,
** but WITHOUT ANY WARRANTY; without even the implied warranty of
** MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
** GNU General Public License for more details.
**
** You should have received a copy of the GNU General Public License
** along with this program; if not, write to the Free Software
** Foundation, Inc., 51 Franklin Street, Fifth Floor, Boston, MA  02110-1301, USA.
**/

package main

import (
	"encoding/binary"
	"encoding/json"
	"fmt"
	"io"
	"net"
	"os"
	"path/filepath"
	"runtime/debug"
	"strconv"
	"strings"
	"time"

	clientcommon "jobarranger2/src/jobarg_server/managers/icon_exec_manager/workers/common"
	"jobarranger2/src/libs/golibs/common"
	"jobarranger2/src/libs/golibs/config_reader/server"
	"jobarranger2/src/libs/golibs/event"
	"jobarranger2/src/libs/golibs/forker"
	"jobarranger2/src/libs/golibs/utils"
)

var (
	runStatusChangeFlag     bool
	sshClientStatusFilePath string
	eventData               common.EventData
	executionData           common.IconExecutionProcessData
	agentlessData           common.IconLessData
	sessionData             common.SessionTable
	envVars                 map[string]string
)

// jaReplaceVariable replaces variables like $AAA or ${AAA} in valueSrc
func jaReplaceVariable(command string) (string, error) {

	if len(clientcommon.BeforeVariableMap) == 0 {
		return "", fmt.Errorf("before variable map is empty")
	}

	var sb strings.Builder
	runes := []rune(command)

	for i := 0; i < len(runes); {
		ch := runes[i]

		if ch != '$' {
			sb.WriteRune(ch)
			i++
			continue
		}

		// "$$" → literal "$"
		if i+1 < len(runes) && runes[i+1] == '$' {
			sb.WriteRune('$')
			i += 2
			continue
		}

		// Start variable
		i++
		braced := false
		if i < len(runes) && runes[i] == '{' {
			braced = true
			i++
		}

		start := i
		for i < len(runes) {
			if runes[i] == '$' || runes[i] == '\r' || runes[i] == '\n' {
				break
			}
			if runes[i] == '}' {
				break
			}
			i++
		}

		name := string(runes[start:i])

		// Invalid variable name
		if name == "" {
			if braced {
				return "", fmt.Errorf("invalid variable ${}")
			}
			return "", fmt.Errorf("invalid variable $")
		}

		if braced {
			if i < len(runes) && runes[i] == '}' {
				i++
			} else {
				return "", fmt.Errorf("missing closing brace for %s", name)
			}
		}

		// Replace variable
		value, ok := clientcommon.BeforeVariableMap[name]
		if !ok {
			return "", fmt.Errorf("variable not found: %s", name)
		}

		sb.WriteString(value)
	}

	return sb.String(), nil
}

func waitForSshClientReady(statusFilePath string, timeout time.Duration) error {
	sshClientReadyTimeout := time.Now().Add(timeout)

	for {
		if time.Now().After(sshClientReadyTimeout) {
			return fmt.Errorf("timeout waiting for ssh client to become ready.")
		}

		data, err := os.ReadFile(statusFilePath)
		if err == nil {
			dataStr := strings.TrimSpace(string(data))
			if dataStr != "" {
				status, convErr := strconv.Atoi(dataStr)
				if convErr != nil {
					return fmt.Errorf("invalid status value in status file: %s", dataStr)
				}

				if status != common.SSHD_CONNECT_SUCCESS {
					return fmt.Errorf("ssh client fails to become ready.")
				}

				return nil
			}
		}

		time.Sleep(100 * time.Millisecond)
	}
}

func sshConnect() (int, error) {
	socketFilename := fmt.Sprintf("%d_%s.sock", executionData.RunJobData.InnerJobnetMainID, agentlessData.SessionID)
	socketFilePath := filepath.Join(filepath.Dir(clientcommon.InFolderPath), "client_data", socketFilename)

	os.Remove(socketFilePath)

	// Duplicate session id check
	// if _, err := os.Stat(socketFilePath); err == nil {
	// 	return common.JA_JOBRESULT_FAIL, fmt.Errorf("session id is already in use. inner_jobnet_main_id: %d session_id: %s", executionData.RunJobData.InnerJobnetMainID, agentlessData.SessionID)
	// }

	// Create ssh client status file
	sshClientStatusFilePath = strings.TrimSuffix(socketFilePath, filepath.Ext(socketFilePath)) + ".status"
	statusFile, err := os.OpenFile(sshClientStatusFilePath, os.O_CREATE|os.O_WRONLY|os.O_TRUNC, 0644)
	if err != nil {
		return common.JA_JOBEXEC_FAIL, fmt.Errorf("failed to create ssh client status file. error : %s", err.Error())
	} else {
		statusFile.Close()
	}

	proc := forker.New(forker.ProcessData{
		ExecPath: common.SshClientExecPath,
		ExecParams: []string{
			agentlessData.HostIp,
			strconv.Itoa(agentlessData.HostPort),
			envVars[agentlessData.LoginUser],
			envVars[agentlessData.LoginPassword],
			envVars[agentlessData.PublicKey],
			envVars[agentlessData.PrivateKey],
			envVars[agentlessData.Passphrase],
			socketFilePath,
			strconv.Itoa(int(clientcommon.Timeout)),
		},
		Detached:   true,
		DirectExec: true,
	})
	cmd, err := proc.StartNewProcess()
	if err != nil {
		return common.JA_JOBEXEC_FAIL, fmt.Errorf("failed to start ssh client. error : %s", err.Error())
	}

	agentlessData.SshClientUdsPath = socketFilePath
	sessionData.SSHClientSocket = socketFilePath
	agentlessData.SshClientPid = cmd.Process.Pid

	// Write agentless data in data file
	if err := clientcommon.WriteStructToFD3(agentlessData); err != nil {
		return common.JA_JOBRESULT_PRE_EXEC_FAIL, fmt.Errorf("failed to write client data for sshConnect. error : %s", err.Error())
	}

	// Create run event
	err = event.CreateNextEvent(eventData, executionData.RunJobData.InnerJobnetID, executionData.RunJobnetData.JobnetID, executionData.RunJobData.InnerJobID)
	if err != nil {
		return common.JA_JOBRESULT_PRE_EXEC_FAIL, fmt.Errorf("failed to create run event. error : %s", err.Error())
	}

	runStatusChangeFlag = true

	return common.JA_JOBRESULT_SUCCEED, nil
}

func sshExec() (resp clientcommon.SshExecResponse, err error) {
	conn, err := net.Dial("unix", sessionData.SSHClientSocket)
	if err != nil {
		return resp, fmt.Errorf("failed to connect socket file. error : %s", err.Error())
	}
	defer conn.Close()

	if strings.TrimSpace(agentlessData.Command) == "" {
		return resp, nil
	}

	if agentlessData.Command, err = jaReplaceVariable(agentlessData.Command); err != nil {
		return resp, fmt.Errorf("failed to get the Job Controller variable value : %s", err.Error())
	}

	sshExecRequest := clientcommon.SshExecRequest{
		Command:      agentlessData.Command,
		RunMode:      agentlessData.RunMode,
		SessionFlag:  agentlessData.SessionFlag,
		PromptString: agentlessData.PromptString,
		LineFeedCode: agentlessData.LineFeedCode,
	}
	req, err := json.Marshal(sshExecRequest)
	if err != nil {
		return resp, fmt.Errorf("failed to marshal SSH request. error : %s", err.Error())
	}

	// Send length-prefixed JSON
	err = binary.Write(conn, binary.BigEndian, uint32(len(req)))
	if err != nil {
		return resp, fmt.Errorf("failed to write request length: %s", err)
	}

	// Send script command
	_, err = conn.Write(req)
	if err != nil {
		return resp, fmt.Errorf("failed to write request data: %s", err)
	}

	if agentlessData.SessionFlag == common.JA_SES_OPERATION_FLAG_CONTINUE {
		// Create run event
		err = event.CreateNextEvent(eventData, executionData.RunJobData.InnerJobnetID, executionData.RunJobnetData.JobnetID, executionData.RunJobData.InnerJobID)
		if err != nil {
			return resp, fmt.Errorf("failed to create icon run event. error : %s", err.Error())
		}
	}

	runStatusChangeFlag = true

	// Read response length
	var length uint32
	err = binary.Read(conn, binary.BigEndian, &length)
	if err != nil {
		return resp, fmt.Errorf("failed to read response length: %s", err)
	}

	// Read response JSON
	respData := make([]byte, length)
	_, err = io.ReadFull(conn, respData)
	if err != nil {
		return resp, fmt.Errorf("failed to read response data: %s", err)
	}

	err = json.Unmarshal(respData, &resp)
	if err != nil {
		return resp, fmt.Errorf("failed to unmarshal SSH response: %s", err)
	}

	// Update process exit code in data file
	agentlessData.JobExitCode = resp.ExitCode
	if err := clientcommon.WriteStructToFD3(agentlessData); err != nil {
		return resp, fmt.Errorf("failed to write client data for sshExec. error : %s", err.Error())
	}

	return resp, nil
}

func sshDisconnect() error {
	conn, err := net.Dial("unix", sessionData.SSHClientSocket)
	if err != nil {
		return fmt.Errorf("failed to connect socket file. error : %s", err.Error())
	}
	defer conn.Close()

	closeCmd := `{"command":"close"}`
	data := []byte(closeCmd)

	// Write length prefix
	length := uint32(len(data))
	err = binary.Write(conn, binary.BigEndian, length)
	if err != nil {
		return fmt.Errorf("failed to write message length to ssh client: %s", err)
	}

	// Write message body
	_, err = conn.Write(data)
	if err != nil {
		return fmt.Errorf("failed to send close request to ssh client: %s", err)
	}

	if agentlessData.SessionFlag == common.JA_SES_OPERATION_FLAG_CLOSE {
		// Write agentless data in data file
		if err := clientcommon.WriteStructToFD3(agentlessData); err != nil {
			return fmt.Errorf("failed to write client data for sshDisconnect. error : %s", err.Error())
		}

		// Create run event
		err = event.CreateNextEvent(eventData, executionData.RunJobData.InnerJobnetID, executionData.RunJobnetData.JobnetID, executionData.RunJobData.InnerJobID)
		if err != nil {
			return fmt.Errorf("failed to create icon run event. error : %s", err.Error())
		}

		runStatusChangeFlag = true
	}

	return nil
}

func sessionOnetime() (clientcommon.SshExecResponse, int, error) {
	var (
		exitCode int
		err      error
		resp     clientcommon.SshExecResponse
	)

	// Start ssh client to connect to sshd
	if exitCode, err = sshConnect(); err != nil {
		return resp, exitCode, err
	}

	// Wait for ssh client to becomes ready
	err = waitForSshClientReady(sshClientStatusFilePath, time.Duration(clientcommon.Timeout)*time.Second)
	if err != nil {
		return resp, common.JA_JOBRESULT_PRE_EXEC_FAIL, err
	}

	// Send execution command to ssh client
	if resp, err = sshExec(); err != nil {
		return resp, common.JA_JOBRESULT_FAIL, err
	}

	// Send disconnect request to ssh client
	if err = sshDisconnect(); err != nil {
		return resp, common.JA_JOBRESULT_FAIL, err
	}
	return resp, common.JA_JOBRESULT_SUCCEED, err
}

func sessionConnect() (clientcommon.SshExecResponse, int, error) {
	var (
		exitCode int
		err      error
		resp     clientcommon.SshExecResponse
	)

	// Start ssh client to connect to sshd
	if exitCode, err = sshConnect(); err != nil {
		return resp, exitCode, err
	}

	// Wait for ssh client to becomes ready
	err = waitForSshClientReady(sshClientStatusFilePath, time.Duration(clientcommon.Timeout)*time.Second)
	if err != nil {
		return resp, common.JA_JOBRESULT_PRE_EXEC_FAIL, err
	}

	// Send execution command to ssh client
	if resp, err = sshExec(); err != nil {
		return resp, common.JA_JOBRESULT_FAIL, err
	}
	return resp, common.JA_JOBRESULT_SUCCEED, err
}

func sessionContinue() (clientcommon.SshExecResponse, int, error) {
	var (
		err  error
		resp clientcommon.SshExecResponse
	)

	// Session check
	// if _, err := os.Stat(sessionData.SSHClientSocket); err != nil {
	// 	if os.IsNotExist(err) {
	// 		return resp, common.JA_JOBRESULT_FAIL, fmt.Errorf("session is not established. inner_jobnet_main_id: %d session_id: %s", executionData.RunJobData.InnerJobnetMainID, agentlessData.SessionID)
	// 	}
	// }

	// Send execution command to ssh client
	if resp, err = sshExec(); err != nil {
		return resp, common.JA_JOBRESULT_PRE_EXEC_FAIL, err
	}
	return resp, common.JA_JOBRESULT_SUCCEED, err
}

func sessionClose() (int, error) {
	var err error

	// Session check
	// if _, err := os.Stat(sessionData.SSHClientSocket); err != nil {
	// 	if os.IsNotExist(err) {
	// 		return common.JA_JOBRESULT_FAIL, fmt.Errorf("session is not established. inner_jobnet_main_id: %d session_id: %s", executionData.RunJobData.InnerJobnetMainID, agentlessData.SessionID)
	// 	}
	// }

	// Send disconnect request to ssh client
	if err = sshDisconnect(); err != nil {
		return common.JA_JOBRESULT_PRE_EXEC_FAIL, err
	}

	return common.JA_JOBRESULT_SUCCEED, err
}

func agentlessIconClient() {
	var (
		innerJobnetMainId, innerJobId uint64
		runCount, methodFlag          int
		flag                          bool
		resp                          clientcommon.SshExecResponse
		err                           error
		exitCode                      int
	)
	// Change json file to struct
	if executionData, err = clientcommon.GetIconExecData(); err != nil {
		fmt.Fprintf(os.Stderr, "[AgentlessIconClient] %v", err)
		os.Exit(common.JA_JOBEXEC_FAIL)
	}

	// Check necessary data
	innerJobnetMainId = executionData.RunJobData.InnerJobnetMainID
	if innerJobnetMainId <= 0 {
		err = fmt.Errorf("invalid inner_jobnet_main_id")
		fmt.Fprintf(os.Stderr, "[AgentlessIconClient] %v", err)
		os.Exit(common.JA_JOBEXEC_FAIL)
	}
	innerJobId = executionData.RunJobData.InnerJobID
	if innerJobId <= 0 {
		err = fmt.Errorf("invalid inner_job_id")
		fmt.Fprintf(os.Stderr, "[AgentlessIconClient] %v", err)
		os.Exit(common.JA_JOBEXEC_FAIL)
	}
	runCount = executionData.RunJobData.RunCount
	if runCount < 0 {
		err = fmt.Errorf("invalid run_count")
		fmt.Fprintf(os.Stderr, "[AgentlessIconClient] %v", err)
		os.Exit(common.JA_JOBEXEC_FAIL)
	}
	methodFlag = int(executionData.RunJobData.MethodFlag)
	if methodFlag < 0 {
		err = fmt.Errorf("invalid method_flag")
		fmt.Fprintf(os.Stderr, "[AgentlessIconClient] %v", err)
		os.Exit(common.JA_JOBEXEC_FAIL)
	}

	// check flag file exist or not
	flag, err = event.CheckRunCountFile(clientcommon.RunCountFolderPath, innerJobnetMainId, innerJobId, runCount, methodFlag)
	if err != nil {
		err = fmt.Errorf("failed to create flag file : %w", err)
		fmt.Fprintf(os.Stderr, "[AgentlessIconClient] %v", err)
		os.Exit(common.JA_JOBEXEC_FAIL)
	} else if !flag { // flag == false (nil && false)
		os.Exit(common.JA_JOBEXEC_IGNORE)
	}

	if err = utils.Convert(executionData.RunJobData.Data, &agentlessData); err != nil {
		fmt.Fprintf(os.Stderr, "[AgentlessIconClient] %v", err)
		os.Exit(common.JA_JOBEXEC_FAIL)
	}

	if err = utils.Convert(executionData.SessionData, &sessionData); err != nil {
		fmt.Fprintf(os.Stderr, "[AgentlessIconClient] %v", err)
		os.Exit(common.JA_JOBEXEC_FAIL)
	}

	eventData = clientcommon.IconRunDataPrep(string(clientcommon.IconClientLess), executionData)

	// flag file creation
	err = event.CreateRunCountFile(clientcommon.RunCountFolderPath, innerJobnetMainId, innerJobId, runCount)
	if err != nil {
		err = fmt.Errorf("failed to create flag file : %w", err)
		fmt.Fprintf(os.Stderr, "[AgentlessIconClient] %v", err)
		os.Exit(common.JA_JOBRESULT_PRE_EXEC_FAIL)
	}

	// set before variables
	// Set before variables
	if err := clientcommon.SetBeforeVariable(executionData); err != nil {
		os.Exit(common.JA_JOBEXEC_FAIL)
	}

	// Search and re-assign env variable
	if envVars, err = clientcommon.ResolveEnvVars(executionData); err != nil {
		os.Exit(common.JA_JOBEXEC_FAIL)
	}

	// Processing start
	switch agentlessData.SessionFlag {
	case common.JA_SES_OPERATION_FLAG_ONETIME:
		resp, exitCode, err = sessionOnetime()
	case common.JA_SES_OPERATION_FLAG_CONNECT:
		resp, exitCode, err = sessionConnect()
	case common.JA_SES_OPERATION_FLAG_CONTINUE:
		resp, exitCode, err = sessionContinue()
	case common.JA_SES_OPERATION_FLAG_CLOSE:
		exitCode, err = sessionClose()
	default:
		fmt.Fprintf(os.Stderr, "[AgentlessIconClient] Invalid Operation flag : %d", agentlessData.SessionFlag)
		os.Exit(common.JA_JOBRESULT_PRE_EXEC_FAIL)
	}

	// Write execution error
	if err != nil {
		fmt.Fprintf(os.Stderr, "[AgentlessIconClient] %v", err)
		os.Exit(exitCode)
	}

	// Write run command error
	if trimmed := strings.TrimSpace(resp.Stderr); trimmed != "" && agentlessData.RunMode == common.JA_RUN_MODE_NON_INTERACTIVE {
		fmt.Fprintf(os.Stderr, "[AgentlessIconClient] %v", trimmed)
	}

	// Write run command output
	fmt.Print(resp.Stdout)
}

func main() {

	// Catch runtime panic errors
	defer func() {
		if r := recover(); r != nil {
			// Output stacktrace
			fmt.Fprintf(os.Stderr, "[AgentlessIconClient] Runtime panic error occurs in client. error : %s", string(debug.Stack()))

			if runStatusChangeFlag {
				os.Exit(common.JA_JOBRESULT_FAIL)
			}

			os.Exit(common.JA_JOBRESULT_PRE_EXEC_FAIL)
		}
	}()

	// Add client_pid in .clientPID file
	pid := os.Getpid()
	err := common.SetClientPid(pid)
	if err != nil {
		fmt.Fprintf(os.Stderr, "[AgentlessIconClient] %v", err)
		os.Exit(common.JA_JOBEXEC_FAIL)
	}

	if err := clientcommon.ParseArgs(); err != nil {
		fmt.Fprintf(os.Stderr, "[AgentlessIconClient] %v", err)
		os.Exit(common.JA_JOBEXEC_FAIL)
	}

	/*
	 // Since the client does not accept a configuration parameter,
	 // manually assign UnixSockParentDir and TmpDir.
	*/
	server.Options.UnixSockParentDir = clientcommon.UdsDirpath
	server.Options.TmpDir = clientcommon.TmpDirPath

	agentlessIconClient()
	os.Exit(common.JA_JOBRESULT_SUCCEED)
}
