// Command verify sets up cross-signing keys for a Matrix bot user.
// This eliminates the "Encrypted by a device not verified by its owner" warning.
//
// Usage:
//
//	go run -tags goolm ./cmd/verify --homeserver https://... --username asistente-2 --password <pass> --token <access_token>
//	go run -tags goolm ./cmd/verify --homeserver https://... --username asistente-2 --token <access_token>   # tries dummy/admin UIA
package main

import (
	"context"
	"crypto/sha256"
	"encoding/hex"
	"fmt"
	"os"
	"path/filepath"
	"strings"

	"github.com/spf13/cobra"
	"maunium.net/go/mautrix"
	"maunium.net/go/mautrix/crypto"
	"maunium.net/go/mautrix/crypto/cryptohelper"
	"maunium.net/go/mautrix/id"
)

func main() {
	var (
		homeserver   string
		username     string
		password     string
		token        string
		storePath    string
		pickleKeyHex string
	)

	root := &cobra.Command{
		Use:   "verify",
		Short: "Set up cross-signing keys for a Matrix bot",
		Long: `Generates and uploads cross-signing keys so the bot's device is verified.
This removes the "Encrypted by a device not verified by its owner" warning.

Requires the bot's access token. Password is optional — if omitted, tries
dummy auth (MSC3967, Synapse 1.79+) then falls back to password if needed.`,
		RunE: func(cmd *cobra.Command, args []string) error {
			homeserver = strings.TrimRight(homeserver, "/")
			serverName := homeserver
			serverName = strings.TrimPrefix(serverName, "https://")
			serverName = strings.TrimPrefix(serverName, "http://")

			userID := id.UserID(fmt.Sprintf("@%s:%s", username, serverName))
			fmt.Printf("→ Setting up cross-signing for %s\n", userID)

			// Create mautrix client
			client, err := mautrix.NewClient(homeserver, userID, token)
			if err != nil {
				return fmt.Errorf("create client: %w", err)
			}

			ctx := context.Background()

			// Resolve device ID
			whoami, err := client.Whoami(ctx)
			if err != nil {
				return fmt.Errorf("whoami: %w", err)
			}
			client.DeviceID = whoami.DeviceID
			fmt.Printf("→ Device ID: %s\n", client.DeviceID)

			// Initialize crypto — use explicit pickle key if provided, else sha256(token)
			var pickleKey []byte
			if pickleKeyHex != "" {
				var err error
				pickleKey, err = hex.DecodeString(pickleKeyHex)
				if err != nil {
					return fmt.Errorf("decode pickle-key hex: %w", err)
				}
			} else {
				sum := sha256.Sum256([]byte(token))
				pickleKey = sum[:]
			}

			dbPath := filepath.Join(storePath, "crypto.db")
			if err := os.MkdirAll(filepath.Dir(dbPath), 0700); err != nil {
				return fmt.Errorf("create store dir: %w", err)
			}

			helper, err := cryptohelper.NewCryptoHelper(client, pickleKey, dbPath)
			if err != nil {
				return fmt.Errorf("create crypto helper: %w", err)
			}
			helper.DBAccountID = username

			if err := helper.Init(ctx); err != nil {
				return fmt.Errorf("init crypto: %w", err)
			}
			defer helper.Close()

			client.Crypto = helper

			// Get the OlmMachine to generate cross-signing keys
			olmMachine := helper.Machine()
			if olmMachine == nil {
				return fmt.Errorf("olm machine not available")
			}

			fmt.Println("→ Generating and uploading cross-signing keys...")

			// Try multiple UIA strategies in order of preference.
			recoveryKey, err := uploadCrossSigningKeys(ctx, olmMachine, password)
			if err != nil {
				// If keys already exist, try to just sign our device
				fmt.Printf("  Note: %v\n", err)
				fmt.Println("→ Attempting to sign own device with existing keys...")
				return signOwnDevice(ctx, olmMachine, client)
			}

			fmt.Println("✓ Cross-signing keys uploaded successfully")

			// Sign own device immediately after uploading keys
			if signErr := signOwnDevice(ctx, olmMachine, client); signErr != nil {
				fmt.Printf("  Warning: could not auto-sign device: %v\n", signErr)
			}

			fmt.Println()
			fmt.Println("─── IMPORTANT: Save the recovery key ───")
			envKey := strings.ToUpper(strings.ReplaceAll(username, "-", "_"))
			fmt.Printf("SSSS_RECOVERY_KEY_%s=%s\n", envKey, recoveryKey)
			fmt.Println()
			fmt.Println("Add this to your .env file and set recovery_key_env in the agent's config.yaml:")
			fmt.Println("  encryption:")
			fmt.Printf("    recovery_key_env: SSSS_RECOVERY_KEY_%s\n", envKey)
			return nil
		},
	}

	root.Flags().StringVar(&homeserver, "homeserver", "", "Matrix homeserver URL")
	root.Flags().StringVar(&username, "username", "", "Bot username (without @ or server)")
	root.Flags().StringVar(&password, "password", "", "Bot password (for UIA auth, optional)")
	root.Flags().StringVar(&token, "token", "", "Bot access token")
	root.Flags().StringVar(&storePath, "store", "./data/verify-crypto/", "Crypto store path")
	root.Flags().StringVar(&pickleKeyHex, "pickle-key", "", "Hex-encoded pickle key (must match agent's pickle key if sharing crypto store)")
	_ = root.MarkFlagRequired("homeserver")
	_ = root.MarkFlagRequired("username")
	_ = root.MarkFlagRequired("token")
	// password is no longer required

	if err := root.Execute(); err != nil {
		os.Exit(1)
	}
}

