package onboarding

import (
	"crypto/rand"
	"crypto/rsa"
	"crypto/sha256"
	"crypto/x509"
	"encoding/base64"
	"encoding/json"
	"encoding/pem"
	"github.com/google/uuid"
	"testing"
	"time"
)

func TestIntegrationCredentialRotationAndReboot(t *testing.T) {
	f := newFixture(t)
	f.pending()
	rotation := RotationInput{Reason: "test rotation", ExpectedVersion: 1}
	req := uuid.NewString()
	rotated, err := f.s.Rotate(f.ctx, f.p, f.tenantPublic, f.d.PublicID, req, req, requestJSON(rotation), rotation)
	must(t, err)
	_, err = f.s.findFactoryDevice(f.ctx, f.d.FactoryID, 2)
	if err == nil {
		t.Fatal("pending key became active")
	}
	receiver, err := rsa.GenerateKey(rand.Reader, 2048)
	must(t, err)
	pub, err := x509.MarshalPKIXPublicKey(&receiver.PublicKey)
	must(t, err)
	job := rotated["enrollment_job_id"].(string)
	input := DeliveryInput{JobID: job, ReceiverKey: string(pem.EncodeToMemory(&pem.Block{Type: "PUBLIC KEY", Bytes: pub}))}
	first, err := f.s.DeliverCredential(f.ctx, f.p, f.tenantPublic, f.d.PublicID, uuid.NewString(), input)
	must(t, err)
	_, err = f.s.pool.Exec(f.ctx, `UPDATE device_enrollment_job SET delivery_expires_at=now()-interval '1 second' WHERE id=$1::uuid`, job)
	must(t, err)
	renewed, err := f.s.DeliverCredential(f.ctx, f.p, f.tenantPublic, f.d.PublicID, uuid.NewString(), input)
	must(t, err)
	if first["ciphertext_b64"] == renewed["ciphertext_b64"] {
		t.Fatal("expired envelope reused")
	}
	release := ReleaseInput{JobID: job, ExpectedVersion: 2, Evidence: "fixture readback"}
	_, err = f.s.ReleaseInventory(f.ctx, f.p, f.tenantPublic, f.d.PublicID, uuid.NewString(), release)
	must(t, err)
	_, err = f.s.findFactoryDevice(f.ctx, f.d.FactoryID, 1)
	if err == nil {
		t.Fatal("old key still active")
	}
	f.d, err = f.s.findFactoryDevice(f.ctx, f.d.FactoryID, 2)
	must(t, err)
	boot := uuid.NewString()
	challenge, err := f.s.TimeChallenge(f.ctx, TimeChallengeInput{DeviceID: f.d.FactoryID, KeyVersion: 2, BootID: boot, ClientNonce: randomToken(16)})
	must(t, err)
	raw, err := base64.RawURLEncoding.DecodeString(challenge["payload_b64"].(string))
	must(t, err)
	var packet map[string]any
	must(t, json.Unmarshal(raw, &packet))
	reg := RegisterInput{ChallengeID: packet["challenge_id"].(string), BootID: boot, Model: f.d.Model, HardwareRevision: f.d.HardwareRevision, FirmwareVersion: "1.2", SDKVersion: "1.0"}
	_, err = f.s.Register(f.ctx, f.d, uuid.NewString(), requestJSON(reg), reg)
	must(t, err)
	f.d, err = f.s.findFactoryDevice(f.ctx, f.d.FactoryID, 2)
	must(t, err)
	_, err = f.s.Confirm(f.ctx, f.d, uuid.NewString(), requestJSON(f.confirm), f.confirm)
	if err == nil {
		t.Fatal("pre-reboot confirmation accepted")
	}
}
func TestIntegrationRegistrationSurvivesRebindAndRejectsOldResult(t *testing.T) {
	f := newFixture(t)
	f.bound()
	j := f.registerSDK()
	var status string
	_, err := f.s.pool.Exec(f.ctx, `UPDATE device_binding SET unbound_at=now() WHERE device_id=$1::uuid AND unbound_at IS NULL`, f.d.ID)
	must(t, err)
	current, err := f.s.findFactoryDevice(f.ctx, f.d.FactoryID, 1)
	must(t, err)
	late := ResultInput{ActivationID: j.ID, AttemptID: j.AttemptID, Stage: "SDK_REGISTER", Outcome: "SUCCESS_OBSERVED"}
	_, err = f.s.ActivationResult(f.ctx, current, uuid.NewString(), requestJSON(late), late)
	if err == nil {
		t.Fatal("old binding result was accepted")
	}
	f.d, err = f.s.findFactoryDevice(f.ctx, f.d.FactoryID, 1)
	must(t, err)
	setupIn := SetupInput{BootID: f.d.BootID, LocalSetupID: uuid.NewString(), Transport: "LOCAL_WEB"}
	setup, err := f.s.Setup(f.ctx, f.d, uuid.NewString(), requestJSON(setupIn), setupIn)
	must(t, err)
	f.setup = setup["setup_id"].(string)
	verifier := randomToken(32)
	sum := sha256.Sum256([]byte(verifier))
	handoffIn := HandoffInput{SetupID: f.setup, LocalID: setupIn.LocalSetupID, Challenge: encode64(sum[:])}
	handoff, err := f.s.Handoff(f.ctx, f.d, uuid.NewString(), requestJSON(handoffIn), handoffIn)
	must(t, err)
	browser, err := f.s.Exchange(f.ctx, ExchangeInput{Token: handoff["handoff_token"].(string), Verifier: verifier})
	must(t, err)
	f.web, err = f.s.webSession(f.ctx, f.s.pool, digest([]byte(browser.Token)))
	must(t, err)
	browser, err = f.s.AuthLink(f.ctx, f.web, f.user)
	must(t, err)
	f.web, err = f.s.webSession(f.ctx, f.s.pool, digest([]byte(browser.Token)))
	must(t, err)
	f.bound()
	must(t, f.s.ProcessActivation(f.ctx, f.event()))
	if f.registrar.calls.Load() != 0 {
		t.Fatal("rebind repeated supplier registration")
	}
	must(t, f.s.pool.QueryRow(f.ctx, `SELECT status FROM device_activation_job WHERE device_id=$1::uuid`, f.d.ID).Scan(&status))
	if status != "REGISTERED_OBSERVED" {
		t.Fatal("rebind did not resume activation")
	}
}
func TestIntegrationCloseRequiresDurableACKAndGrace(t *testing.T) {
	f := newFixture(t)
	f.bound()
	poll, err := f.s.Poll(f.ctx, f.d, PollInput{SetupID: f.setup, BootID: f.d.BootID})
	must(t, err)
	st := poll["state"].(State)
	ack := StateAckInput{Revision: st.Revision, Generation: st.Generation, Hash: st.Hash, LocalGeneration: 1}
	_, err = f.s.StateAck(f.ctx, f.d, uuid.NewString(), requestJSON(ack), ack)
	must(t, err)
	var saved time.Time
	must(t, f.s.pool.QueryRow(f.ctx, `SELECT accepted_at FROM device_state_ack WHERE device_id=$1::uuid`, f.d.ID).Scan(&saved))
	_, err = f.s.StateAck(f.ctx, f.d, uuid.NewString(), requestJSON(ack), ack)
	must(t, err)
	var again time.Time
	must(t, f.s.pool.QueryRow(f.ctx, `SELECT accepted_at FROM device_state_ack WHERE device_id=$1::uuid`, f.d.ID).Scan(&again))
	if !again.Equal(saved) {
		t.Fatal("ACK retry extended close grace")
	}
	closeIn := CloseInput{SetupID: f.setup, Reason: "SUCCESS", Revision: st.Revision}
	_, err = f.s.Close(f.ctx, f.d, uuid.NewString(), requestJSON(closeIn), closeIn)
	if err == nil {
		t.Fatal("closed before 30 second grace")
	}
	var binding string
	must(t, f.s.pool.QueryRow(f.ctx, `SELECT public_id::text FROM device_binding WHERE device_id=$1::uuid AND unbound_at IS NULL`, f.d.ID).Scan(&binding))
	webAck := WebAckInput{BindingID: binding, Generation: st.Generation, Revision: st.Revision}
	_, err = f.s.WebAck(f.ctx, f.web, uuid.NewString(), requestJSON(webAck), webAck)
	must(t, err)
	f.s.now = func() time.Time { return time.Now().Add(4 * time.Second) }
	_, err = f.s.Close(f.ctx, f.d, uuid.NewString(), requestJSON(closeIn), closeIn)
	must(t, err)
}