package normalize

import (
	"bytes"
	"encoding/json"
	"fmt"
	"io"
	"os"
	"strings"
)

// InputShape identifies which top-level input shape was parsed.
type InputShape string

const (
	ShapeObjectWithSegments InputShape = "object_with_segments"
	ShapeBareSegmentsArray  InputShape = "bare_segments_array"
)

// ParsedTranscript is the validated normalize input model.
type ParsedTranscript struct {
	Shape             InputShape
	InputSegmentCount int
	Stats             NormalizeStats
	Segments          []InputSegment
}

// InputSegment is a validated segment from normalize input.
type InputSegment struct {
	InputIndex         int
	OriginalID         *int
	StartPresent       bool
	EndPresent         bool
	SpeakerPresent     bool
	TextPresent        bool
	Start              float64
	End                float64
	Speaker            string
	Text               string
	Categories         []string
	Source             string
	SourceSegmentIndex *int
	SourceRef          string
	DerivedFrom        []string
	OverlapGroupID     *int
}

// NormalizeStats captures deterministic repair/drop outcomes from normalize input processing.
type NormalizeStats struct {
	TimingFieldsRepaired int
	TimingOrderSwapped   int
	SpeakerFilled        int
	SegmentsDroppedText  int
}

type inputSegmentPayload struct {
	ID                 *int     `json:"id"`
	Start              *float64 `json:"start"`
	End                *float64 `json:"end"`
	Speaker            *string  `json:"speaker"`
	Text               *string  `json:"text"`
	Categories         []string `json:"categories"`
	Source             string   `json:"source"`
	SourceSegmentIndex *int     `json:"source_segment_index"`
	SourceRef          string   `json:"source_ref"`
	DerivedFrom        []string `json:"derived_from"`
	OverlapGroupID     *int     `json:"overlap_group_id"`
}

// ParseFile parses normalize input JSON from file path.
func ParseFile(path string) (ParsedTranscript, error) {
	file, err := os.Open(path)
	if err != nil {
		return ParsedTranscript{}, err
	}
	defer file.Close()

	return ParseReader(file)
}

// ParseReader parses normalize input JSON from a reader.
func ParseReader(reader io.Reader) (ParsedTranscript, error) {
	var raw json.RawMessage
	decoder := json.NewDecoder(reader)
	decoder.UseNumber()
	if err := decoder.Decode(&raw); err != nil {
		return ParsedTranscript{}, fmt.Errorf("decode normalize input JSON: %w", err)
	}
	if err := ensureSingleValue(decoder); err != nil {
		return ParsedTranscript{}, err
	}

	trimmed := bytes.TrimSpace(raw)
	if len(trimmed) == 0 {
		return ParsedTranscript{}, fmt.Errorf("normalize input is empty")
	}

	switch trimmed[0] {
	case '{':
		return parseObjectShape(trimmed)
	case '[':
		segments, stats, inputCount, err := parseSegmentsArray(trimmed)
		if err != nil {
			return ParsedTranscript{}, err
		}
		return ParsedTranscript{
			Shape:             ShapeBareSegmentsArray,
			InputSegmentCount: inputCount,
			Stats:             stats,
			Segments:          segments,
		}, nil
	default:
		return ParsedTranscript{}, fmt.Errorf("normalize input must be a top-level object with \"segments\" or a top-level segment array")
	}
}

func ensureSingleValue(decoder *json.Decoder) error {
	var extra json.RawMessage
	err := decoder.Decode(&extra)
	if err == io.EOF {
		return nil
	}
	if err == nil {
		return fmt.Errorf("normalize input must contain exactly one top-level JSON value")
	}
	return fmt.Errorf("decode normalize input JSON: %w", err)
}

func parseObjectShape(raw []byte) (ParsedTranscript, error) {
	var object map[string]json.RawMessage
	if err := json.Unmarshal(raw, &object); err != nil {
		return ParsedTranscript{}, fmt.Errorf("decode normalize object input: %w", err)
	}

	segmentsRaw, exists := object["segments"]
	if !exists {
		return ParsedTranscript{}, fmt.Errorf("normalize object input must contain a \"segments\" field")
	}

	segments, stats, inputCount, err := parseSegmentsArray(segmentsRaw)
	if err != nil {
		return ParsedTranscript{}, err
	}

	return ParsedTranscript{
		Shape:             ShapeObjectWithSegments,
		InputSegmentCount: inputCount,
		Stats:             stats,
		Segments:          segments,
	}, nil
}

