package main

import (
	// "bytes"
	// "bufio"
	"bytes"
	"embed"
	"encoding/json"
	"fmt"
	"html/template"
	"io"
	"log"
	"net/http"
	"net/url"
	"os"
	"os/exec"
	"os/user"
	"path/filepath"
	"regexp"
	"strings"

	"github.com/creack/pty"
	"github.com/gorilla/websocket"
)

type PageData struct {
	CurrentDir  string
	Hostname 	string
	CurrentUsername string
	CommandLog  []CommandOutput
}

type CommandOutput struct {
	Command string
	Output  string
	Error   string
}

type CommandNode struct {
	Operator 	string
	Left 		*CommandNode
	Right 		*CommandNode
	Command 	string
}

//go:embed templates/*
var templateFiles embed.FS
//go:embed static/*
var staticFiles embed.FS

var commandLog []CommandOutput
var variables = make(map[string]string)

var tmpl = template.Must(template.ParseFS(templateFiles, "templates/index.html"))

func expandCommandSubstitution(command string, currentDir string) (string, error) {
	// Do not expand single quoted strings
	reSingleQuotes := regexp.MustCompile(`'([^']*)'`)
	singleQuotedStrings := reSingleQuotes.FindAllString(command, -1)

	placeholder := "\x00PLACEHOLDER\x00"
	command = reSingleQuotes.ReplaceAllLiteralString(command, placeholder)

	// Expand variables
	reVar := regexp.MustCompile(`\$(\w+)`)
	command = reVar.ReplaceAllStringFunc(command, func(match string) string {
		varName := match[1:]
		if val, exists := variables[varName]; exists {
			return val
		}
		if envVal, exists := os.LookupEnv(varName); exists {
			return envVal
		}

		return match
	})

	// Restore single-quoted strings and avoid expansion inside them
	for _, sq := range singleQuotedStrings {
		command = strings.Replace(command, placeholder, sq, 1)
	}

	// expand $(...) command substitution
	re := regexp.MustCompile(`\$\(([^()]+)\)`)
	for {
		match := re.FindStringSubmatch(command)
		if match == nil {
			break
		}

		subCommand := strings.TrimSpace(match[1])
		subOutput, err := executeCommandTree(parseCommandTree(subCommand), currentDir)
		if err != nil {
			return "", err
		}

		command = strings.Replace(command, match[0], strings.TrimSpace(subOutput), 1)
	}
	return command, nil
}

func setVariable(command string) bool {
	re := regexp.MustCompile(`^(\w+)=(.*)$`)
	match := re.FindStringSubmatch(command)
	if match == nil {
		return false
	}

	varName := match[1]
	value := strings.Trim(match[2], `"`)

	if strings.HasPrefix(command, "export ") {
		os.Setenv(varName, value)
	} else {
		variables[varName] = value
	}

	return true
}

func parseCommandTree(command string) *CommandNode {
	if strings.Contains(command, ";") {
		parts := strings.SplitN(command, ";", 2)
		return &CommandNode{
			Operator: ";",
			Left: parseCommandTree(strings.TrimSpace(parts[0])),
			Right: parseCommandTree(strings.TrimSpace(parts[1])),
		}
	} else if strings.Contains(command, ">") {
		parts := strings.SplitN(command, ">", 2)
		return &CommandNode{
			Operator: ">",
			Left: parseCommandTree(strings.TrimSpace(parts[0])),
			Right: &CommandNode{Command: strings.TrimSpace(parts[1])},
		}
	} else if strings.Contains(command, "<") {
		parts := strings.SplitN(command, "<", 2)
		return &CommandNode{
			Operator: "<",
			Left: parseCommandTree(strings.TrimSpace(parts[0])),
			Right: &CommandNode{Command: strings.TrimSpace(parts[1])},
		}
	} else if strings.Contains(command, "|") {
		parts := strings.SplitN(command, "|", 2)
		return &CommandNode{
			Operator: "|",
			Left: parseCommandTree(strings.TrimSpace(parts[0])),
			Right: parseCommandTree(strings.TrimSpace(parts[1])),
		}
	} else if strings.Contains(command, "&&") {
		parts := strings.SplitN(command, "&&", 2)
		return &CommandNode{
			Operator: "&&",
			Left: parseCommandTree(strings.TrimSpace(parts[0])),
			Right: parseCommandTree(strings.TrimSpace(parts[1])),
		}
	} else if strings.Contains(command, "||") {
		parts := strings.SplitN(command, "||", 2)
		return &CommandNode{
			Operator: "||",
			Left: parseCommandTree(strings.TrimSpace(parts[0])),
			Right: parseCommandTree(strings.TrimSpace(parts[1])),
		}
	}
	return &CommandNode{Command: command}
}

