handlers_add.go

package pdp

import (
	"context"
	"encoding/json"
	"errors"
	"fmt"
	"math/big"
	"net/http"
	"path"
	"strconv"
	"strings"
	"time"

	"github.com/ethereum/go-ethereum/common"
	"github.com/ethereum/go-ethereum/core/types"
	"github.com/go-chi/chi/v5"
	"github.com/ipfs/go-cid"

	commputils "github.com/filecoin-project/go-commp-utils/v2"
	commcid "github.com/filecoin-project/go-fil-commcid"
	"github.com/filecoin-project/go-state-types/abi"

	"github.com/filecoin-project/curio/harmony/harmonydb"
	"github.com/filecoin-project/curio/pdp/contract"
)

type SubPieceEntry struct {
	SubPieceCID   string `json:"subPieceCid"`
	subPieceCIDv1 string
}

type AddPieceRequest struct {
	PieceCID   string `json:"pieceCid"`
	pieceCIDv1 string
	SubPieces  []SubPieceEntry `json:"subPieces"`
}
type PieceData struct {
	Data []byte // CID
}

// Map to store subPieceCID -> [pieceInfo, pdp_pieceref.id, subPieceOffset]
type SubPieceInfo struct {
	PieceCIDv1     cid.Cid
	PaddedSize     abi.PaddedPieceSize
	RawSize        uint64 // RawSize is the size of the piece with no padding applied
	PDPPieceRefID  int64
	SubPieceOffset uint64
}

