package spool

import (
	"bytes"
	"context"
	"database/sql"
	"errors"
	"fmt"
	"time"

	"github.com/klauspost/compress/zstd"
	rvboxv1 "github.com/rvbox/rvbox/gen/go/rvbox/v1"
	"github.com/rvbox/rvbox/internal/domain"
	"google.golang.org/protobuf/proto"
)

const MaxOutputChunkBytes = 64 << 10

var ErrPinnedOutput = errors.New("client output is pinned awaiting server acknowledgement")

type OutputInput struct {
	Stream     rvboxv1.StreamKind
	Raw        []byte
	ObservedAt time.Time
}

// AppendOutput stores one raw pipe chunk as a Zstandard OutputChunk. Under
// offline pressure it replaces only the oldest unsequenced output with an
// ordered loss marker; assigned wire events are immutable and never evicted.
func (store *Store) AppendOutput(ctx context.Context, issueUUID domain.UUID, input OutputInput) (Event, error) {
	if (input.Stream != rvboxv1.StreamKind_STREAM_STDOUT && input.Stream != rvboxv1.StreamKind_STREAM_STDERR) || uint64(len(input.Raw)) > MaxOutputChunkBytes || input.ObservedAt.IsZero() {
		return Event{}, errors.New("invalid raw output chunk")
	}
	payload, err := encodeOutputChunk(input.Stream, input.Raw)
	if err != nil {
		return Event{}, err
	}
	eventInput := EventInput{Kind: EventKindOutput, Compression: 2, RawBytes: uint64(len(input.Raw)), Payload: payload, Output: true, CreatedAt: input.ObservedAt}
	for {
		event, err := store.AppendEvent(ctx, issueUUID, eventInput)
		if !errors.Is(err, ErrCapacityExhausted) {
			return event, err
		}
		rotated, rotateErr := store.RotateOldestUnassignedOutput(ctx, issueUUID, input.ObservedAt)
		if rotateErr != nil {
			return Event{}, rotateErr
		}
		if !rotated {
			return Event{}, err
		}
	}
}

// RotateOldestUnassignedOutput replaces exactly one earliest contiguous run.
// Retaining the first removed local ordinal for the marker preserves durable
// command-local order while leaving every assigned event pinned.
func (store *Store) RotateOldestUnassignedOutput(ctx context.Context, issueUUID domain.UUID, observedAt time.Time) (bool, error) {
	if !validUUID(issueUUID) || observedAt.IsZero() {
		return false, errors.New("invalid output rotation request")
	}
	tx, err := store.db.BeginTx(ctx, nil)
	if err != nil {
		return false, err
	}
	defer tx.Rollback()
	var nextOrdinal, outputCharged, totalCharged, closeout uint64
	err = tx.QueryRowContext(ctx, `SELECT next_local_ordinal, output_charged_bytes, total_charged_bytes, closeout_remaining_bytes FROM commands WHERE issue_uuid = ?`, issueUUID[:]).Scan(&nextOrdinal, &outputCharged, &totalCharged, &closeout)
	if err == sql.ErrNoRows {
		return false, ErrUnknownCommand
	}
	if err != nil {
		return false, err
	}
	rows, err := tx.QueryContext(ctx, `SELECT local_ordinal, raw_bytes, charged_bytes, payload FROM events WHERE issue_uuid = ? AND event_seq IS NULL AND output = 1 ORDER BY local_ordinal`, issueUUID[:])
	if err != nil {
		return false, err
	}
	var run []outputRow
	for rows.Next() {
		var current outputRow
		if err := rows.Scan(&current.ordinal, &current.rawBytes, &current.charge, &current.payload); err != nil {
			_ = rows.Close()
			return false, err
		}
		if len(run) > 0 && current.ordinal != run[len(run)-1].ordinal+1 {
			break
		}
		run = append(run, current)
	}
	if err := rows.Close(); err != nil {
		return false, err
	}
	if len(run) == 0 {
		return false, ErrPinnedOutput
	}
	markerPayload, removedRaw, removedCompressed, removedCharge, err := outputLossMarker(run)
	if err != nil {
		return false, err
	}
	markerCharge, err := EstimateCharge(ChargeInput{EncodedBytes: uint64(len(markerPayload)), SQLiteRows: 1, IndexEntries: 2})
	if err != nil {
		return false, err
	}
	if removedCharge > totalCharged || removedCharge > outputCharged {
		return false, fmt.Errorf("client spool output charge counter mismatch")
	}
	clientTotal, err := clientTotalCharge(ctx, tx)
	if err != nil {
		return false, err
	}
	if removedCharge > clientTotal {
		return false, fmt.Errorf("client spool aggregate charge counter mismatch")
	}
	remainingTotal := totalCharged - removedCharge
	remainingOutput := outputCharged - removedCharge
	remainingClient := clientTotal - removedCharge
	decision, err := CheckReservation(store.quotaLimits, ReservationState{CommandOutputCharged: remainingOutput, CommandTotalCharged: remainingTotal, ClientTotalCharged: remainingClient, CloseoutRemaining: closeout}, ReservationRequest{ChargedBytes: markerCharge})
	if err != nil {
		return false, err
	}
	for _, current := range run {
		if _, err := tx.ExecContext(ctx, `DELETE FROM events WHERE issue_uuid = ? AND local_ordinal = ? AND event_seq IS NULL`, issueUUID[:], current.ordinal); err != nil {
			return false, err
		}
	}
	digest := immutableDigest(markerPayload)
	if _, err := tx.ExecContext(ctx, `INSERT INTO events(issue_uuid, local_ordinal, event_kind, compression, raw_bytes, charged_bytes, output, payload, payload_sha256, created_at) VALUES (?, ?, ?, 2, ?, ?, 0, ?, ?, ?)`, issueUUID[:], run[0].ordinal, EventKindOutputTruncation, removedRaw, markerCharge, markerPayload, digest[:], observedAt.UnixNano()); err != nil {
		return false, err
	}
	if _, err := tx.ExecContext(ctx, `UPDATE commands SET output_charged_bytes = ?, total_charged_bytes = ?, closeout_remaining_bytes = ?, next_local_ordinal = ? WHERE issue_uuid = ?`, decision.CommandOutputCharged, decision.CommandTotalCharged, decision.CloseoutRemaining, nextOrdinal, issueUUID[:]); err != nil {
		return false, err
	}
	if err := updateClientTotalCharge(ctx, tx, decision.ClientTotalCharged); err != nil {
		return false, err
	}
	if err := tx.Commit(); err != nil {
		return false, err
	}
	_ = removedCompressed // retained in the durable marker payload.
	return true, nil
}

