package filestream

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

	"github.com/hashicorp/go-retryablehttp"

	"github.com/wandb/wandb/core/internal/wboperation"
)

// startProcessingUpdates asynchronously ingests updates.
//
// This returns a channel of requests to send.
func (fs *fileStream) startProcessingUpdates(
	updates <-chan Update,
) <-chan *FileStreamRequest {
	requests := make(chan *FileStreamRequest)

	go func() {
		defer close(requests)

		fs.logger.Debug("filestream: open", "path", fs.path)

		for update := range updates {
			err := update.Apply(UpdateContext{
				MakeRequest: func(req *FileStreamRequest) {
					requests <- req
				},

				Settings: fs.settings,

				Logger:  fs.logger,
				Printer: fs.printer,
			})

			if err != nil {
				fs.logFatalAndStopWorking(err)
				break
			}
		}

		// Flush input channel if we exited early.
		for range updates {
		}
	}()

	return requests
}

// startTransmitting makes requests to the filestream API.
//
// It ingests a channel of requests and outputs a channel of API responses.
//
// Requests are batched to reduce the total number of HTTP requests.
// An empty "heartbeat" request is sent when there are no updates for too long,
// guaranteeing that a request is sent at least once every period specified
// by `heartbeatStopwatch`.
func (fs *fileStream) startTransmitting(
	requests <-chan *FileStreamRequest,
	initialOffsets FileStreamOffsetMap,
) <-chan map[string]any {
	state := &FileStreamState{
		MaxRequestSizeBytes: max(
			int(fs.settings.GetFileStreamMaxBytes()),
			defaultMaxRequestSizeBytes,
		),
		MaxFileLineSize: max(
			int(fs.settings.GetFileStreamMaxLineBytes()),
			defaultMaxFileLineBytes,
		),
	}

	if initialOffsets != nil {
		state.HistoryLineNum = initialOffsets[HistoryChunk]
		state.EventsLineNum = initialOffsets[EventsChunk]
		state.SummaryLineNum = initialOffsets[SummaryChunk]
		state.ConsoleLineOffset = initialOffsets[OutputChunk]
	}

	transmissions := CollectLoop{
		Logger:            fs.logger,
		Printer:           fs.printer,
		TransmitRateLimit: fs.transmitRateLimit,
	}.Start(state, requests)

	feedback := TransmitLoop{
		HeartbeatStopwatch:     fs.heartbeatStopwatch,
		Send:                   fs.send,
		LogFatalAndStopWorking: fs.logFatalAndStopWorking,
	}.Start(transmissions)

	return feedback
}

// startProcessingFeedback processes feedback from the filestream API.
//
// This increments the wait group and decrements it after completing
// all work.
func (fs *fileStream) startProcessingFeedback(
	feedback <-chan map[string]any,
	wg *sync.WaitGroup,
) {
	wg.Add(1)
	go func() {
		defer wg.Done()

		for res := range feedback {
			if v, ok := res["stopped"]; ok {
				if b, ok := v.(bool); ok {
					fs.stopState.CompareAndSwap(false, b)
				}
			}
		}
	}()
}

func (fs *fileStream) send(
	data *FileStreamRequestJSON,
	feedbackChan chan<- map[string]any,
) error {
	// Stop working after death to avoid data corruption.
	if fs.isDead() {
		return fmt.Errorf("filestream: can't send because I am dead")
	}

	jsonData, err := json.Marshal(data)
	if err != nil {
		return fmt.Errorf("filestream: json marshal error in send(): %v", err)
	}
	fs.logger.Debug("filestream: post request", "request", string(jsonData))

	op := fs.trackUploadOperation(data)
	defer op.Finish()

	req, err := retryablehttp.NewRequestWithContext(
		op.Context(context.Background()),
		http.MethodPost,
		fs.baseURL.JoinPath(fs.path).String(),
		bytes.NewReader(jsonData),
	)
	if err != nil {
		return fmt.Errorf("filestream: error constructing request: %v", err)
	}
	req.Header.Set("Content-Type", "application/json")

	resp, err := fs.apiClient.Do(req)

	switch {
	case err != nil:
		return fmt.Errorf(
			"filestream: error making HTTP request: %v. got response: %v",
			err,
			resp,
		)
	case resp.StatusCode < 200 || resp.StatusCode > 300:
		// If we reach here, that means all retries were exhausted. This could
		// mean, for instance, that the user's internet connection broke.
		body, _ := io.ReadAll(io.LimitReader(resp.Body, 1<<10))
		_ = resp.Body.Close()

		return fmt.Errorf(
			"filestream: failed to upload: %v url=%v: %s",
			resp.Status,
			req.URL,
			string(body),
		)
	}

	defer func(Body io.ReadCloser) {
		if err = Body.Close(); err != nil {
			fs.logger.CaptureError(
				fmt.Errorf("filestream: error closing response body: %v", err))
		}
	}(resp.Body)

	var res map[string]interface{}
	err = json.NewDecoder(resp.Body).Decode(&res)
	if err != nil {
		fs.logger.CaptureError(
			fmt.Errorf("filestream: json decode error: %v", err))
	}
	feedbackChan <- res
	fs.logger.Debug("filestream: post response", "response", res)
	return nil
}

// trackUploadOperation returns a WandbOperation for tracking
// a filestream upload.
func (fs *fileStream) trackUploadOperation(
	data *FileStreamRequestJSON,
) *wboperation.WandbOperation {
	parts := make([]string, 0, 3)

	if history, ok := data.Files[HistoryFileName]; ok && len(history.Content) > 0 {
		parts = append(parts,
			fmt.Sprintf("history steps %d-%d",
				history.Offset,
				history.Offset+len(history.Content)-1))
	}

	if summary, ok := data.Files[SummaryFileName]; ok && len(summary.Content) > 0 {
		parts = append(parts, "summary")
	}

	if console, ok := data.Files[OutputFileName]; ok && len(console.Content) > 0 {
		parts = append(parts,
			fmt.Sprintf("console lines %d-%d",
				console.Offset,
				console.Offset+len(console.Content)-1))
	}

	if len(parts) == 0 {
		// Shouldn't happen, but guard against future bugs.
		return fs.operations.New("uploading data")
	} else {
		return fs.operations.New("uploading " + strings.Join(parts, ", "))
	}
}