func executeCommandTree(node *CommandNode, currentDir string) (string, error) {
	if node == nil {
		return "", nil
	}

	switch node.Operator {
	case "&&":
		leftOutput, err := executeCommandTree(node.Left, currentDir)
		if err != nil {
			return "", err
		}
		if leftOutput != "" {
			return executeCommandTree(node.Right, currentDir)
		}
		return executeCommandTree(node.Right, currentDir)

	// TODO: or operator does not work correctly, fix it
	case "||":
		leftOutput, err := executeCommandTree(node.Left, currentDir)
		if err == nil && leftOutput != "" {
			return leftOutput, nil
		}
		return executeCommandTree(node.Right, currentDir)

	case "|":
		leftOutput, err := executeCommandTree(node.Left, currentDir)
		if err != nil {
			return "", err
		}
		cmdArgs := parseCommandWithQuotes(node.Right.Command)
		if len(cmdArgs) == 0 {
			return "", fmt.Errorf("Invalid Command")
		}
		cmd := exec.Command(cmdArgs[0], cmdArgs[1:]...)
		cmd.Dir = currentDir
		cmd.Stdin = strings.NewReader(leftOutput)
		var outputBuffer bytes.Buffer
		cmd.Stdout = &outputBuffer
		cmd.Stderr = &outputBuffer
		if err := cmd.Run(); err != nil {
			return "", err
		}
		return outputBuffer.String(), nil

	case ">":
		leftOutput, err := executeCommandTree(node.Left, currentDir)
		if err != nil {
			return "", err
		}
		file, err := os.Create(strings.TrimSpace(node.Right.Command))
		if err != nil {
			return "", err
		}
		defer file.Close()
		_, err = file.WriteString(leftOutput)
		return "", err

	case "<":
		file, err := os.Open(strings.TrimSpace(node.Right.Command))
		if err != nil {
			return "", err
		}
		defer file.Close()
		cmdArgs := parseCommandWithQuotes(node.Left.Command)
		if len(cmdArgs) == 0 {
			return "", fmt.Errorf("Invalid command")
		}
		cmd := exec.Command(cmdArgs[0], cmdArgs[1:]...)
		cmd.Dir = currentDir
		cmd.Stdin = file
		var outputBuffer bytes.Buffer
		cmd.Stdout = &outputBuffer
		cmd.Stderr = &outputBuffer
		if err := cmd.Run(); err != nil {
			return "", err
		}
		return outputBuffer.String(), nil

	case ";":
		leftOutput, _ := executeCommandTree(node.Left, currentDir)
		rightOutput, _ := executeCommandTree(node.Right, currentDir)
		return leftOutput + rightOutput, nil
	}
	return executeSimpleCommand(node.Command, currentDir).Output, nil
}

func executeCommand(command string, args []string, dir string) (string, error) {
	cmd := exec.Command(command, args...)
	cmd.Dir = dir
	output, err := cmd.CombinedOutput()
	return string(output), err
}

func executeSimpleCommand(command string, currentDir string) CommandOutput {
	var cmdOutput CommandOutput
	// args := strings.Fields(command)
	args := parseCommandWithQuotes(command)
	if len(args) == 0 {
		return cmdOutput
	}

	cmd := exec.Command(args[0], args[1:]...)
	cmd.Dir = currentDir
	output, err := cmd.CombinedOutput()
	cmdOutput.Command = command
	cmdOutput.Output = string(output)
	if err != nil {
		cmdOutput.Error = err.Error()
	}
	return cmdOutput
}

