package nostr

import (
	"encoding/json"
	"fmt"
)

// Envelope represents a Nostr protocol message.
type Envelope interface {
	Label() string
	MarshalJSON() ([]byte, error)
}

// EventEnvelope wraps an event for the EVENT message.
// Used both client→relay (publish) and relay→client (subscription).
type EventEnvelope struct {
	SubscriptionID string // Only for relay→client messages
	Event          *Event
}

func (e EventEnvelope) Label() string { return "EVENT" }

func (e EventEnvelope) MarshalJSON() ([]byte, error) {
	if e.SubscriptionID != "" {
		return json.Marshal([]interface{}{"EVENT", e.SubscriptionID, e.Event})
	}
	return json.Marshal([]interface{}{"EVENT", e.Event})
}

// ReqEnvelope represents a REQ message (client→relay).
type ReqEnvelope struct {
	SubscriptionID string
	Filters        []Filter
}

func (e ReqEnvelope) Label() string { return "REQ" }

func (e ReqEnvelope) MarshalJSON() ([]byte, error) {
	arr := make([]interface{}, 2+len(e.Filters))
	arr[0] = "REQ"
	arr[1] = e.SubscriptionID
	for i, f := range e.Filters {
		arr[2+i] = f
	}
	return json.Marshal(arr)
}

// CloseEnvelope represents a CLOSE message (client→relay).
type CloseEnvelope struct {
	SubscriptionID string
}

func (e CloseEnvelope) Label() string { return "CLOSE" }

func (e CloseEnvelope) MarshalJSON() ([]byte, error) {
	return json.Marshal([]interface{}{"CLOSE", e.SubscriptionID})
}

// OKEnvelope represents an OK message (relay→client).
type OKEnvelope struct {
	EventID string
	OK      bool
	Message string
}

func (e OKEnvelope) Label() string { return "OK" }

func (e OKEnvelope) MarshalJSON() ([]byte, error) {
	return json.Marshal([]interface{}{"OK", e.EventID, e.OK, e.Message})
}

// EOSEEnvelope represents an EOSE (End of Stored Events) message (relay→client).
type EOSEEnvelope struct {
	SubscriptionID string
}

func (e EOSEEnvelope) Label() string { return "EOSE" }

func (e EOSEEnvelope) MarshalJSON() ([]byte, error) {
	return json.Marshal([]interface{}{"EOSE", e.SubscriptionID})
}

// ClosedEnvelope represents a CLOSED message (relay→client).
type ClosedEnvelope struct {
	SubscriptionID string
	Message        string
}

func (e ClosedEnvelope) Label() string { return "CLOSED" }

func (e ClosedEnvelope) MarshalJSON() ([]byte, error) {
	return json.Marshal([]interface{}{"CLOSED", e.SubscriptionID, e.Message})
}

// NoticeEnvelope represents a NOTICE message (relay→client).
type NoticeEnvelope struct {
	Message string
}

func (e NoticeEnvelope) Label() string { return "NOTICE" }

func (e NoticeEnvelope) MarshalJSON() ([]byte, error) {
	return json.Marshal([]interface{}{"NOTICE", e.Message})
}

// ParseEnvelope parses a raw JSON message into the appropriate envelope type.
func ParseEnvelope(data []byte) (Envelope, error) {
	var arr []json.RawMessage
	if err := json.Unmarshal(data, &arr); err != nil {
		return nil, fmt.Errorf("invalid envelope: %w", err)
	}

	if len(arr) < 2 {
		return nil, fmt.Errorf("envelope too short")
	}

	var label string
	if err := json.Unmarshal(arr[0], &label); err != nil {
		return nil, fmt.Errorf("invalid envelope label: %w", err)
	}

	switch label {
	case "EVENT":
		return parseEventEnvelope(arr)
	case "REQ":
		return parseReqEnvelope(arr)
	case "CLOSE":
		return parseCloseEnvelope(arr)
	case "OK":
		return parseOKEnvelope(arr)
	case "EOSE":
		return parseEOSEEnvelope(arr)
	case "CLOSED":
		return parseClosedEnvelope(arr)
	case "NOTICE":
		return parseNoticeEnvelope(arr)
	default:
		return nil, fmt.Errorf("unknown envelope type: %s", label)
	}
}