func (p *PDPService) transformAddPiecesRequest(ctx context.Context, serviceLabel string, pieces []AddPieceRequest) ([]PieceData, map[string]*SubPieceInfo, error) {
	// Collect all subPieceCids to fetch their info in a batch
	subPieceCidSet := make(map[string]struct{})
	for _, addPieceReq := range pieces {
		if addPieceReq.PieceCID == "" {
			return nil, nil, errors.New("PieceCID is required for each piece")
		}

		if len(addPieceReq.SubPieces) == 0 {
			return nil, nil, errors.New("at least one subPiece is required per piece")
		}

		for i, subPieceEntry := range addPieceReq.SubPieces {
			if subPieceEntry.SubPieceCID == "" {
				return nil, nil, errors.New("subPieceCid is required for each subPiece")
			}
			info, err := ParsePieceCid(subPieceEntry.SubPieceCID)
			if err != nil {
				return nil, nil, fmt.Errorf("invalid SubPiece: %w", err)
			}
			pieceCidString := info.CidV1.String()

			addPieceReq.SubPieces[i].subPieceCIDv1 = pieceCidString // save it for to query subPieceInfoMap later

			if _, exists := subPieceCidSet[pieceCidString]; exists {
				return nil, nil, errors.New("duplicate subPieceCid in request")
			}

			subPieceCidSet[pieceCidString] = struct{}{}
		}
	}

	// Convert set to slice
	subPieceCidList := make([]string, 0, len(subPieceCidSet))
	for cidStr := range subPieceCidSet {
		subPieceCidList = append(subPieceCidList, cidStr)
	}

	subPieceInfoMap := make(map[string]*SubPieceInfo)

	// Start a DB transaction
	_, err := p.db.BeginTransaction(ctx, func(tx *harmonydb.Tx) (bool, error) {
		// Step 4: Get pdp_piecerefs matching all subPiece cids + make sure those refs belong to serviceLabel
		rows, err := tx.Query(`
            SELECT ppr.piece_cid, ppr.id AS pdp_pieceref_id, ppr.piece_ref,
                   pp.piece_padded_size, pp.piece_raw_size
            FROM pdp_piecerefs ppr
            JOIN parked_piece_refs pprf ON pprf.ref_id = ppr.piece_ref
            JOIN parked_pieces pp ON pp.id = pprf.piece_id
            WHERE ppr.service = $1 AND ppr.piece_cid = ANY($2)
            ORDER BY ppr.created_at ASC, ppr.id ASC
        `, serviceLabel, subPieceCidList)
		if err != nil {
			return false, err
		}
		defer rows.Close()

		foundSubPieces := make(map[string]struct{})
		for rows.Next() {
			var pieceCIDStr string
			var pdpPieceRefID, pieceRefID int64
			var piecePaddedSize uint64
			var pieceRawSize uint64

			err := rows.Scan(&pieceCIDStr, &pdpPieceRefID, &pieceRefID, &piecePaddedSize, &pieceRawSize)
			if err != nil {
				return false, err
			}
			if _, found := foundSubPieces[pieceCIDStr]; found {
				continue
			}

			// Parse the piece CID
			pieceCID, err := cid.Decode(pieceCIDStr)
			if err != nil {
				return false, fmt.Errorf("invalid piece CID in database: %s", pieceCIDStr)
			}

			subPieceInfoMap[pieceCIDStr] = &SubPieceInfo{
				PieceCIDv1:     pieceCID,
				PaddedSize:     abi.PaddedPieceSize(piecePaddedSize),
				RawSize:        pieceRawSize,
				PDPPieceRefID:  pdpPieceRefID,
				SubPieceOffset: 0, // Will compute offset later
			}

			foundSubPieces[pieceCIDStr] = struct{}{}
		}

		// Check if all subPiece CIDs were found
		for _, cidStr := range subPieceCidList {
			if _, found := foundSubPieces[cidStr]; !found {
				return false, fmt.Errorf("subPiece CID %s not found or does not belong to service %s", cidStr, serviceLabel)
			}
		}

		// Now, for each AddPieceRequest, validate PieceCid and prepare data for ETH transaction
		for i, addPieceReq := range pieces {
			// Collect pieceInfos for subPieces
			pieceInfos := make([]abi.PieceInfo, len(addPieceReq.SubPieces))

			var totalOffset uint64 = 0
			for j, subPieceEntry := range addPieceReq.SubPieces {
				subPieceInfo, exists := subPieceInfoMap[subPieceEntry.subPieceCIDv1]
				if !exists {
					return false, fmt.Errorf("subPiece CID %s not found in subPiece info map", subPieceEntry.subPieceCIDv1)
				}

				// Update SubPieceOffset
				subPieceInfo.SubPieceOffset = totalOffset
				subPieceInfoMap[subPieceEntry.subPieceCIDv1] = subPieceInfo // Update the map

				pieceInfos[j] = abi.PieceInfo{
					Size:     subPieceInfo.PaddedSize,
					PieceCID: subPieceInfo.PieceCIDv1,
				}

				totalOffset += uint64(subPieceInfo.PaddedSize)
			}

			// Use GenerateUnsealedCID to generate PieceCid from subPieces
			proofType := abi.RegisteredSealProof_StackedDrg64GiBV1_1 // Proof type sets max piece size, nothing else
			generatedPieceCid, _, err := commputils.PieceAggregateCommP(proofType, pieceInfos)
			if err != nil {
				return false, fmt.Errorf("failed to generate PieceCid: %v", err)
			}

			// Compare generated PieceCid with provided PieceCid
			providedInfo, err := ParsePieceCid(addPieceReq.PieceCID)
			if err != nil {
				return false, fmt.Errorf("invalid provided PieceCid: %v", err)
			}
			pieces[i].pieceCIDv1 = providedInfo.CidV1.String()

			if !providedInfo.CidV1.Equals(generatedPieceCid) {
				return false, fmt.Errorf("provided PieceCid does not match generated PieceCid: %s != %s", providedInfo.CidV1, generatedPieceCid)
			}
		}

		// All validations passed, commit the transaction
		return true, nil
	}, harmonydb.OptionRetry())
	if err != nil {
		return nil, nil, fmt.Errorf("failed to validate subPieces: %w", err)
	}

	// Prepare PieceData array for Ethereum transaction
	// Define a Struct that matches the Solidity PieceData struct
	var pieceDataArray []PieceData

	for _, addPieceReq := range pieces {
		// Convert PieceCid to bytes
		pieceCidV2, err := cid.Decode(addPieceReq.PieceCID)
		if err != nil {
			return nil, nil, fmt.Errorf("invalid PieceCid: %w", err)
		}
		_, rawSize, err := commcid.PieceCidV1FromV2(pieceCidV2)
		if err != nil {
			return nil, nil, fmt.Errorf("invalid CommPv2: %w", err)
		}
		height, _, err := commcid.PayloadSizeToV1TreeHeightAndPadding(rawSize)
		if err != nil {
			return nil, nil, fmt.Errorf("computing height and padding: %w", err)
		}
		if height > 50 {
			return nil, nil, errors.New("invalid height")
		}

		// Get raw size by summing up the sizes of subPieces
		var totalSize uint64 = 0
		prevSubPieceSize := subPieceInfoMap[addPieceReq.SubPieces[0].subPieceCIDv1].PaddedSize
		for i, subPieceEntry := range addPieceReq.SubPieces {
			subPieceInfo := subPieceInfoMap[subPieceEntry.subPieceCIDv1]
			if subPieceInfo.PaddedSize > prevSubPieceSize {
				return nil, nil, fmt.Errorf("subPieces must be in descending order of size, piece %d %s is larger than prev subPiece %s",
					i, subPieceEntry.SubPieceCID, addPieceReq.SubPieces[i-1].SubPieceCID)
			}

			prevSubPieceSize = subPieceInfo.PaddedSize
			totalSize += uint64(subPieceInfo.RawSize)
		}
		// sanity check that the rawSize in the CommPv2 matches the totalSize of the subPieces
		if rawSize != totalSize {
			return nil, nil, fmt.Errorf("raw size mismatch: expected %d, got %d", totalSize, rawSize)
		}

		/* TODO: this doesn't work, do we need it?
		// sanity check that height and totalSize match
		computedHeight := bits.LeadingZeros64(totalSize-1) - 5
		if computedHeight != int(height) {
			http.Error(w, fmt.Sprintf("Height mismatch: expected %d, got %d for total size %d", computedHeight, height, totalSize), http.StatusBadRequest)
		}
		*/

		// Prepare PieceData for Ethereum transaction
		pieceData := PieceData{
			Data: pieceCidV2.Bytes(),
		}

		pieceDataArray = append(pieceDataArray, pieceData)
	}
	return pieceDataArray, subPieceInfoMap, nil
}