func changeDirectory(command string, args []string, currentDir *string) CommandOutput {
	var cmdOutput CommandOutput

	if len(args) == 0 {
		homeDir, _ := os.UserHomeDir()
		err := os.Chdir(homeDir)
		if err != nil {
			cmdOutput = CommandOutput {
				Command: command,
				Error: "Failed to change to home directory: " + err.Error(),
			}
		} else {
			cmdOutput = CommandOutput {
				Command: command,
				Output: "Changed to home directory: " +  homeDir,
			}
			*currentDir = homeDir
		}
	} else {
		newDir := args[0]
		err := os.Chdir(newDir)
		if err != nil {
			cmdOutput = CommandOutput {
				Command: command + " " + newDir,
				Error: "Failed to change to directory: " + err.Error(),
			}
		} else {
			cmdOutput = CommandOutput {
				Command: command + " " + newDir,
				Output: "Changed to directory: " + newDir,
			}
			*currentDir = newDir
		}
	}
	return cmdOutput
}

func fileUploadHandler(w http.ResponseWriter, r *http.Request) {
	if r.Method != http.MethodPost {
		http.Error(w, "Invalid request method", http.StatusMethodNotAllowed)
		return
	}

	file, header, err := r.FormFile("file")
	if err != nil {
		http.Error(w, "Failed to upload file", http.StatusInternalServerError)
		return
	}
	defer file.Close()

	uploadDir := r.FormValue("uploadDir")
	if uploadDir == "" {
		uploadDir = "./"
	}

	os.MkdirAll(uploadDir, os.ModePerm)
	destPath := filepath.Join(uploadDir, header.Filename)
	out, err := os.Create(destPath)
	if err != nil {
		http.Error(w, "Failed to save file", http.StatusInternalServerError)
		return
	}
	defer out.Close()

	_, err = io.Copy(out, file)
	if err != nil {
		http.Error(w, "Failed to write file", http.StatusInternalServerError)
		return
	}

	w.Write([]byte("File uploaded successfully"))
}

func fileDownloadHandler(w http.ResponseWriter, r *http.Request) {
	filePath := r.URL.Query().Get("filePath")
	if filePath == ""{
		http.Error(w, "Filepath is required", http.StatusBadRequest)
		return
	}
	absPath, err := filepath.Abs(filePath)
	if err != nil || !fileExists(absPath) {
		http.Error(w, "File not found", http.StatusNotFound)
		return
	}

	w.Header().Set("Content-Disposition", "attachement; filename=\""+filepath.Base(absPath)+"\"")
	w.Header().Set("Content-Type", "application/octet-stream")
	http.ServeFile(w, r, absPath)
}

func fileExists(filePath string) bool {
	_, err := os.Stat(filePath)
	return !os.IsNotExist(err)
}

func processCommand(command string, currentDir string) CommandOutput {
	var errorMsg string

	// Handle variable assignments before execution
	if setVariable(command) {
		return CommandOutput{Command: command, Output: "", Error: ""}
	}

	// This handles command substitution which is $(...) in bash
	expandedCommand, err := expandCommandSubstitution(command, currentDir)
	if err != nil {
		errorMsg = err.Error()
		return CommandOutput{Command: command, Output: "", Error: errorMsg}
	}

	tree := parseCommandTree(expandedCommand)
	output, err := executeCommandTree(tree, currentDir)
	if err != nil {
		errorMsg = err.Error()
	}
	return CommandOutput{Command: command, Output: output, Error: errorMsg}
}

func parseCommandWithQuotes(command string) []string {
	var args []string
	var current strings.Builder
	inSingleQuote, inDoubleQuote := false, false

	for i := 0; i < len(command); i++ {
		c := command[i]

		switch c {
		case '\'':
			if !inDoubleQuote {
				inSingleQuote = !inSingleQuote
				continue
			}
		case '"':
			if !inSingleQuote {
				inDoubleQuote = !inDoubleQuote
				continue
			}
		case ' ':
			if !inSingleQuote && !inDoubleQuote {
				if current.Len() > 0 {
					args = append(args, current.String())
					current.Reset()
				}
				continue
			}
		case '\\':
			// Handle escaped characters
			if i+1 < len(command) {
				next := command[i+1]
				if next == '"' || next == '\'' || next == '\\' {
					current.WriteByte(next)
					i++
					continue
				}
			}
		}

		current.WriteByte(c)
	}

	// Add the last argument
	if current.Len() > 0 {
		args = append(args, current.String())
	}

	return args
}

