package ssh

import (
	"context"
	"fmt"
	"os"
	"runtime"
	"slices"
	"time"

	"github.com/docker/docker/pkg/ioutils"
	"github.com/mattn/go-colorable"
	"github.com/pkg/errors"
	"github.com/samber/lo"
	"github.com/spf13/cobra"
	fly "github.com/superfly/fly-go"
	"github.com/superfly/fly-go/flaps"
	"github.com/superfly/flyctl/agent"
	"github.com/superfly/flyctl/internal/appconfig"
	"github.com/superfly/flyctl/internal/command"
	"github.com/superfly/flyctl/internal/flag"
	"github.com/superfly/flyctl/internal/flapsutil"
	"github.com/superfly/flyctl/internal/flyutil"
	"github.com/superfly/flyctl/internal/prompt"
	"github.com/superfly/flyctl/internal/sentry"
	"github.com/superfly/flyctl/iostreams"
	"github.com/superfly/flyctl/ip"
	"github.com/superfly/flyctl/ssh"
	"github.com/superfly/flyctl/terminal"
)

func stdArgsSSH(cmd *cobra.Command) {
	flag.Add(cmd,
		flag.Org(),
		flag.App(),
		flag.AppConfig(),
		flag.String{
			Name:        "command",
			Shorthand:   "C",
			Default:     "",
			Description: "command to run on SSH session",
		},
		flag.String{
			Name:        "machine",
			Default:     "",
			Description: "Run the console in the existing machine with the specified ID",
		},
		flag.Bool{
			Name:        "select",
			Shorthand:   "s",
			Default:     false,
			Description: "select available instances",
		},
		flag.Region(),
		flag.Bool{
			Name:        "quiet",
			Shorthand:   "q",
			Description: "Don't print progress indicators for WireGuard",
		},
		flag.String{
			Name:        "address",
			Shorthand:   "A",
			Description: "Address of VM to connect to",
		},
		flag.String{
			Name:        "container",
			Description: "Container to connect to",
		},
		flag.Bool{
			Name:        "no-container",
			Description: "Connect to the machine itself rather than to one of its containers",
		},
		flag.Bool{
			Name:        "pty",
			Description: "Allocate a pseudo-terminal (default: on when no command is provided)",
		},
		flag.String{
			Name:        "user",
			Shorthand:   "u",
			Description: "Unix username to connect as",
			Default:     DefaultSshUsername,
		},
		flag.ProcessGroup(""),
	)
}

func quiet(ctx context.Context) bool {
	return flag.GetBool(ctx, "quiet")
}

func lookupAddressAndContainer(ctx context.Context, cli *agent.Client, dialer agent.Dialer, app *fly.AppCompact, console bool) (addr string, container string, err error) {
	selectedMachine, err := selectMachine(ctx, app)
	if err != nil {
		return "", "", err
	}

	container, err = selectContainer(ctx, selectedMachine)
	if err != nil {
		return "", "", err
	}

	if addr = flag.GetString(ctx, "address"); addr != "" {
		return addr, container, nil
	}

	if addr == "" {
		if console && len(flag.Args(ctx)) != 0 {
			addr = flag.Args(ctx)[0]
		} else {
			addr = selectedMachine.PrivateIP
		}
	}

	// wait for the addr to be resolved in dns unless it's an ip address
	if !ip.IsV6(addr) {
		if err := cli.WaitForDNS(ctx, dialer, app.Organization.Slug, addr, ""); err != nil {
			captureError(ctx, err, app)

			return "", "", errors.Wrapf(err, "host unavailable at %s", addr)
		}
	}

	return
}

func newConsole() *cobra.Command {
	const (
		long  = `Connect to a running instance of the current app.`
		short = long
		usage = "console"
	)

	cmd := command.New(usage, short, long, runConsole, command.RequireSession, command.RequireAppName)

	cmd.Args = cobra.MaximumNArgs(1)

	stdArgsSSH(cmd)

	return cmd
}

// SessionTarget selects where a session runs on the remote machine.
type SessionTarget = ssh.SessionTarget

func captureError(ctx context.Context, err error, app *fly.AppCompact) {
	// ignore cancelled errors
	if errors.Is(err, context.Canceled) {
		return
	}

	sentry.CaptureException(err,
		sentry.WithTraceID(ctx),
		sentry.WithTag("feature", "ssh-console"),
		sentry.WithContexts(map[string]sentry.Context{
			"app": map[string]any{
				"name": app.Name,
			},
			"organization": map[string]any{
				"name": app.Organization.Slug,
			},
		}),
	)
}

