package orchestrator

import (
	"bytes"
	"context"
	"encoding/json"
	"fmt"
	"io"
	"net"
	"net/http"
)

type networkOverride struct {
	IfaceID     string `json:"iface_id"`
	HostDevName string `json:"host_dev_name"`
}

type snapshotLoadRequest struct {
	MemFilePath      string            `json:"mem_file_path"`
	SnapshotPath     string            `json:"snapshot_path"`
	ResumeVM         bool              `json:"resume_vm,omitempty"`
	NetworkOverrides []networkOverride `json:"network_overrides,omitempty"`
}

// loadSnapshotWithNetworkOverride calls PUT /snapshot/load on the Firecracker
// Unix socket, remapping the first network interface to tapName.
// This bypasses the SDK's LoadSnapshotHandler which doesn't expose
// network_overrides (added in Firecracker v1.15, SDK v1.0.0 omits it).
func loadSnapshotWithNetworkOverride(ctx context.Context, sockPath, memPath, vmstatePath, tapName string) error {
	payload := snapshotLoadRequest{
		MemFilePath:  memPath,
		SnapshotPath: vmstatePath,
		ResumeVM:     false, // Changed: We pause here so MMDS can be configured BEFORE Resume.
		NetworkOverrides: []networkOverride{
			{IfaceID: "1", HostDevName: tapName},
		},
	}

	data, err := json.Marshal(payload)
	if err != nil {
		return fmt.Errorf("marshal snapshot load params: %w", err)
	}

	httpClient := &http.Client{
		Transport: &http.Transport{
			DialContext: func(ctx context.Context, _, _ string) (net.Conn, error) {
				return net.Dial("unix", sockPath)
			},
		},
	}

	req, err := http.NewRequestWithContext(ctx, http.MethodPut,
		"http://localhost/snapshot/load", bytes.NewReader(data))
	if err != nil {
		return fmt.Errorf("build snapshot load request: %w", err)
	}
	req.Header.Set("Content-Type", "application/json")
	req.Header.Set("Accept", "application/json")

	resp, err := httpClient.Do(req)
	if err != nil {
		return fmt.Errorf("snapshot load request: %w", err)
	}
	defer resp.Body.Close()

	if resp.StatusCode != http.StatusNoContent {
		body, _ := io.ReadAll(resp.Body)
		return fmt.Errorf("snapshot load failed (%d): %s", resp.StatusCode, body)
	}
	return nil
}