package main

import (
	"context"
	"fmt"
	"log/slog"

	"atcr.io/pkg/atproto/did"
	"atcr.io/pkg/auth/oauth"
	"atcr.io/pkg/labeler"

	"github.com/bluesky-social/indigo/atproto/atcrypto"
	"github.com/spf13/cobra"
)

var plcCmd = &cobra.Command{
	Use:   "plc",
	Short: "PLC directory management commands",
}

var plcConfigFile string

var (
	plcAddRotationKeyFirst bool
	plcAddRotationKeyLast  bool
)

var plcAddRotationKeyCmd = &cobra.Command{
	Use:   "add-rotation-key [multibase-key]",
	Short: "Add a rotation key to this labeler's PLC identity",
	Long: `Add an additional rotation key to the labeler's did:plc document.

If a multibase-encoded private key (K-256 or P-256, starting with 'z') is supplied as
the positional argument, that key is added. If no argument is given, a fresh K-256
keypair is generated and the private half is printed to stdout. Save it offline as
your recovery key, since it will not be shown again.

By default the new key is inserted at the highest priority position (--first), which
allows it to override ops signed by lower-priority keys within PLC's 72-hour recovery
window. Pass --last to append at the lowest priority instead.

The labeler's configured rotation key is used to sign the PLC update.

  atcr-labeler plc add-rotation-key --config config.yaml             # generate + print
  atcr-labeler plc add-rotation-key --config config.yaml --last      # append, low priority
  atcr-labeler plc add-rotation-key --config config.yaml z...        # use supplied key`,
	Args: cobra.MaximumNArgs(1),
	RunE: func(cmd *cobra.Command, args []string) error {
		firstSet := cmd.Flags().Changed("first")
		lastSet := cmd.Flags().Changed("last")
		if firstSet && lastSet {
			return fmt.Errorf("--first and --last are mutually exclusive")
		}
		prepend := !plcAddRotationKeyLast

		cfg, err := labeler.LoadConfig(plcConfigFile)
		if err != nil {
			return fmt.Errorf("failed to load config: %w", err)
		}
		if cfg.Labeler.DIDMethod != "plc" {
			return fmt.Errorf("this command only works with did:plc (labeler.did_method is %q)", cfg.Labeler.DIDMethod)
		}

		ctx := context.Background()
		labelerDID, rotationKey, signingKey, err := loadLabelerPLCIdentity(ctx, cfg)
		if err != nil {
			return err
		}

		var newKey atcrypto.PrivateKeyExportable
		if len(args) == 1 {
			newKey, err = atcrypto.ParsePrivateMultibase(args[0])
			if err != nil {
				return fmt.Errorf("failed to parse key argument: %w", err)
			}
		}

		res, err := did.AddRotationKey(ctx, did.AddRotationKeyOptions{
			DID:                 labelerDID,
			PLCDirectoryURL:     cfg.PLCDirectoryURL(),
			RotationKey:         rotationKey,
			SigningKey:          signingKey,
			VerificationKeyName: "atproto_label",
			NewKey:              newKey,
			Prepend:             prepend,
		})
		if err != nil {
			return err
		}

		if res.AlreadyPresent {
			fmt.Printf("Key %s is already a rotation key for %s (priority %d of %d)\n",
				res.NewKeyDIDKey, labelerDID, res.ExistingAt, res.TotalKeys)
			return nil
		}

		if res.Generated {
			fmt.Println("=========================================================================")
			fmt.Println("GENERATED NEW ROTATION KEY. SAVE THIS NOW. IT WILL NOT BE SHOWN AGAIN.")
			fmt.Println("Store it offline (password manager, paper, hardware token).")
			fmt.Println()
			fmt.Printf("Private key (multibase): %s\n", res.NewKey.Multibase())
			fmt.Printf("Public key  (did:key):   %s\n", res.NewKeyDIDKey)
			fmt.Println("=========================================================================")
		}

		slog.Info("Added rotation key to PLC identity",
			"did", labelerDID,
			"new_key", res.NewKeyDIDKey,
			"priority", res.InsertedAt,
			"total_rotation_keys", res.TotalKeys,
			"generated", res.Generated,
		)
		fmt.Printf("Added rotation key %s to %s (priority %d of %d)\n",
			res.NewKeyDIDKey, labelerDID, res.InsertedAt, res.TotalKeys)
		return nil
	},
}