func subPieceCidV1ListFromPieces(pieces []AddPieceRequest) ([]string, error) {
	seen := make(map[string]struct{})
	list := make([]string, 0)
	for _, piece := range pieces {
		for _, subPiece := range piece.SubPieces {
			info, err := ParsePieceCid(subPiece.SubPieceCID)
			if err != nil {
				return nil, err
			}
			cidStr := info.CidV1.String()
			if _, ok := seen[cidStr]; ok {
				continue
			}
			seen[cidStr] = struct{}{}
			list = append(list, cidStr)
		}
	}
	return list, nil
}

func (p *PDPService) handleAddPieceToDataSet(w http.ResponseWriter, r *http.Request) {
	ctx := r.Context()

	workCtx, cancel := context.WithTimeout(context.Background(), 5*time.Minute)
	defer cancel()

	// Step 1: Verify that the request is authorized using ECDSA JWT
	serviceLabel, err := p.AuthService(r)
	if err != nil {
		httpServerError(w, http.StatusUnauthorized, "Unauthorized: "+err.Error(), err)
		return
	}

	// Step 2: Extract dataSetId from the URL
	dataSetIdStr := chi.URLParam(r, "dataSetId")
	if dataSetIdStr == "" {
		httpServerError(w, http.StatusBadRequest, "Missing data set ID in URL", err)
		return
	}

	// Convert dataSetId to uint64
	dataSetIdUint64, err := strconv.ParseUint(dataSetIdStr, 10, 64)
	if err != nil {
		httpServerError(w, http.StatusBadRequest, "Invalid data set ID format", err)
		return
	}

	// Step 3: Parse the request body

	// AddPiecesPayload defines the structure for the entire add pieces request payload
	type AddPiecesPayload struct {
		Pieces    []AddPieceRequest `json:"pieces"`
		ExtraData *string           `json:"extraData,omitempty"`
	}

	var payload AddPiecesPayload
	err = json.NewDecoder(r.Body).Decode(&payload)
	if err != nil {
		httpServerError(w, http.StatusBadRequest, "Invalid request body: "+err.Error(), err)
		return
	}
	defer func() {
		_ = r.Body.Close()
	}()

	if len(payload.Pieces) == 0 {
		httpServerError(w, http.StatusBadRequest, "At least one piece must be provided", err)
		return
	}
	if len(payload.Pieces) > MaxAddPiecesBatchSize {
		errMsg := fmt.Sprintf("piece count (%d) exceeds the maximum allowed per AddPieces call (%d)", len(payload.Pieces), MaxAddPiecesBatchSize)
		httpServerError(w, http.StatusBadRequest, errMsg, err)
		return
	}

	subPieceCidV1List, err := subPieceCidV1ListFromPieces(payload.Pieces)
	if err != nil {
		httpServerError(w, http.StatusBadRequest, "Invalid subPieceCid: "+err.Error(), err)
		return
	}

	if err = verifyDataSetForService(ctx, p.db, serviceLabel, dataSetIdUint64); err != nil {
		switch {
		case errors.Is(err, ErrDataSetNotFound), errors.Is(err, ErrDataSetTerminated):
			if p.recordIPOffense(w, r, OffenseBadDataSetAdd) {
				return
			}
			if discardErr := discardOrphanPiecrefsForSubPieces(ctx, p.db, serviceLabel, subPieceCidV1List); discardErr != nil {
				log.Warnw("failed to discard orphan piecerefs after bad data set addPieces",
					"dataSetId", dataSetIdUint64, "error", discardErr)
			}
			if errors.Is(err, ErrDataSetNotFound) {
				httpServerError(w, http.StatusNotFound, "Data set not found", err)
			} else {
				http.Error(w, err.Error(), http.StatusConflict)
			}
		default:
			httpServerError(w, http.StatusInternalServerError, "Failed to retrieve data set: "+err.Error(), err)
		}
		return
	}

	if p.refuseUnallowlistedAuthorizer(w, ctx, dataSetIdUint64) {
		return
	}

	// Convert dataSetId to *big.Int
	dataSetId := new(big.Int).SetUint64(dataSetIdUint64)

	extraDataBytes, err := decodeExtraData(payload.ExtraData)
	if err != nil {
		httpServerError(w, http.StatusBadRequest, "Invalid extraData format (must be hex encoded): "+err.Error(), err)
		return
	}

	// Step 4: Prepare piece information
	pieceDataArray, subPieceInfoMap, err := p.transformAddPiecesRequest(ctx, serviceLabel, payload.Pieces)
	if err != nil {
		log.Warnf("Failed to process AddPieces request data: %+v", err)
		httpServerError(w, http.StatusBadRequest, "Failed to process request: "+err.Error(), err)
		return
	}

	// Step 5: Prepare the Ethereum transaction data outside the DB transaction
	// Obtain the ABI of the PDPVerifier contract
	abiData, err := contract.PDPVerifierMetaData.GetAbi()
	if err != nil {
		httpServerError(w, http.StatusInternalServerError, "Failed to get contract ABI: "+err.Error(), err)
		return
	}

	// Step 6: Prepare the Ethereum transaction
	// Pack the method call data
	// The extraDataBytes variable is now correctly populated above
	data, err := abiData.Pack("addPieces", dataSetId, common.Address{}, pieceDataArray, extraDataBytes)
	if err != nil {
		httpServerError(w, http.StatusInternalServerError, "Failed to pack method call: "+err.Error(), err)
		return
	}

	// Step 7: Get the sender address from 'eth_keys' table where role = 'pdp' limit 1
	fromAddress, err := p.getSenderAddress(ctx)
	if err != nil {
		httpServerError(w, http.StatusInternalServerError, "Failed to get sender address: "+err.Error(), err)
		return
	}

	if err := p.preflightAuthorizerCall(ctx, fromAddress, contract.ContractAddresses().PDPVerifier, data, nil); err != nil {
		httpServerError(w, http.StatusBadRequest, "addPieces validation failed: "+err.Error(), err)
		return
	}

	// Prepare the transaction (nonce will be set to 0, SenderETH will assign it)
	txEth := types.NewTransaction(
		0,
		contract.ContractAddresses().PDPVerifier,
		big.NewInt(0),
		0,
		nil,
		data,
	)

	// Step 8: Resolve indexing intent from cached pdp_data_sets.ipfs_indexing (chain-fill if NULL).
	mustIndex, err := ResolveDatasetIPFSIndexing(workCtx, p.db, p.ethClient, dataSetIdUint64)
	if err != nil {
		log.Errorw("Failed to resolve indexing requirements", "error", err, "dataSetId", dataSetId)
		httpServerError(w, http.StatusInternalServerError, "Internal server error", err)
		return
	}
	if mustIndex {
		log.Infow("Data set has withIPFSIndexing enabled, pieces will be indexed", "dataSetId", dataSetId)
	}

	// Step 9: Send the transaction
	reason := "pdp-addpieces"
	txHash, err := p.sender.Send(workCtx, fromAddress, txEth, reason)
	if err != nil {
		log.Errorf("Failed to send transaction: %+v", err)
		httpServerError(w, http.StatusInternalServerError, "Failed to send transaction: "+err.Error(), err)
		return
	}

	// Step 10: Insert database tracking records
	txHashLower := strings.ToLower(txHash.Hex())
	log.Infow("PDP AddPieces: Inserting transaction tracking",
		"txHash", txHashLower,
		"dataSetId", dataSetIdUint64,
		"pieceCount", len(payload.Pieces))

	comm, err := p.db.BeginTransaction(workCtx, func(txdb *harmonydb.Tx) (bool, error) {
		// Insert into message_waits_eth
		log.Debugw("Inserting AddPieces into message_waits_eth",
			"txHash", txHashLower,
			"status", "pending")
		n, err := txdb.Exec(`
            INSERT INTO message_waits_eth (signed_tx_hash, tx_status)
            VALUES ($1, $2)
        `, txHashLower, "pending")
		if err != nil {
			log.Errorw("Failed to insert AddPieces into message_waits_eth",
				"txHash", txHashLower,
				"error", err)
			return false, err // Return false to rollback the transaction
		}

		if n != 1 {
			log.Errorw("Failed to insert AddPieces into message_waits_eth",
				"txHash", txHashLower,
				"expected_rows", 1,
				"actual_rows", n)
			return false, fmt.Errorf("expected 1 row to be inserted, got %d", n)
		}

		// Insert into pdp_data_set_pieces
		err = p.insertPieceAdds(txdb, &dataSetIdUint64, txHashLower, payload.Pieces, subPieceInfoMap)
		if err != nil {
			return false, err
		}

		if mustIndex {
			log.Debugw("Data set metadata exists, marking all subpieces as needing indexing", "dataSetId", dataSetId)
			subPieceRefIDs := make([]int64, 0, len(subPieceInfoMap))
			for _, info := range subPieceInfoMap {
				subPieceRefIDs = append(subPieceRefIDs, info.PDPPieceRefID)
			}
			if err := EnableIndexingForPiecesInTx(txdb, serviceLabel, subPieceRefIDs); err != nil {
				return false, err
			}
		}

		// Return true to commit the transaction
		return true, nil
	}, harmonydb.OptionRetry())
	if err != nil {
		log.Errorw("Failed to insert into database", "error", err, "txHash", txHashLower, "subPieces", subPieceInfoMap)
		httpServerError(w, http.StatusInternalServerError, "Internal server error", err)
		return
	}

	if !comm {
		log.Errorw("Failed to commit database transaction", "txHash", txHashLower)
		httpServerError(w, http.StatusInternalServerError, "Internal server error", err)
		return
	}

	// Step 10: Respond with 201 Created
	w.Header().Set("Location", path.Join("/pdp/data-sets", dataSetIdStr, "pieces/added", txHashLower))
	w.WriteHeader(http.StatusCreated)
}

