package keybase

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

	"enclave/internal/crypto/mpc"

	"github.com/ncruces/go-sqlite3"
)

func RegisterMPCFunctions(conn *sqlite3.Conn) error {
	if err := registerGenerateFunction(conn); err != nil {
		return fmt.Errorf("register mpc_generate: %w", err)
	}
	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 registerGenerateFunction(conn *sqlite3.Conn) error {
	return conn.CreateFunction("mpc_generate", 0, 0, func(ctx sqlite3.Context, args ...sqlite3.Value) {
		enc, err := mpc.NewSimpleEnclave()
		if err != nil {
			ctx.ResultError(fmt.Errorf("mpc_generate: %w", err))
			return
		}

		idBytes := make([]byte, 8)
		rand.Read(idBytes)
		enclaveID := fmt.Sprintf("enc_%x", idBytes)

		kb := Get()
		if kb == nil {
			ctx.ResultError(fmt.Errorf("mpc_generate: keybase not initialized"))
			return
		}

		am, err := NewActionManager()
		if err != nil {
			ctx.ResultError(fmt.Errorf("mpc_generate: action manager: %w", err))
			return
		}

		_, err = am.CreateEnclave(context.Background(), NewEnclaveInput{
			EnclaveID:    enclaveID,
			PublicKeyHex: enc.PubKeyHex(),
			PublicKey:    enc.PubKeyBytes(),
			ValShare:     enc.GetShare1(),
			UserShare:    enc.GetShare2(),
			Nonce:        enc.GetNonce(),
			Curve:        string(enc.GetCurve()),
		})
		if err != nil {
			ctx.ResultError(fmt.Errorf("mpc_generate: create enclave: %w", err))
			return
		}

		ctx.ResultText(enclaveID)
	})
}

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 := loadSimpleEnclaveFromDB(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 := loadSimpleEnclaveFromDB(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 := updateSimpleEnclaveInDB(enclaveID, newEnclave); err != nil {
			ctx.ResultError(fmt.Errorf("mpc_refresh: update failed: %w", err))
			return
		}

		ctx.ResultText(enclaveID)
	})
}

func loadSimpleEnclaveFromDB(enclaveID string) (*mpc.SimpleEnclave, 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)
	}

	return mpc.ImportSimpleEnclave(
		dbEnc.PublicKey,
		dbEnc.ValShare,
		dbEnc.UserShare,
		dbEnc.Nonce,
		mpc.CurveName(dbEnc.Curve),
	)
}

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

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

	_, err = kb.db.ExecContext(context.Background(), `
		UPDATE mpc_enclaves 
		SET val_share = ?, user_share = ?, nonce = ?, rotated_at = datetime('now')
		WHERE id = ?
	`, enclave.GetShare1(), enclave.GetShare2(), enclave.GetNonce(), dbEnc.ID)

	return err
}