var plcListRotationKeysCmd = &cobra.Command{
	Use:   "list-rotation-keys",
	Short: "List rotation keys in this labeler's PLC document",
	Long: `Fetch the labeler's did:plc document from the PLC directory and print its
rotation keys in priority order (index 0 is highest priority and can override
ops signed by lower-priority keys within PLC's 72-hour recovery window).

The key matching the local labeler.rotation_key is marked as LOCAL.`,
	Args: cobra.NoArgs,
	RunE: func(cmd *cobra.Command, args []string) error {
		cfg, err := labeler.LoadConfig(plcConfigFile)
		if err != nil {
			return fmt.Errorf("failed to load config: %w", err)
		}
		if cfg.Labeler.DIDMethod != "plc" {
			return fmt.Errorf("this command only works with did:plc (labeler.did_method is %q)", cfg.Labeler.DIDMethod)
		}

		ctx := context.Background()
		labelerDID, _, _, err := loadLabelerPLCIdentity(ctx, cfg)
		if err != nil {
			return err
		}

		var localRotationKey atcrypto.PrivateKey
		if cfg.Labeler.RotationKey != "" {
			localRotationKey, err = atcrypto.ParsePrivateMultibase(cfg.Labeler.RotationKey)
			if err != nil {
				return fmt.Errorf("failed to parse rotation_key from config: %w", err)
			}
		}

		res, err := did.ListRotationKeys(ctx, did.ListRotationKeysOptions{
			DID:              labelerDID,
			PLCDirectoryURL:  cfg.PLCDirectoryURL(),
			LocalRotationKey: localRotationKey,
		})
		if err != nil {
			return err
		}
		printRotationKeys(res)
		return nil
	},
}

// loadLabelerPLCIdentity is the shared "load DID + rotation key + signing key" helper
// used by every PLC command. Mirrors loadHoldPLCIdentity over in cmd/hold/plc.go.
func loadLabelerPLCIdentity(ctx context.Context, cfg *labeler.Config) (string, atcrypto.PrivateKey, *atcrypto.PrivateKeyK256, error) {
	labelerDID, _, err := labeler.LoadIdentity(ctx, cfg)
	if err != nil {
		return "", nil, nil, err
	}

	if cfg.Labeler.RotationKey == "" {
		return "", nil, nil, fmt.Errorf("labeler.rotation_key must be set to sign PLC updates")
	}
	rotationKey, err := atcrypto.ParsePrivateMultibase(cfg.Labeler.RotationKey)
	if err != nil {
		return "", nil, nil, fmt.Errorf("failed to parse rotation_key from config: %w", err)
	}

	signingKey, err := oauth.GenerateOrLoadPDSKey(cfg.SigningKeyPath())
	if err != nil {
		return "", nil, nil, fmt.Errorf("failed to load signing key: %w", err)
	}
	return labelerDID, rotationKey, signingKey, nil
}

// printRotationKeys is the shared CLI output for list-rotation-keys, kept identical to
// the hold version since the formatting is service-agnostic.
func printRotationKeys(res *did.ListRotationKeysResult) {
	fmt.Printf("DID: %s\n", res.DID)
	fmt.Printf("PLC directory: %s\n", res.Directory)
	fmt.Printf("Rotation keys (%d):\n", len(res.Keys))
	for i, k := range res.Keys {
		marker := ""
		switch {
		case len(res.Keys) == 1:
			marker = "(only key)"
		case i == 0:
			marker = "(highest priority)"
		case i == len(res.Keys)-1:
			marker = "(lowest priority)"
		}
		localTag := ""
		if res.LocalDIDKey != "" && k == res.LocalDIDKey {
			localTag = " [LOCAL — labeler.rotation_key]"
		}
		fmt.Printf("  [%d] %s %s%s\n", i, k, marker, localTag)
	}

	if res.LocalDIDKey != "" && !res.LocalPresent {
		fmt.Printf("\nWARNING: local rotation_key (%s) is NOT present in the PLC document.\n", res.LocalDIDKey)
		fmt.Println("This service cannot sign PLC updates. Possible compromise or out-of-band rotation.")
	}
}

func init() {
	plcCmd.PersistentFlags().StringVarP(&plcConfigFile, "config", "c", "", "path to YAML configuration file")

	plcAddRotationKeyCmd.Flags().BoolVar(&plcAddRotationKeyFirst, "first", true, "insert at highest priority (default)")
	plcAddRotationKeyCmd.Flags().BoolVar(&plcAddRotationKeyLast, "last", false, "insert at lowest priority")

	plcCmd.AddCommand(plcAddRotationKeyCmd)
	plcCmd.AddCommand(plcListRotationKeysCmd)
}