package promptkit

import (
	"encoding/json"
	"time"
)

// 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(struct {
		PromptID             string                `json:"prompt_id"`
		PromptVersion        string                `json:"prompt_version,omitempty"`
		PromptHash           string                `json:"prompt_hash,omitempty"`
		SelectedProfileID    string                `json:"selected_profile_id"`
		SelectedBackendID    string                `json:"selected_backend_id,omitempty"`
		EffectiveModelParams ExecutionTarget       `json:"effective_model_params"`
		OutputContract       OutputContract        `json:"output_contract"`
		StructuredOutput     *StructuredOutputSpec `json:"structured_output,omitempty"`
		InputHashes          map[string]string     `json:"input_hashes,omitempty"`
		SessionID            string                `json:"session_id,omitempty"`
		RenderedPromptHash   string                `json:"rendered_prompt_hash"`
		Messages             []RenderedMessage     `json:"messages"`
		StartTime            *time.Time            `json:"start_time,omitempty"`
		EndTime              *time.Time            `json:"end_time,omitempty"`
		DurationMS           *int64                `json:"duration_ms,omitempty"`
	}{
		PromptID:             r.PromptID,
		PromptVersion:        r.PromptVersion,
		PromptHash:           r.PromptHash,
		SelectedProfileID:    r.SelectedProfileID,
		SelectedBackendID:    r.SelectedBackendID,
		EffectiveModelParams: r.EffectiveModelParams,
		OutputContract:       r.OutputContract,
		StructuredOutput:     r.StructuredOutput,
		InputHashes:          r.InputHashes,
		SessionID:            r.SessionID,
		RenderedPromptHash:   r.RenderedPromptHash,
		Messages:             r.Messages,
		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{
		RunID:                r.RunID,
		Artifact:             r.Artifact,
		RawOutput:            r.RawOutput,
		Validation:           r.Validation,
		PromptID:             r.PromptID,
		PromptVersion:        r.PromptVersion,
		PromptHash:           r.PromptHash,
		SessionID:            r.SessionID,
		RenderedPromptHash:   r.RenderedPromptHash,
		SelectedProfileID:    r.SelectedProfileID,
		SelectedBackendID:    r.SelectedBackendID,
		ModelName:            r.ModelName,
		Endpoint:             r.Endpoint,
		EffectiveModelParams: r.EffectiveModelParams,
		InputHashes:          r.InputHashes,
		Usage:                r.Usage,
		StartTime:            startTime,
		EndTime:              endTime,
		DurationMS:           durationMS,
	})
}

// UnmarshalJSON implements json.Unmarshaler for RunResult. It decodes
// duration_ms into Duration with millisecond precision.
func (r *RunResult) UnmarshalJSON(data []byte) error {
	var wire runResultJSON
	if err := json.Unmarshal(data, &wire); err != nil {
		return err
	}

	*r = RunResult{
		RunID:                wire.RunID,
		Artifact:             wire.Artifact,
		RawOutput:            wire.RawOutput,
		Validation:           wire.Validation,
		PromptID:             wire.PromptID,
		PromptVersion:        wire.PromptVersion,
		PromptHash:           wire.PromptHash,
		SessionID:            wire.SessionID,
		RenderedPromptHash:   wire.RenderedPromptHash,
		SelectedProfileID:    wire.SelectedProfileID,
		SelectedBackendID:    wire.SelectedBackendID,
		ModelName:            wire.ModelName,
		Endpoint:             wire.Endpoint,
		EffectiveModelParams: wire.EffectiveModelParams,
		InputHashes:          wire.InputHashes,
		Usage:                wire.Usage,
		Duration:             time.Duration(valueOrZero(wire.DurationMS)) * time.Millisecond,
	}
	if wire.StartTime != nil {
		r.StartTime = *wire.StartTime
	}
	if wire.EndTime != nil {
		r.EndTime = *wire.EndTime
	}
	return nil
}

type runResultJSON struct {
	RunID                string            `json:"run_id"`
	Artifact             Artifact          `json:"artifact"`
	RawOutput            string            `json:"raw_output"`
	Validation           ValidationResult  `json:"validation"`
	PromptID             string            `json:"prompt_id"`
	PromptVersion        string            `json:"prompt_version,omitempty"`
	PromptHash           string            `json:"prompt_hash,omitempty"`
	SessionID            string            `json:"session_id,omitempty"`
	RenderedPromptHash   string            `json:"rendered_prompt_hash"`
	SelectedProfileID    string            `json:"selected_profile_id"`
	SelectedBackendID    string            `json:"selected_backend_id,omitempty"`
	ModelName            string            `json:"model_name"`
	Endpoint             string            `json:"endpoint"`
	EffectiveModelParams ExecutionTarget   `json:"effective_model_params"`
	InputHashes          map[string]string `json:"input_hashes,omitempty"`
	Usage                TokenUsage        `json:"usage"`
	StartTime            *time.Time        `json:"start_time,omitempty"`
	EndTime              *time.Time        `json:"end_time,omitempty"`
	DurationMS           *int64            `json:"duration_ms,omitempty"`
}

func valueOrZero(value *int64) int64 {
	if value == nil {
		return 0
	}
	return *value
}