type outputRow struct {
	ordinal  uint64
	rawBytes uint64
	charge   uint64
	payload  []byte
}

func encodeOutputChunk(stream rvboxv1.StreamKind, raw []byte) ([]byte, error) {
	encoder, err := zstd.NewWriter(nil, zstd.WithEncoderConcurrency(1))
	if err != nil {
		return nil, err
	}
	compressed := encoder.EncodeAll(raw, nil)
	encoder.Close()
	chunk := &rvboxv1.OutputChunk{Stream: stream, Compression: rvboxv1.Compression_COMPRESSION_ZSTD, Data: compressed, UncompressedSize: uint64(len(raw)), CompressedSize: uint64(len(compressed))}
	return proto.Marshal(chunk)
}

func outputLossMarker(rows []outputRow) ([]byte, uint64, uint64, uint64, error) {
	var raw, compressed, charged uint64
	for _, row := range rows {
		var chunk rvboxv1.OutputChunk
		if err := proto.Unmarshal(row.payload, &chunk); err != nil || chunk.Compression != rvboxv1.Compression_COMPRESSION_ZSTD || chunk.UncompressedSize != row.rawBytes || chunk.CompressedSize != uint64(len(chunk.Data)) {
			return nil, 0, 0, 0, ErrStoredPayloadChecksum
		}
		if raw > ^uint64(0)-row.rawBytes || compressed > ^uint64(0)-chunk.CompressedSize || charged > ^uint64(0)-row.charge {
			return nil, 0, 0, 0, ErrCapacityExhausted
		}
		raw += row.rawBytes
		compressed += chunk.CompressedSize
		charged += row.charge
	}
	marker := &rvboxv1.OutputTruncation{RemovedCompressedBytes: &compressed, RemovedUncompressedBytes: raw, Reason: "CLIENT_SPOOL", Source: rvboxv1.OutputTruncationSource_OUTPUT_TRUNCATION_SOURCE_CLIENT_SPOOL}
	payload, err := proto.Marshal(marker)
	return payload, raw, compressed, charged, err
}

func decodeOutputLossMarker(payload []byte) (*rvboxv1.OutputTruncation, error) {
	var marker rvboxv1.OutputTruncation
	if err := proto.Unmarshal(payload, &marker); err != nil || marker.Source != rvboxv1.OutputTruncationSource_OUTPUT_TRUNCATION_SOURCE_CLIENT_SPOOL || marker.FirstRemovedEventSeq != nil || marker.LastRemovedEventSeq != nil || marker.RemovedCompressedBytes == nil {
		return nil, ErrStoredPayloadChecksum
	}
	return &marker, nil
}

func outputChunkRaw(payload []byte) ([]byte, error) {
	var chunk rvboxv1.OutputChunk
	if err := proto.Unmarshal(payload, &chunk); err != nil || (chunk.Stream != rvboxv1.StreamKind_STREAM_STDOUT && chunk.Stream != rvboxv1.StreamKind_STREAM_STDERR) || chunk.Compression != rvboxv1.Compression_COMPRESSION_ZSTD || chunk.CompressedSize != uint64(len(chunk.Data)) {
		return nil, ErrStoredPayloadChecksum
	}
	decoder, err := zstd.NewReader(bytes.NewReader(chunk.Data), zstd.WithDecoderConcurrency(1), zstd.WithDecoderMaxMemory(MaxOutputChunkBytes+1))
	if err != nil {
		return nil, ErrStoredPayloadChecksum
	}
	defer decoder.Close()
	raw, err := decoder.DecodeAll(chunk.Data, nil)
	if err != nil || uint64(len(raw)) != chunk.UncompressedSize || uint64(len(raw)) > MaxOutputChunkBytes {
		return nil, ErrStoredPayloadChecksum
	}
	return raw, nil
}