package schema

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

type SourceSegment struct {
	ID         *int     `json:"id,omitempty"`
	Speaker    string   `json:"speaker"`
	Start      float64  `json:"start"`
	End        float64  `json:"end"`
	Text       string   `json:"text"`
	Categories []string `json:"categories,omitempty"`
}

type Segment struct {
	ID         int      `json:"id"`
	Speaker    string   `json:"speaker"`
	Start      float64  `json:"start"`
	End        float64  `json:"end"`
	Text       string   `json:"text"`
	Categories []string `json:"categories,omitempty"`
}

type Transcript struct {
	Segments []Segment `json:"segments"`
}

type SourceTranscript struct {
	Segments []SourceSegment `json:"segments"`
}

func ParseSourceTranscriptJSON(raw []byte) (*SourceTranscript, error) {
	if !json.Valid(raw) {
		return nil, &ParseError{Message: "transcript is not valid JSON"}
	}

	var top any
	if err := json.Unmarshal(raw, &top); err != nil {
		return nil, &ParseError{Message: fmt.Sprintf("failed to parse transcript JSON: %v", err)}
	}

	segmentsRaw, err := extractSegments(top)
	if err != nil {
		return nil, err
	}

	var segments []SourceSegment
	if err := json.Unmarshal(segmentsRaw, &segments); err != nil {
		return nil, &ParseError{Message: fmt.Sprintf("failed to parse segments: %v", err)}
	}
	if len(segments) == 0 {
		return nil, &ParseError{Message: "transcript must contain at least one segment"}
	}

	if err := validateSourceSegments(segments); err != nil {
		return nil, err
	}

	return &SourceTranscript{Segments: segments}, nil
}

func ParseTranscriptJSON(raw []byte) (*Transcript, error) {
	source, err := ParseSourceTranscriptJSON(raw)
	if err != nil {
		return nil, err
	}

	segments := make([]Segment, len(source.Segments))
	for i, s := range source.Segments {
		segments[i] = Segment{
			ID:         *s.ID,
			Speaker:    s.Speaker,
			Start:      s.Start,
			End:        s.End,
			Text:       s.Text,
			Categories: s.Categories,
		}
	}

	if err := validateSequentialIDs(segments); err != nil {
		return nil, err
	}

	return &Transcript{Segments: segments}, nil
}

func ParseTranscriptJSONLenient(raw []byte) (*Transcript, error) {
	source, err := ParseSourceTranscriptJSON(raw)
	if err != nil {
		return nil, err
	}

	segments := make([]Segment, len(source.Segments))
	for i, s := range source.Segments {
		id := i + 1
		if s.ID != nil {
			id = *s.ID
		}
		segments[i] = Segment{
			ID:         id,
			Speaker:    s.Speaker,
			Start:      s.Start,
			End:        s.End,
			Text:       s.Text,
			Categories: s.Categories,
		}
	}

	return &Transcript{Segments: segments}, nil
}

func extractSegments(top any) (json.RawMessage, error) {
	switch v := top.(type) {
	case []any:
		raw, err := json.Marshal(v)
		if err != nil {
			return nil, &ParseError{Message: "failed to re-encode segment array"}
		}
		return raw, nil
	case map[string]any:
		segs, ok := v["segments"]
		if !ok {
			return nil, &ParseError{Message: "transcript object must contain a segments array"}
		}
		segsArray, ok := segs.([]any)
		if !ok {
			return nil, &ParseError{Message: "segments field must be an array"}
		}
		raw, err := json.Marshal(segsArray)
		if err != nil {
			return nil, &ParseError{Message: "failed to re-encode segments array"}
		}
		return raw, nil
	default:
		return nil, &ParseError{Message: "transcript must be a JSON array or an object with a segments array"}
	}
}

func validateSourceSegments(segments []SourceSegment) error {
	for i, s := range segments {
		segLabel := segmentLabel(i, s.ID)

		if s.Speaker == "" {
			return &ValidationError{Field: fmt.Sprintf("%s.speaker", segLabel), Message: "must not be empty"}
		}

		if s.Text == "" {
			return &ValidationError{Field: fmt.Sprintf("%s.text", segLabel), Message: "must not be empty"}
		}

		if err := validateTime(s.Start, fmt.Sprintf("%s.start", segLabel)); err != nil {
			return err
		}
		if err := validateTime(s.End, fmt.Sprintf("%s.end", segLabel)); err != nil {
			return err
		}

		if s.End < s.Start {
			return &ValidationError{
				Field:   fmt.Sprintf("%s.end", segLabel),
				Message: fmt.Sprintf("end (%g) must be greater than or equal to start (%g)", s.End, s.Start),
			}
		}

		for j, cat := range s.Categories {
			if cat == "" {
				return &ValidationError{
					Field:   fmt.Sprintf("%s.categories[%d]", segLabel, j),
					Message: "must not be empty",
				}
			}
		}
	}

	seenIDs := make(map[int]int)
	for i, s := range segments {
		if s.ID != nil {
			if firstIdx, exists := seenIDs[*s.ID]; exists {
				return &ValidationError{
					Field:   fmt.Sprintf("segment[%d].id", i),
					Message: fmt.Sprintf("duplicate id %d (first used at segment[%d])", *s.ID, firstIdx),
				}
			}
			seenIDs[*s.ID] = i
		}
	}

	return nil
}

func validateTime(t float64, field string) error {
	if math.IsNaN(t) || math.IsInf(t, 0) {
		return &ValidationError{Field: field, Message: "must be a finite number"}
	}
	if t < 0 {
		return &ValidationError{Field: field, Message: "must be non-negative"}
	}
	return nil
}

func validateSequentialIDs(segments []Segment) error {
	for i, s := range segments {
		expected := i + 1
		if s.ID != expected {
			return &ValidationError{
				Field:   fmt.Sprintf("segment[%d].id", i),
				Message: fmt.Sprintf("must be sequential starting at 1 (got %d, expected %d)", s.ID, expected),
			}
		}
	}
	return nil
}

func segmentLabel(index int, id *int) string {
	if id != nil {
		return fmt.Sprintf("segment[%d] (id=%d)", index, *id)
	}
	return fmt.Sprintf("segment[%d]", index)
}

func TranscriptToJSON(t *Transcript) ([]byte, error) {
	payload := make([]map[string]any, len(t.Segments))
	for i, s := range t.Segments {
		payload[i] = map[string]any{
			"id":      s.ID,
			"speaker": s.Speaker,
			"start":   s.Start,
			"end":     s.End,
			"text":    s.Text,
		}
		if len(s.Categories) > 0 {
			payload[i]["categories"] = s.Categories
		}
	}
	return json.MarshalIndent(payload, "", "  ")
}