func parseEventEnvelope(arr []json.RawMessage) (*EventEnvelope, error) {
	env := &EventEnvelope{}

	if len(arr) == 2 {
		// Client→relay: ["EVENT", event]
		var event Event
		if err := json.Unmarshal(arr[1], &event); err != nil {
			return nil, fmt.Errorf("invalid event: %w", err)
		}
		env.Event = &event
	} else if len(arr) == 3 {
		// Relay→client: ["EVENT", subscription_id, event]
		if err := json.Unmarshal(arr[1], &env.SubscriptionID); err != nil {
			return nil, fmt.Errorf("invalid subscription ID: %w", err)
		}
		var event Event
		if err := json.Unmarshal(arr[2], &event); err != nil {
			return nil, fmt.Errorf("invalid event: %w", err)
		}
		env.Event = &event
	} else {
		return nil, fmt.Errorf("invalid EVENT envelope length: %d", len(arr))
	}

	return env, nil
}

func parseReqEnvelope(arr []json.RawMessage) (*ReqEnvelope, error) {
	if len(arr) < 3 {
		return nil, fmt.Errorf("REQ envelope must have at least 3 elements")
	}

	env := &ReqEnvelope{}
	if err := json.Unmarshal(arr[1], &env.SubscriptionID); err != nil {
		return nil, fmt.Errorf("invalid subscription ID: %w", err)
	}

	for i := 2; i < len(arr); i++ {
		var filter Filter
		if err := json.Unmarshal(arr[i], &filter); err != nil {
			return nil, fmt.Errorf("invalid filter at index %d: %w", i-2, err)
		}
		env.Filters = append(env.Filters, filter)
	}

	return env, nil
}

func parseCloseEnvelope(arr []json.RawMessage) (*CloseEnvelope, error) {
	if len(arr) != 2 {
		return nil, fmt.Errorf("CLOSE envelope must have exactly 2 elements")
	}

	env := &CloseEnvelope{}
	if err := json.Unmarshal(arr[1], &env.SubscriptionID); err != nil {
		return nil, fmt.Errorf("invalid subscription ID: %w", err)
	}

	return env, nil
}

func parseOKEnvelope(arr []json.RawMessage) (*OKEnvelope, error) {
	if len(arr) != 4 {
		return nil, fmt.Errorf("OK envelope must have exactly 4 elements")
	}

	env := &OKEnvelope{}
	if err := json.Unmarshal(arr[1], &env.EventID); err != nil {
		return nil, fmt.Errorf("invalid event ID: %w", err)
	}
	if err := json.Unmarshal(arr[2], &env.OK); err != nil {
		return nil, fmt.Errorf("invalid OK status: %w", err)
	}
	if err := json.Unmarshal(arr[3], &env.Message); err != nil {
		return nil, fmt.Errorf("invalid message: %w", err)
	}

	return env, nil
}

func parseEOSEEnvelope(arr []json.RawMessage) (*EOSEEnvelope, error) {
	if len(arr) != 2 {
		return nil, fmt.Errorf("EOSE envelope must have exactly 2 elements")
	}

	env := &EOSEEnvelope{}
	if err := json.Unmarshal(arr[1], &env.SubscriptionID); err != nil {
		return nil, fmt.Errorf("invalid subscription ID: %w", err)
	}

	return env, nil
}

func parseClosedEnvelope(arr []json.RawMessage) (*ClosedEnvelope, error) {
	if len(arr) != 3 {
		return nil, fmt.Errorf("CLOSED envelope must have exactly 3 elements")
	}

	env := &ClosedEnvelope{}
	if err := json.Unmarshal(arr[1], &env.SubscriptionID); err != nil {
		return nil, fmt.Errorf("invalid subscription ID: %w", err)
	}
	if err := json.Unmarshal(arr[2], &env.Message); err != nil {
		return nil, fmt.Errorf("invalid message: %w", err)
	}

	return env, nil
}

func parseNoticeEnvelope(arr []json.RawMessage) (*NoticeEnvelope, error) {
	if len(arr) != 2 {
		return nil, fmt.Errorf("NOTICE envelope must have exactly 2 elements")
	}

	env := &NoticeEnvelope{}
	if err := json.Unmarshal(arr[1], &env.Message); err != nil {
		return nil, fmt.Errorf("invalid message: %w", err)
	}

	return env, nil
}