func runConsole(ctx context.Context) error {
	client := flyutil.ClientFromContext(ctx)
	appName := appconfig.NameFromContext(ctx)

	if !quiet(ctx) {
		terminal.Debugf("Retrieving app info for %s\n", appName)
	}

	app, err := client.GetAppCompact(ctx, appName)
	if err != nil {
		return fmt.Errorf("get app: %w", err)
	}

	network, err := client.GetAppNetwork(ctx, app.Name)
	if err != nil {
		return fmt.Errorf("get app network: %w", err)
	}

	agentclient, dialer, err := agent.BringUpAgent(ctx, client, app, *network, quiet(ctx))
	if err != nil {
		return err
	}

	addr, container, err := lookupAddressAndContainer(ctx, agentclient, dialer, app, true)
	if err != nil {
		return err
	}

	// TODO: eventually remove the exception for sh and bash.
	cmd := flag.GetString(ctx, "command")
	allocPTY := cmd == "" || flag.GetBool(ctx, "pty")
	if !allocPTY && (cmd == "sh" || cmd == "/bin/sh" || cmd == "bash" || cmd == "/bin/bash") {
		terminal.Warn(
			"Allocating a pseudo-terminal since the command provided is a shell. " +
				"This behavior will change in the future; please use --pty explicitly if this is what you want.",
		)
		allocPTY = true
	}

	params := &ConnectParams{
		Ctx:            ctx,
		Org:            app.Organization,
		Dialer:         dialer,
		Username:       flag.GetString(ctx, "user"),
		DisableSpinner: quiet(ctx),
		Container:      container,
		Machine:        flag.GetBool(ctx, "no-container"),
		AppNames:       []string{app.Name},
	}
	sshc, err := Connect(params, addr)
	if err != nil {
		captureError(ctx, err, app)

		return err
	}

	target := SessionTarget{Container: params.Container, Machine: params.Machine}

	if err := Console(ctx, sshc, cmd, allocPTY, target); err != nil {
		captureError(ctx, err, app)

		return err
	}

	return nil
}

func Console(ctx context.Context, sshClient *ssh.Client, cmd string, allocPTY bool, target SessionTarget) error {
	currentStdin, currentStdout, currentStderr, err := setupConsole()
	defer func() error {
		if err := cleanupConsole(currentStdin, currentStdout, currentStderr); err != nil {
			return err
		}

		return nil
	}()

	sessIO := &ssh.SessionIO{
		Stdin: os.Stdin,
		// "colorable" package should be used after the console setup performed above.
		// Otherwise, virtual terminal emulation provided by the package will break UTF-8 encoding.
		// If flyctl targets Windows 10+ only then we can avoid using this package at all
		// because Windows 10+ already provides virtual terminal support.
		Stdout:   ioutils.NewWriteCloserWrapper(colorable.NewColorableStdout(), func() error { return nil }),
		Stderr:   ioutils.NewWriteCloserWrapper(colorable.NewColorableStderr(), func() error { return nil }),
		AllocPTY: allocPTY,
		TermEnv:  determineTermEnv(),
	}

	if err := sshClient.Shell(ctx, sessIO, cmd, target); err != nil {
		return errors.Wrap(err, "ssh shell")
	}

	return err
}

