package state

import (
	"encoding/json"
	"fmt"
	"github.com/Fuwn/plutia/internal/config"
	"github.com/Fuwn/plutia/internal/storage"
	"github.com/Fuwn/plutia/internal/types"
	"time"
)

type Engine struct {
	store storage.Store
	mode  string
}

func New(store storage.Store, mode string) *Engine {
	return &Engine{store: store, mode: mode}
}

func (e *Engine) Apply(op types.ParsedOperation, blockRef *types.BlockRefV1) error {
	existing, ok, err := e.store.GetState(op.DID)

	if err != nil {
		return fmt.Errorf("load state: %w", err)
	}

	var prior *types.StateV1

	if ok {
		prior = &existing
	}

	next, err := ComputeNextState(op, prior)

	if err != nil {
		return err
	}

	if pebbleStore, ok := e.store.(*storage.PebbleStore); ok {
		includeOpRef := e.mode == config.ModeMirror && blockRef != nil

		if err := pebbleStore.ApplyOperationBatch(next, blockRef, includeOpRef); err != nil {
			return fmt.Errorf("apply operation batch: %w", err)
		}

		return nil
	}

	if err := e.store.PutState(next); err != nil {
		return fmt.Errorf("put state: %w", err)
	}

	if err := e.store.SetChainHead(op.DID, op.Sequence); err != nil {
		return fmt.Errorf("set chain head: %w", err)
	}

	if err := e.store.AddDIDSequence(op.DID, op.Sequence); err != nil {
		return fmt.Errorf("append did sequence: %w", err)
	}

	if e.mode == config.ModeMirror && blockRef != nil {
		if err := e.store.PutOpSeqRef(op.Sequence, *blockRef); err != nil {
			return fmt.Errorf("put opseq ref: %w", err)
		}
	}

	if err := e.store.SetGlobalSeq(op.Sequence); err != nil {
		return fmt.Errorf("set global seq: %w", err)
	}

	return nil
}

func ComputeNextState(op types.ParsedOperation, existing *types.StateV1) (types.StateV1, error) {
	if existing != nil && op.Sequence <= existing.LatestOpSeq {
		if existing.ChainTipHash == op.CID {
			return *existing, nil
		}

		return types.StateV1{}, fmt.Errorf("non-monotonic sequence for %s: got %d <= %d", op.DID, op.Sequence, existing.LatestOpSeq)
	}

	doc, err := materializeDIDDocument(op.DID, op.Payload)

	if err != nil {
		return types.StateV1{}, err
	}

	next := types.StateV1{
		Version:      1,
		DID:          op.DID,
		DIDDocument:  doc,
		ChainTipHash: op.CID,
		LatestOpSeq:  op.Sequence,
		UpdatedAt:    time.Now().UTC(),
	}
	next.RotationKeys = extractRotationKeys(op.Payload)

	if len(next.RotationKeys) == 0 && existing != nil {
		next.RotationKeys = append([]string(nil), existing.RotationKeys...)
	}

	return next, nil
}

func materializeDIDDocument(did string, payload map[string]any) ([]byte, error) {
	for _, k := range []string{"did_document", "didDoc", "document", "didDocument"} {
		if v, ok := payload[k]; ok {
			b, err := json.Marshal(v)

			if err != nil {
				return nil, fmt.Errorf("marshal did document field %s: %w", k, err)
			}

			return types.CanonicalizeJSON(b)
		}
	}

	if v, ok := payload["alsoKnownAs"]; ok {
		doc := map[string]any{
			"id":          did,
			"alsoKnownAs": v,
		}
		b, err := json.Marshal(doc)

		if err != nil {
			return nil, fmt.Errorf("marshal derived did document: %w", err)
		}

		return types.CanonicalizeJSON(b)
	}

	doc := map[string]any{"id": did}

	if typ, _ := payload["type"].(string); typ == "plc_tombstone" || typ == "tombstone" {
		doc["deactivated"] = true
	}

	if handle, ok := payload["handle"].(string); ok && handle != "" {
		doc["alsoKnownAs"] = []string{"at://" + handle}
	}

	if serviceEndpoint, ok := payload["service"].(string); ok && serviceEndpoint != "" {
		doc["service"] = []map[string]any{
			{
				"id":              "#atproto_pds",
				"type":            "AtprotoPersonalDataServer",
				"serviceEndpoint": serviceEndpoint,
			},
		}
	}

	if signingKey, ok := payload["signingKey"].(string); ok && signingKey != "" {
		doc["verificationMethod"] = []map[string]any{
			{
				"id":                 "#atproto",
				"type":               "Multikey",
				"controller":         did,
				"publicKeyMultibase": signingKey,
			},
		}
		doc["authentication"] = []string{"#atproto"}
	}

	b, err := json.Marshal(doc)

	if err != nil {
		return nil, fmt.Errorf("marshal synthesized did document: %w", err)
	}

	return types.CanonicalizeJSON(b)
}

func extractRotationKeys(payload map[string]any) []string {
	keys := make([]string, 0, 4)
	seen := map[string]struct{}{}
	add := func(v string) {
		if v == "" {
			return
		}

		if _, ok := seen[v]; ok {
			return
		}

		seen[v] = struct{}{}
		keys = append(keys, v)
	}

	if arr, ok := payload["rotationKeys"].([]any); ok {
		for _, v := range arr {
			if s, ok := v.(string); ok {
				add(s)
			}
		}
	}

	if recovery, ok := payload["recoveryKey"].(string); ok {
		add(recovery)
	}

	if signing, ok := payload["signingKey"].(string); ok {
		add(signing)
	}

	return keys
}