func handler(w http.ResponseWriter, r *http.Request) {
	currentDir, _ := os.Getwd()
	displayDir := filepath.Base(currentDir)
	currentUser, _ := user.Current()
	currentUsername := currentUser.Username
	hostname, _ := os.Hostname()

	if r.Method == http.MethodPost {
		input := r.FormValue("command")
		parts := strings.Fields(input)
		if len(parts) > 0 {
			command := parts[0]
			args := parts[1:]
			if  command == "cd" {
				cmdOutput := changeDirectory(command, args, &currentDir)
				displayDir = filepath.Base(currentDir)
				commandLog = append(commandLog, cmdOutput)
			} else if command == "download" {
				if len(args) < 1{
					http.Error(w, "Usage: download <filePath>", http.StatusBadRequest)
					return
				}
				http.Redirect(w, r, "/download?filePath="+url.QueryEscape(args[0]), http.StatusFound)
				return
			} else {
				cmdOutput := processCommand(input, currentDir)
				commandLog = append(commandLog, cmdOutput)
			}
		}
	}

	data := PageData{
		CurrentDir: displayDir,
		Hostname: hostname,
		CurrentUsername: currentUsername,
		CommandLog: commandLog,
	}
	tmpl.Execute(w, data)
}

func main() {
	http.HandleFunc("/", handler)
	http.HandleFunc("/upload", fileUploadHandler)
	http.HandleFunc("/download", fileDownloadHandler)
	http.HandleFunc("/terminal", terminalHandler)
	http.Handle("/static/", http.FileServer(http.FS(staticFiles)))
	fmt.Println("Starting server on :8080")
	http.ListenAndServe(":8080", nil)
}

var upgrader = websocket.Upgrader {
	ReadBufferSize: 1024,
	WriteBufferSize: 1024,
	CheckOrigin: func(r *http.Request) bool { return true },
}

func terminalHandler (w http.ResponseWriter, r *http.Request) {
	ws, err := upgrader.Upgrade(w, r, nil)
	if err != nil {
		log.Printf("Websocket upgrade error: %v", err)
		return
	}
	defer ws.Close()

	cmd := exec.Command("bash")
	cmd.Env = append(os.Environ(), "TERM=xterm-256color")
	ptmx, err := pty.Start(cmd)
	if err != nil {
		log.Printf("Error starting PTY: %v", err)
		return
	}
	defer func() { _ = ptmx.Close() }()
	// ptmx.Write([]byte("stty raw -echo\n"))
	// ptmx.Write([]byte("stty -echo\n"))
	// ptmx.Write([]byte("export SHELL=bash 1>&2 2>/dev/null; export TERM=xterm-256color 1>&2 2>/dev/null"))

	go func() {
		buf := make([]byte, 1024)
		for {
			n, err := ptmx.Read(buf)
			if err != nil {
				break
			}
			if err := ws.WriteMessage(websocket.BinaryMessage, buf[:n]); err != nil {
				break
			}
		}
	}()

	for {
			_, message, err := ws.ReadMessage()
		if err != nil {
			break
		}

		// Check if the message is binary and starts with control prefix 0xFF.
		if len(message) > 0 && message[0] == 0xFF {
			var resizeMsg struct {
				Type string `json:"type"`
				Cols uint16 `json:"cols"`
				Rows uint16 `json:"rows"`
			}
			if err := json.Unmarshal(message[1:], &resizeMsg); err == nil && resizeMsg.Type == "resize" {
				err := pty.Setsize(ptmx, &pty.Winsize{
					Cols: resizeMsg.Cols,
					Rows: resizeMsg.Rows,
					X:    0,
					Y:    0,
				})
				if err != nil {
					log.Printf("Error resizing PTY: %v", err)
				} else {
					log.Printf("Resized PTY to cols: %d, rows: %d", resizeMsg.Cols, resizeMsg.Rows)
				}
				continue // Do not write this message to the PTY.
			}
		}

		// Otherwise, treat it as normal input.
		if _, err := ptmx.Write(message); err != nil {
			break
		}
	}
}