package promptkit

import (
	"encoding/json"
	"fmt"
	"math"
	"time"
)

const (
	minDurationMilliseconds = int64(time.Duration(math.MinInt64) / time.Millisecond)
	maxDurationMilliseconds = int64(time.Duration(math.MaxInt64) / time.Millisecond)
)

// MarshalJSON implements json.Marshaler for PreparedRun. It uses RFC 3339
// timestamps, integer duration_ms, and omits zero timing values.
func (r PreparedRun) MarshalJSON() ([]byte, error) {
	var startTime, endTime *time.Time
	if !r.StartTime.IsZero() {
		startTime = &r.StartTime
	}
	if !r.EndTime.IsZero() {
		endTime = &r.EndTime
	}

	var durationMS *int64
	if r.DurationMS != 0 {
		durationMS = &r.DurationMS
	}

	return json.Marshal(preparedRunJSON{
		preparedRunJSONFields: preparedRunJSONFields(r),
		StartTime:             startTime,
		EndTime:               endTime,
		DurationMS:            durationMS,
	})
}

// MarshalJSON implements json.Marshaler for RunResult. It encodes Duration as
// integer milliseconds in duration_ms and omits zero timing values.
func (r RunResult) MarshalJSON() ([]byte, error) {
	var startTime, endTime *time.Time
	if !r.StartTime.IsZero() {
		startTime = &r.StartTime
	}
	if !r.EndTime.IsZero() {
		endTime = &r.EndTime
	}

	var durationMS *int64
	if r.Duration != 0 {
		value := r.Duration.Milliseconds()
		durationMS = &value
	}

	return json.Marshal(runResultJSON{
		runResultJSONFields: runResultJSONFields(r),
		StartTime:           startTime,
		EndTime:             endTime,
		DurationMS:          durationMS,
	})
}

// UnmarshalJSON implements json.Unmarshaler for RunResult. It decodes
// duration_ms into Duration with millisecond precision. A duration_ms outside
// the range representable by time.Duration returns an error without changing
// the receiver.
func (r *RunResult) UnmarshalJSON(data []byte) error {
	var wire runResultJSON
	if err := json.Unmarshal(data, &wire); err != nil {
		return err
	}

	result := RunResult(wire.runResultJSONFields)
	if wire.DurationMS != nil {
		if *wire.DurationMS < minDurationMilliseconds || *wire.DurationMS > maxDurationMilliseconds {
			return fmt.Errorf(
				"decode RunResult duration_ms: %d cannot be represented as time.Duration",
				*wire.DurationMS,
			)
		}
		result.Duration = time.Duration(*wire.DurationMS) * time.Millisecond
	}
	if wire.StartTime != nil {
		result.StartTime = *wire.StartTime
	}
	if wire.EndTime != nil {
		result.EndTime = *wire.EndTime
	}
	*r = result
	return nil
}

type preparedRunJSONFields PreparedRun

type preparedRunJSON struct {
	preparedRunJSONFields
	StartTime  *time.Time `json:"start_time,omitempty"`
	EndTime    *time.Time `json:"end_time,omitempty"`
	DurationMS *int64     `json:"duration_ms,omitempty"`
}

type runResultJSONFields RunResult

type runResultJSON struct {
	runResultJSONFields
	StartTime  *time.Time `json:"start_time,omitempty"`
	EndTime    *time.Time `json:"end_time,omitempty"`
	DurationMS *int64     `json:"duration_ms,omitempty"`
}