package keybase

import (
	"context"
	"encoding/hex"
	"fmt"

	"enclave/internal/crypto/mpc"

	"github.com/ncruces/go-sqlite3"
	"github.com/sonr-io/crypto/core/protocol"
)

// RegisterMPCFunctions registers custom SQLite functions for MPC operations:
//   - mpc_sign(enclave_id TEXT, data BLOB) -> BLOB
//   - mpc_verify(public_key_hex TEXT, data BLOB, signature BLOB) -> INTEGER (0 or 1)
//   - mpc_refresh(enclave_id TEXT) -> TEXT
func RegisterMPCFunctions(conn *sqlite3.Conn) error {
	if err := registerSignFunction(conn); err != nil {
		return fmt.Errorf("register mpc_sign: %w", err)
	}
	if err := registerVerifyFunction(conn); err != nil {
		return fmt.Errorf("register mpc_verify: %w", err)
	}
	if err := registerRefreshFunction(conn); err != nil {
		return fmt.Errorf("register mpc_refresh: %w", err)
	}
	return nil
}

func registerSignFunction(conn *sqlite3.Conn) error {
	return conn.CreateFunction("mpc_sign", 2, sqlite3.DETERMINISTIC, func(ctx sqlite3.Context, args ...sqlite3.Value) {
		if len(args) != 2 {
			ctx.ResultError(fmt.Errorf("mpc_sign requires 2 arguments: enclave_id, data"))
			return
		}

		enclaveID := args[0].Text()
		data := args[1].RawBlob()

		if enclaveID == "" {
			ctx.ResultError(fmt.Errorf("mpc_sign: enclave_id cannot be empty"))
			return
		}
		if len(data) == 0 {
			ctx.ResultError(fmt.Errorf("mpc_sign: data cannot be empty"))
			return
		}

		enclave, err := loadEnclaveFromDB(enclaveID)
		if err != nil {
			ctx.ResultError(fmt.Errorf("mpc_sign: %w", err))
			return
		}

		signature, err := enclave.Sign(data)
		if err != nil {
			ctx.ResultError(fmt.Errorf("mpc_sign: signing failed: %w", err))
			return
		}

		ctx.ResultBlob(signature)
	})
}

func registerVerifyFunction(conn *sqlite3.Conn) error {
	return conn.CreateFunction("mpc_verify", 3, sqlite3.DETERMINISTIC|sqlite3.INNOCUOUS, func(ctx sqlite3.Context, args ...sqlite3.Value) {
		if len(args) != 3 {
			ctx.ResultError(fmt.Errorf("mpc_verify requires 3 arguments: public_key_hex, data, signature"))
			return
		}

		pubKeyHex := args[0].Text()
		data := args[1].RawBlob()
		signature := args[2].RawBlob()

		if pubKeyHex == "" {
			ctx.ResultError(fmt.Errorf("mpc_verify: public_key_hex cannot be empty"))
			return
		}
		if len(data) == 0 {
			ctx.ResultError(fmt.Errorf("mpc_verify: data cannot be empty"))
			return
		}
		if len(signature) == 0 {
			ctx.ResultError(fmt.Errorf("mpc_verify: signature cannot be empty"))
			return
		}

		pubKeyBytes, err := hex.DecodeString(pubKeyHex)
		if err != nil {
			ctx.ResultError(fmt.Errorf("mpc_verify: invalid public key hex: %w", err))
			return
		}

		valid, err := mpc.VerifyWithPubKey(pubKeyBytes, data, signature)
		if err != nil {
			ctx.ResultError(fmt.Errorf("mpc_verify: verification error: %w", err))
			return
		}

		if valid {
			ctx.ResultInt(1)
		} else {
			ctx.ResultInt(0)
		}
	})
}

func registerRefreshFunction(conn *sqlite3.Conn) error {
	return conn.CreateFunction("mpc_refresh", 1, 0, func(ctx sqlite3.Context, args ...sqlite3.Value) {
		if len(args) != 1 {
			ctx.ResultError(fmt.Errorf("mpc_refresh requires 1 argument: enclave_id"))
			return
		}

		enclaveID := args[0].Text()
		if enclaveID == "" {
			ctx.ResultError(fmt.Errorf("mpc_refresh: enclave_id cannot be empty"))
			return
		}

		enclave, err := loadEnclaveFromDB(enclaveID)
		if err != nil {
			ctx.ResultError(fmt.Errorf("mpc_refresh: %w", err))
			return
		}

		newEnclave, err := enclave.Refresh()
		if err != nil {
			ctx.ResultError(fmt.Errorf("mpc_refresh: refresh failed: %w", err))
			return
		}

		if err := updateEnclaveInDB(enclaveID, newEnclave); err != nil {
			ctx.ResultError(fmt.Errorf("mpc_refresh: update failed: %w", err))
			return
		}

		ctx.ResultText(enclaveID)
	})
}

func loadEnclaveFromDB(enclaveID string) (mpc.Enclave, error) {
	kb := Get()
	if kb == nil {
		return nil, fmt.Errorf("keybase not initialized")
	}

	dbEnc, err := kb.queries.GetEnclaveByID(context.Background(), enclaveID)
	if err != nil {
		return nil, fmt.Errorf("enclave not found: %w", err)
	}

	valShare, err := protocol.DecodeMessage(string(dbEnc.ValShare))
	if err != nil {
		return nil, fmt.Errorf("decode val_share: %w", err)
	}

	userShare, err := protocol.DecodeMessage(string(dbEnc.UserShare))
	if err != nil {
		return nil, fmt.Errorf("decode user_share: %w", err)
	}

	return &mpc.EnclaveData{
		PubHex:    dbEnc.PublicKeyHex,
		PubBytes:  dbEnc.PublicKey,
		ValShare:  valShare,
		UserShare: userShare,
		Nonce:     dbEnc.Nonce,
		Curve:     mpc.CurveName(dbEnc.Curve),
	}, nil
}

func updateEnclaveInDB(enclaveID string, enclave mpc.Enclave) error {
	kb := Get()
	if kb == nil {
		return fmt.Errorf("keybase not initialized")
	}

	data := enclave.GetData()

	valShareStr, err := protocol.EncodeMessage(data.ValShare)
	if err != nil {
		return fmt.Errorf("encode val_share: %w", err)
	}

	userShareStr, err := protocol.EncodeMessage(data.UserShare)
	if err != nil {
		return fmt.Errorf("encode user_share: %w", err)
	}

	dbEnc, err := kb.queries.GetEnclaveByID(context.Background(), enclaveID)
	if err != nil {
		return fmt.Errorf("get enclave: %w", err)
	}

	// Raw SQL required: SQLC-generated code lacks full enclave update query
	_, err = kb.db.ExecContext(context.Background(), `
		UPDATE mpc_enclaves 
		SET val_share = ?, user_share = ?, nonce = ?, rotated_at = datetime('now')
		WHERE id = ?
	`, []byte(valShareStr), []byte(userShareStr), data.Nonce, dbEnc.ID)

	return err
}