func (p *PDPService) insertPieceAdds(txdb *harmonydb.Tx, dataSetId *uint64, txHash string, pieces []AddPieceRequest, subPieceInfoMap map[string]*SubPieceInfo) error {
	for addMessageIndex, addPieceReq := range pieces {
		for _, subPieceEntry := range addPieceReq.SubPieces {
			subPieceInfo := subPieceInfoMap[subPieceEntry.subPieceCIDv1]

			// Insert into pdp_data_set_pieces
			n, err := txdb.Exec(`
                    INSERT INTO pdp_data_set_piece_adds (
                        data_set,
                        piece,
                        add_message_hash,
                        add_message_index,
                        sub_piece,
                        sub_piece_offset,
                        sub_piece_size,
                        pdp_pieceref
                    )
                    VALUES ($1, $2, $3, $4, $5, $6, $7, $8)
                `,
				dataSetId,
				addPieceReq.pieceCIDv1,
				txHash,
				addMessageIndex,
				subPieceEntry.subPieceCIDv1,
				subPieceInfo.SubPieceOffset,
				subPieceInfo.PaddedSize,
				subPieceInfo.PDPPieceRefID,
			)
			if err != nil {
				return err
			}
			if n != 1 {
				return fmt.Errorf("expected 1 row to be inserted into pdp_data_set_piece_adds, got %d", n)
			}
		}
	}
	return nil
}