package services

import (
	"context"
	"crypto/md5"
	"encoding/base64"
	"fmt"
	"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 computeFingerprint(pubKey string) string {
	parts := strings.Fields(pubKey)
	if len(parts) < 2 {
		return ""
	}
	raw, err := base64.StdEncoding.DecodeString(parts[1])
	if err != nil {
		return ""
	}
	sum := md5.Sum(raw)
	var pairs []string
	for _, b := range sum {
		pairs = append(pairs, fmt.Sprintf("%02x", b))
	}
	return "MD5:" + strings.Join(pairs, ":")
}

func setKeyMeta(k *models.Key) {
	k.HasPrivateKey = k.PrivateKeyEncrypted != ""
	k.HasPassphrase = k.PassphraseEncrypted != ""
}

func CreateKey(instanceID, label, publicKey, source, generatedByServerID, privateKey, passphrase string) (*models.Key, error) {
	key := &models.Key{
		InstanceID:          instanceID,
		KeyID:               uuid.NewString(),
		Label:               label,
		PublicKey:           publicKey,
		Fingerprint:         computeFingerprint(publicKey),
		Source:              source,
		GeneratedByServerID: generatedByServerID,
		CreatedAt:           time.Now(),
	}
	if privateKey != "" {
		enc, err := encryptPrivateKey(privateKey)
		if err != nil {
			return nil, fmt.Errorf("encrypt private key: %w", err)
		}
		key.PrivateKeyEncrypted = enc
	}
	if passphrase != "" {
		enc, err := encryptString(passphrase)
		if err != nil {
			return nil, fmt.Errorf("encrypt passphrase: %w", err)
		}
		key.PassphraseEncrypted = enc
	}
	ctx, cancel := context.WithTimeout(context.Background(), 5*time.Second)
	defer cancel()

	_, err := db.Col("keys").InsertOne(ctx, key)
	if err != nil {
		return nil, err
	}
	setKeyMeta(key)
	return key, nil
}

func GetKey(instanceID, keyID string) (*models.Key, error) {
	ctx, cancel := context.WithTimeout(context.Background(), 5*time.Second)
	defer cancel()

	var key models.Key
	err := db.Col("keys").FindOne(ctx, bson.M{"key_id": keyID, "instance_id": instanceID}).Decode(&key)
	if err != nil {
		return nil, err
	}
	setKeyMeta(&key)
	return &key, nil
}

func GetPrivateKey(instanceID, keyID string) (string, error) {
	ctx, cancel := context.WithTimeout(context.Background(), 5*time.Second)
	defer cancel()

	var key models.Key
	if err := db.Col("keys").FindOne(ctx, bson.M{"key_id": keyID, "instance_id": instanceID}).Decode(&key); err != nil {
		return "", err
	}
	if key.PrivateKeyEncrypted == "" {
		return "", fmt.Errorf("no private key stored for this key")
	}
	return decryptPrivateKey(key.PrivateKeyEncrypted)
}

func GetPassphrase(keyID string) (string, error) {
	ctx, cancel := context.WithTimeout(context.Background(), 5*time.Second)
	defer cancel()

	var key models.Key
	if err := db.Col("keys").FindOne(ctx, bson.M{"key_id": keyID}).Decode(&key); err != nil {
		return "", err
	}
	if key.PassphraseEncrypted == "" {
		return "", nil
	}
	return decryptString(key.PassphraseEncrypted)
}

type KeyWithCount struct {
	models.Key    `bson:",inline"`
	AssignedCount int `bson:"-" json:"assigned_count"`
}

func ListKeys(instanceID string) ([]KeyWithCount, error) {
	ctx, cancel := context.WithTimeout(context.Background(), 5*time.Second)
	defer cancel()

	cursor, err := db.Col("keys").Find(ctx, bson.M{"instance_id": instanceID})
	if err != nil {
		return nil, err
	}
	defer cursor.Close(ctx)

	var keys []models.Key
	if err := cursor.All(ctx, &keys); err != nil {
		return nil, err
	}

	result := make([]KeyWithCount, 0, len(keys))
	for _, k := range keys {
		setKeyMeta(&k)
		count, _ := db.Col("assignments").CountDocuments(ctx, bson.M{
			"instance_id": instanceID,
			"key_id":      k.KeyID,
			"revoked_at":  nil,
		})
		result = append(result, KeyWithCount{Key: k, AssignedCount: int(count)})
	}
	return result, nil
}