func selectMachine(ctx context.Context, app *fly.AppCompact) (machine *fly.Machine, err error) {
	out := iostreams.FromContext(ctx).Out
	flapsClient := flapsutil.ClientFromContext(ctx)

	machines, err := flapsClient.ListActive(ctx, app.Name)
	if err != nil {
		return nil, err
	}

	machines = lo.Filter(machines, func(m *fly.Machine, _ int) bool {
		return m.State == "started"
	})

	if len(machines) < 1 {
		return nil, fmt.Errorf("app %s has no started VMs.\nIt may be unhealthy or not have been deployed yet.\nTry the following command to verify:\n\nfly status", app.Name)
	}

	if region := flag.GetRegion(ctx); region != "" {
		machines = lo.Filter(machines, func(m *fly.Machine, _ int) bool {
			return m.Region == region
		})
		if len(machines) < 1 {
			return nil, fmt.Errorf("app %s has no VMs in region %s", app.Name, region)
		}
	}

	if group := flag.GetProcessGroup(ctx); group != "" {
		machines = lo.Filter(machines, func(m *fly.Machine, _ int) bool {
			return m.ProcessGroup() == group
		})
		if len(machines) < 1 {
			return nil, fmt.Errorf("app %s has no VMs in process group %s", app.Name, group)
		}
	}

	var namesWithRegion []string
	machineID := flag.GetString(ctx, "machine")
	var selectedMachine *fly.Machine
	multipleGroups := len(lo.UniqBy(machines, func(m *fly.Machine) string { return m.ProcessGroup() })) > 1

	for _, machine := range machines {
		if machine.ID == machineID {
			selectedMachine = machine
		}

		nameWithRegion := fmt.Sprintf("%s: %s %s %s", machine.Region, machine.ID, machine.PrivateIP, machine.Name)

		role := ""
		for _, check := range machine.Checks {
			if check.Name == "role" {
				if check.Status == fly.Passing {
					role = check.Output
				} else {
					role = "error"
				}
			}
		}

		if role != "" {
			nameWithRegion += fmt.Sprintf(" (%s)", role)
		}

		if multipleGroups {
			nameWithRegion += fmt.Sprintf(" (%s)", machine.ProcessGroup())
		}
		namesWithRegion = append(namesWithRegion, nameWithRegion)
	}

	if machineID != "" && selectedMachine == nil {
		return nil, fmt.Errorf("--machine=%q not found/started", machineID)
	}

	if flag.GetBool(ctx, "select") {
		if flag.IsSpecified(ctx, "machine") {
			return nil, errors.New("--machine can't be used with -s/--select")
		}

		selected := 0

		if len(namesWithRegion) > 1 {
			if err = prompt.Select(ctx, &selected, "Select VM:", "", namesWithRegion...); err != nil {
				return nil, fmt.Errorf("selecting VM: %w", err)
			}
		} else if len(machines) == 1 {
			fmt.Fprintf(out, "Only one machine available, selecting %s in region %s\n", machines[0].ID, machines[0].Region)
		}

		selectedMachine = machines[selected]
	}

	if selectedMachine != nil {
		if selectedMachine.State != "started" {
			fmt.Fprintf(out, "Starting machine %s..", selectedMachine.ID)
			_, err := flapsClient.Start(ctx, app.Name, selectedMachine.ID, "")
			if err != nil {
				return nil, err
			}

			err = flapsClient.Wait(ctx, app.Name, selectedMachine.ID, flaps.WithWaitStates("started"), flaps.WithWaitTimeout(60*time.Second))

			if err != nil {
				return nil, err
			}
		}
	}

	// No VM was selected or passed as an argument, so just pick the first one for now
	// Later, we might want to use 'nearest.of' but also resolve the machine IP to be able to start it

	if selectedMachine == nil {
		selectedMachine = machines[0]

		if len(machines) > 1 || flag.GetBool(ctx, "select") {
			fmt.Fprintf(out, "No machine specified, using %s in region %s\n", selectedMachine.ID, selectedMachine.Region)
		}
	}

	return selectedMachine, nil
}

// selectContainer selects a container from the machine's config.
func selectContainer(ctx context.Context, machine *fly.Machine) (container string, err error) {
	containers := machine.Config.Containers
	container = flag.GetString(ctx, "container")

	// --no-container asks for the machine itself, so there is nothing to select.
	if flag.GetBool(ctx, "no-container") {
		if container != "" {
			return "", errors.New("--container and --no-container are mutually exclusive")
		}

		return "", nil
	}

	if len(containers) == 0 {
		if container == "" {
			return "", nil
		} else {
			return "", fmt.Errorf("no containers found for machine %s", machine.ID)
		}
	} else if len(containers) == 1 {
		if container == "" || container == containers[0].Name {
			return containers[0].Name, nil
		} else {
			return "", fmt.Errorf("container named %s is not present in machine %s, try running with --select to see a list", container, machine.ID)
		}
	} else {
		var availableContainers []string
		for _, c := range containers {
			availableContainers = append(availableContainers, c.Name)
		}

		if container == "" {
			selected := 0
			if len(availableContainers) > 1 {
				if flag.GetBool(ctx, "select") {
					if err = prompt.Select(ctx, &selected, "Select container:", "", availableContainers...); err != nil {
						return "", fmt.Errorf("selecting container: %w", err)
					}
				} else {
					fmt.Printf("No container specified, using %s\n", availableContainers[0])
				}
			}

			return availableContainers[selected], nil
		} else {
			if slices.Contains(availableContainers, container) {
				return container, nil
			} else {
				return "", fmt.Errorf("container named %s is not present in machine %s, try running with --select to see a list", container, machine.ID)
			}
		}
	}

}

const defaultTermEnv = "xterm"

func determineTermEnv() string {
	switch runtime.GOOS {
	case "aix":
		return determineTermEnvFromLocalEnv()
	case "darwin":
		return determineTermEnvFromLocalEnv()
	case "dragonfly":
		return determineTermEnvFromLocalEnv()
	case "freebsd":
		return determineTermEnvFromLocalEnv()
	case "illumos":
		return determineTermEnvFromLocalEnv()
	case "linux":
		return determineTermEnvFromLocalEnv()
	case "netbsd":
		return determineTermEnvFromLocalEnv()
	case "openbsd":
		return determineTermEnvFromLocalEnv()
	case "solaris":
		return determineTermEnvFromLocalEnv()
	default:
		return defaultTermEnv
	}
}

func determineTermEnvFromLocalEnv() string {
	if term := os.Getenv("TERM"); term != "" {
		return term
	} else {
		return defaultTermEnv
	}
}