func parseSegmentsArray(raw []byte) ([]InputSegment, NormalizeStats, int, error) {
	var segmentValues []json.RawMessage
	if err := json.Unmarshal(raw, &segmentValues); err != nil {
		return nil, NormalizeStats{}, 0, fmt.Errorf("normalize input \"segments\" must be an array")
	}

	segments := make([]InputSegment, len(segmentValues))
	for index, segmentRaw := range segmentValues {
		segment, err := decodeSegment(index, segmentRaw)
		if err != nil {
			return nil, NormalizeStats{}, 0, err
		}
		segments[index] = segment
	}
	normalized, stats, err := normalizeSegments(segments)
	return normalized, stats, len(segmentValues), err
}

func decodeSegment(index int, raw []byte) (InputSegment, error) {
	var payload inputSegmentPayload
	if err := json.Unmarshal(raw, &payload); err != nil {
		return InputSegment{}, fmt.Errorf("segment %d: invalid segment object: %w", index, err)
	}

	var start float64
	if payload.Start != nil {
		start = *payload.Start
	}
	var end float64
	if payload.End != nil {
		end = *payload.End
	}

	speaker := ""
	if payload.Speaker != nil {
		speaker = strings.TrimSpace(*payload.Speaker)
	}

	text := ""
	if payload.Text != nil {
		text = *payload.Text
	}

	return InputSegment{
		InputIndex:         index,
		OriginalID:         payload.ID,
		StartPresent:       payload.Start != nil,
		EndPresent:         payload.End != nil,
		SpeakerPresent:     payload.Speaker != nil,
		TextPresent:        payload.Text != nil,
		Start:              start,
		End:                end,
		Speaker:            speaker,
		Text:               text,
		Categories:         append([]string(nil), payload.Categories...),
		Source:             payload.Source,
		SourceSegmentIndex: payload.SourceSegmentIndex,
		SourceRef:          payload.SourceRef,
		DerivedFrom:        append([]string(nil), payload.DerivedFrom...),
		OverlapGroupID:     payload.OverlapGroupID,
	}, nil
}

func normalizeSegments(segments []InputSegment) ([]InputSegment, NormalizeStats, error) {
	stats := NormalizeStats{}
	filtered := make([]InputSegment, 0, len(segments))
	for _, segment := range segments {
		if strings.TrimSpace(segment.Text) == "" {
			stats.SegmentsDroppedText++
			continue
		}
		filtered = append(filtered, segment)
	}

	if len(filtered) == 0 {
		return filtered, stats, nil
	}

	hasStart := make([]bool, len(filtered))
	hasEnd := make([]bool, len(filtered))
	for index, segment := range filtered {
		hasStart[index] = segment.StartPresent
		hasEnd[index] = segment.EndPresent
	}
	for index := range filtered {
		if !hasStart[index] && !hasEnd[index] {
			midpoint := inferTimestampFromNeighbors(filtered, hasStart, hasEnd, index)
			filtered[index].Start = midpoint
			filtered[index].End = midpoint
			stats.TimingFieldsRepaired += 2
			hasStart[index] = true
			hasEnd[index] = true
			continue
		}
		if hasStart[index] && !hasEnd[index] {
			filtered[index].End = filtered[index].Start
			stats.TimingFieldsRepaired++
			hasEnd[index] = true
		}
		if !hasStart[index] && hasEnd[index] {
			filtered[index].Start = filtered[index].End
			stats.TimingFieldsRepaired++
			hasStart[index] = true
		}
	}

	for index := range filtered {
		if filtered[index].Start < 0 {
			return nil, NormalizeStats{}, fmt.Errorf("segment %d has start %v; start must be >= 0", filtered[index].InputIndex, filtered[index].Start)
		}
		if filtered[index].End < filtered[index].Start {
			filtered[index].Start, filtered[index].End = filtered[index].End, filtered[index].Start
			stats.TimingOrderSwapped++
		}
		if strings.TrimSpace(filtered[index].Speaker) == "" {
			filtered[index].Speaker = "Unknown_Speaker"
			stats.SpeakerFilled++
		}
	}

	return filtered, stats, nil
}

func inferTimestampFromNeighbors(segments []InputSegment, hasStart []bool, hasEnd []bool, index int) float64 {
	prevEnd, hasPrev := nearestPreviousEnd(segments, hasEnd, index)
	nextStart, hasNext := nearestNextStart(segments, hasStart, index)
	if hasPrev && hasNext {
		return (prevEnd + nextStart) / 2
	}
	if hasPrev {
		return prevEnd
	}
	if hasNext {
		return nextStart
	}
	return 0
}

func nearestPreviousEnd(segments []InputSegment, hasEnd []bool, index int) (float64, bool) {
	for i := index - 1; i >= 0; i-- {
		if hasEnd[i] {
			return segments[i].End, true
		}
	}
	return 0, false
}

func nearestNextStart(segments []InputSegment, hasStart []bool, index int) (float64, bool) {
	for i := index + 1; i < len(segments); i++ {
		if hasStart[i] {
			return segments[i].Start, true
		}
	}
	return 0, false
}