func DeleteKey(instanceID, keyID string) error {
	ctx, cancel := context.WithTimeout(context.Background(), 5*time.Second)
	defer cancel()

	var key models.Key
	if err := db.Col("keys").FindOne(ctx, bson.M{"key_id": keyID, "instance_id": instanceID}).Decode(&key); err != nil {
		return err
	}

	if _, err := db.Col("keys").DeleteOne(ctx, bson.M{"key_id": keyID, "instance_id": instanceID}); err != nil {
		return err
	}
	if _, err := db.Col("assignments").DeleteMany(ctx, bson.M{"key_id": keyID, "instance_id": instanceID}); err != nil {
		return err
	}

	if key.Source == "generated" && key.GeneratedByServerID != "" {
		DispatchDeleteKey(key.GeneratedByServerID, key.Label)
	}
	return nil
}

func AssignKey(instanceID, keyID, serverID string) (*models.Assignment, error) {
	ctx, cancel := context.WithTimeout(context.Background(), 5*time.Second)
	defer cancel()

	if _, err := GetKey(instanceID, keyID); err != nil {
		return nil, fmt.Errorf("key not found")
	}
	if _, err := GetServer(instanceID, serverID); err != nil {
		return nil, fmt.Errorf("server not found")
	}

	var existing models.Assignment
	err := db.Col("assignments").FindOne(ctx, bson.M{
		"instance_id": instanceID,
		"key_id":      keyID,
		"server_id":   serverID,
		"revoked_at":  nil,
	}).Decode(&existing)
	if err == nil {
		return &existing, nil
	}

	a := &models.Assignment{
		InstanceID: instanceID,
		KeyID:      keyID,
		ServerID:   serverID,
		AssignedAt: time.Now(),
	}
	_, err = db.Col("assignments").InsertOne(ctx, a)
	if err != nil {
		return nil, err
	}
	return a, nil
}

func RevokeAssignment(instanceID, keyID, serverID string) error {
	ctx, cancel := context.WithTimeout(context.Background(), 5*time.Second)
	defer cancel()

	now := time.Now()
	_, err := db.Col("assignments").UpdateOne(ctx,
		bson.M{"instance_id": instanceID, "key_id": keyID, "server_id": serverID, "revoked_at": nil},
		bson.M{"$set": bson.M{"revoked_at": now}},
	)
	return err
}

func GetAssignmentsForKey(instanceID, keyID string) ([]models.Assignment, error) {
	ctx, cancel := context.WithTimeout(context.Background(), 5*time.Second)
	defer cancel()

	cursor, err := db.Col("assignments").Find(ctx, bson.M{"instance_id": instanceID, "key_id": keyID, "revoked_at": nil})
	if err != nil {
		return nil, err
	}
	defer cursor.Close(ctx)

	var assignments []models.Assignment
	if err := cursor.All(ctx, &assignments); err != nil {
		return nil, err
	}
	return assignments, nil
}

type AssignmentWithServer struct {
	models.Assignment `bson:",inline"`
	Server            *models.Server `json:"server,omitempty"`
}

func GetAssignmentsWithServers(instanceID, keyID string) ([]AssignmentWithServer, error) {
	ctx, cancel := context.WithTimeout(context.Background(), 5*time.Second)
	defer cancel()

	cursor, err := db.Col("assignments").Find(ctx, bson.M{"instance_id": instanceID, "key_id": keyID})
	if err != nil {
		return nil, err
	}
	defer cursor.Close(ctx)

	var assignments []models.Assignment
	if err := cursor.All(ctx, &assignments); err != nil {
		return nil, err
	}

	result := make([]AssignmentWithServer, 0, len(assignments))
	for _, a := range assignments {
		item := AssignmentWithServer{Assignment: a}
		var srv models.Server
		if err := db.Col("servers").FindOne(ctx, bson.M{"server_id": a.ServerID, "instance_id": instanceID}).Decode(&srv); err == nil {
			item.Server = &srv
		}
		result = append(result, item)
	}
	return result, nil
}

type AssignmentWithKey struct {
	models.Assignment `bson:",inline"`
	Key               *models.Key `json:"key,omitempty"`
}

func GetAssignmentsWithKeysForServer(instanceID, serverID string) ([]AssignmentWithKey, error) {
	ctx, cancel := context.WithTimeout(context.Background(), 5*time.Second)
	defer cancel()

	cursor, err := db.Col("assignments").Find(ctx, bson.M{"instance_id": instanceID, "server_id": serverID})
	if err != nil {
		return nil, err
	}
	defer cursor.Close(ctx)

	var assignments []models.Assignment
	if err := cursor.All(ctx, &assignments); err != nil {
		return nil, err
	}

	result := make([]AssignmentWithKey, 0, len(assignments))
	for _, a := range assignments {
		var key models.Key
		if err := db.Col("keys").FindOne(ctx, bson.M{"key_id": a.KeyID, "instance_id": instanceID}).Decode(&key); err != nil {
			continue
		}
		setKeyMeta(&key)
		result = append(result, AssignmentWithKey{Assignment: a, Key: &key})
	}
	return result, nil
}