// uploadCrossSigningKeys tries multiple UIA strategies to upload cross-signing keys.
// Order: password (if provided) → dummy (MSC3967) → password with empty string.
func uploadCrossSigningKeys(ctx context.Context, mach *crypto.OlmMachine, password string) (string, error) {
	type strategy struct {
		name string
		fn   func(*mautrix.RespUserInteractive) interface{}
	}

	var strategies []strategy

	// If password provided, try it first
	if password != "" {
		strategies = append(strategies, strategy{
			name: "password auth",
			fn: func(uiResp *mautrix.RespUserInteractive) interface{} {
				return &mautrix.ReqUIAuthLogin{
					BaseAuthData: mautrix.BaseAuthData{
						Type:    mautrix.AuthTypePassword,
						Session: uiResp.Session,
					},
					User:     mach.Client.UserID.String(),
					Password: password,
				}
			},
		})
	}

	// Try dummy auth (MSC3967 — works on first upload with Synapse 1.79+)
	strategies = append(strategies, strategy{
		name: "dummy auth (MSC3967)",
		fn: func(uiResp *mautrix.RespUserInteractive) interface{} {
			return &mautrix.BaseAuthData{
				Type:    mautrix.AuthTypeDummy,
				Session: uiResp.Session,
			}
		},
	})

	// If no password was given, also try password auth with empty as last resort
	if password == "" {
		strategies = append(strategies, strategy{
			name: "empty password auth",
			fn: func(uiResp *mautrix.RespUserInteractive) interface{} {
				return &mautrix.ReqUIAuthLogin{
					BaseAuthData: mautrix.BaseAuthData{
						Type:    mautrix.AuthTypePassword,
						Session: uiResp.Session,
					},
					User:     mach.Client.UserID.String(),
					Password: " ", // non-empty to avoid omitempty dropping the field
				}
			},
		})
	}

	var lastErr error
	for _, s := range strategies {
		fmt.Printf("  Trying %s...\n", s.name)
		recoveryKey, _, err := mach.GenerateAndUploadCrossSigningKeys(ctx, s.fn, "")
		if err == nil {
			fmt.Printf("  ✓ Succeeded with %s\n", s.name)
			return recoveryKey, nil
		}
		fmt.Printf("  ✗ %s failed: %v\n", s.name, err)
		lastErr = err
	}

	return "", lastErr
}

func signOwnDevice(ctx context.Context, mach *crypto.OlmMachine, client *mautrix.Client) error {
	// Force-fetch own device keys from the server so the local store has
	// the correct signing key. Without this, SignOwnDevice fails with
	// "received update for device with different signing key (expected , got X)".
	devices, err := mach.FetchKeys(ctx, []id.UserID{client.UserID}, true)
	if err != nil {
		return fmt.Errorf("fetch own device keys: %w", err)
	}

	userDevices, ok := devices[client.UserID]
	if !ok {
		return fmt.Errorf("own user %s not found in fetched keys", client.UserID)
	}
	device, ok := userDevices[client.DeviceID]
	if !ok {
		return fmt.Errorf("own device %s not found in fetched keys", client.DeviceID)
	}

	if err := mach.SignOwnDevice(ctx, device); err != nil {
		return fmt.Errorf("sign own device: %w", err)
	}
	fmt.Printf("✓ Device %s signed with cross-signing key\n", client.DeviceID)
	return nil
}