package services

import (
	"context"
	"crypto/hmac"
	"crypto/sha256"
	"encoding/base64"
	"fmt"
	"strconv"
	"strings"
	"time"

	"gitea.hostxtra.co.uk/mrhid6/vantage/server/internal/db"
	"gitea.hostxtra.co.uk/mrhid6/vantage/server/internal/models"
	"github.com/google/uuid"
	"go.mongodb.org/mongo-driver/v2/bson"
)

func sessionHMACKey() ([]byte, error) {

	k, err := encryptionKey()
	if err != nil {
		return nil, err
	}
	mac := hmac.New(sha256.New, k)
	mac.Write([]byte("vantage-console-session-v1"))
	return mac.Sum(nil), nil
}

func b64(b []byte) string { return base64.RawURLEncoding.EncodeToString(b) }

func SignSessionToken(sessionID string, ttl time.Duration) (string, error) {
	key, err := sessionHMACKey()
	if err != nil {
		return "", err
	}
	exp := time.Now().Add(ttl).Unix()
	payload := fmt.Sprintf("%s.%d", b64([]byte(sessionID)), exp)
	mac := hmac.New(sha256.New, key)
	mac.Write([]byte(payload))
	return payload + "." + b64(mac.Sum(nil)), nil
}

func VerifySessionToken(token string) (string, error) {
	parts := strings.Split(token, ".")
	if len(parts) != 3 {
		return "", fmt.Errorf("malformed token")
	}
	payload := parts[0] + "." + parts[1]
	key, err := sessionHMACKey()
	if err != nil {
		return "", err
	}
	mac := hmac.New(sha256.New, key)
	mac.Write([]byte(payload))
	want := mac.Sum(nil)
	got, err := base64.RawURLEncoding.DecodeString(parts[2])
	if err != nil || !hmac.Equal(want, got) {
		return "", fmt.Errorf("invalid signature")
	}
	exp, err := strconv.ParseInt(parts[1], 10, 64)
	if err != nil {
		return "", fmt.Errorf("invalid expiry")
	}
	if time.Now().Unix() > exp {
		return "", fmt.Errorf("token expired")
	}
	sid, err := base64.RawURLEncoding.DecodeString(parts[0])
	if err != nil {
		return "", fmt.Errorf("invalid session id")
	}
	return string(sid), nil
}

type GuacParams struct {
	Protocol string
	Params   map[string]string
}

func portOr(v, def int) int {
	if v == 0 {
		return def
	}
	return v
}

// TargetPort is the port on the managed server that the agent will dial on its
// own loopback address.
func TargetPort(srv *models.Server, protocol string) (int, error) {
	switch protocol {
	case "ssh":
		return portOr(srv.SSHPort, 22), nil
	case "rdp":
		return portOr(srv.RDPPort, 3389), nil
	case "vnc":
		return 5900, nil
	default:
		return 0, fmt.Errorf("unsupported protocol %q", protocol)
	}
}

// BuildGuacParams points guacd at the relay listener, never at the managed
// server: on a cloud deployment the server's address is not routable from here.
func BuildGuacParams(protocol, sshUser, privateKey, passphrase, rdpUser, rdpPass, relayHost string, relayPort int) (*GuacParams, error) {
	port := strconv.Itoa(relayPort)
	switch protocol {
	case "ssh":
		p := map[string]string{
			"hostname": relayHost,
			"port":     port,
		}
		if sshUser == "" {
			sshUser = "root"
		}
		p["username"] = sshUser
		if privateKey != "" {
			p["private-key"] = privateKey
		}
		if passphrase != "" {
			p["passphrase"] = passphrase
		}
		return &GuacParams{Protocol: "ssh", Params: p}, nil
	case "rdp":
		return &GuacParams{Protocol: "rdp", Params: map[string]string{
			"hostname":    relayHost,
			"port":        port,
			"username":    rdpUser,
			"password":    rdpPass,
			"security":    "any",
			"ignore-cert": "true",
		}}, nil
	case "vnc":
		return &GuacParams{Protocol: "vnc", Params: map[string]string{
			"hostname": relayHost,
			"port":     port,
			"password": rdpPass,
		}}, nil
	default:
		return nil, fmt.Errorf("unsupported protocol %q", protocol)
	}
}

func CreateConsoleSession(instanceID, serverID, protocol, keyID, user, clientIP string) (*models.ConsoleSession, error) {
	ctx, cancel := context.WithTimeout(context.Background(), 5*time.Second)
	defer cancel()

	s := &models.ConsoleSession{
		InstanceID: instanceID,
		SessionID:  uuid.NewString(),
		ServerID:   serverID,
		Protocol:   protocol,
		KeyID:      keyID,
		User:       user,
		ClientIP:   clientIP,
		StartedAt:  time.Now(),
	}
	if _, err := db.Col("console_sessions").InsertOne(ctx, s); err != nil {
		return nil, err
	}
	return s, nil
}

func GetConsoleSession(instanceID, sessionID string) (*models.ConsoleSession, error) {
	ctx, cancel := context.WithTimeout(context.Background(), 5*time.Second)
	defer cancel()
	var s models.ConsoleSession
	if err := db.Col("console_sessions").FindOne(ctx, bson.M{"session_id": sessionID, "instance_id": instanceID}).Decode(&s); err != nil {
		return nil, err
	}
	return &s, nil
}

func StashConsoleRDPCreds(instanceID, sessionID, username, password string) error {
	ctx, cancel := context.WithTimeout(context.Background(), 5*time.Second)
	defer cancel()
	u, err := encryptString(username)
	if err != nil {
		return err
	}
	p, err := encryptString(password)
	if err != nil {
		return err
	}
	_, err = db.Col("console_sessions").UpdateOne(ctx,
		bson.M{"session_id": sessionID, "instance_id": instanceID},
		bson.M{"$set": bson.M{"rdp_user_enc": u, "rdp_pass_enc": p}},
	)
	return err
}

func ConsumeConsoleRDPCreds(instanceID, sessionID string) (username, password string, err error) {
	s, err := GetConsoleSession(instanceID, sessionID)
	if err != nil {
		return "", "", err
	}
	if s.RDPUserEnc == "" && s.RDPPassEnc == "" {
		return "", "", nil
	}
	if s.RDPUserEnc != "" {
		if username, err = decryptString(s.RDPUserEnc); err != nil {
			return "", "", err
		}
	}
	if s.RDPPassEnc != "" {
		if password, err = decryptString(s.RDPPassEnc); err != nil {
			return "", "", err
		}
	}
	ctx, cancel := context.WithTimeout(context.Background(), 5*time.Second)
	defer cancel()
	_, _ = db.Col("console_sessions").UpdateOne(ctx,
		bson.M{"session_id": sessionID, "instance_id": instanceID},
		bson.M{"$unset": bson.M{"rdp_user_enc": "", "rdp_pass_enc": ""}},
	)
	return username, password, nil
}

func SetConsoleSSHUser(instanceID, sessionID, username string) error {
	ctx, cancel := context.WithTimeout(context.Background(), 5*time.Second)
	defer cancel()
	_, err := db.Col("console_sessions").UpdateOne(ctx,
		bson.M{"session_id": sessionID, "instance_id": instanceID},
		bson.M{"$set": bson.M{"ssh_username": username}})
	return err
}

func ConsumeSessionToken(instanceID, sessionID string) error {
	ctx, cancel := context.WithTimeout(context.Background(), 5*time.Second)
	defer cancel()
	now := time.Now()
	res, err := db.Col("console_sessions").UpdateOne(ctx,
		bson.M{"session_id": sessionID, "instance_id": instanceID, "token_consumed_at": nil},
		bson.M{"$set": bson.M{"token_consumed_at": now}},
	)
	if err != nil {
		return err
	}
	if res.MatchedCount == 0 {
		return fmt.Errorf("session token already used")
	}
	return nil
}

func EndConsoleSession(instanceID, sessionID string) error {
	ctx, cancel := context.WithTimeout(context.Background(), 5*time.Second)
	defer cancel()
	now := time.Now()
	_, err := db.Col("console_sessions").UpdateOne(ctx,
		bson.M{"session_id": sessionID, "instance_id": instanceID, "ended_at": nil},
		bson.M{"$set": bson.M{"ended_at": now}},
	)
	return err
}
