package main

import (
	"context"
	"encoding/json"
	"fmt"
	"io"
	"net/http"
	"net/url"
	"strings"
	"time"
)

type whisperSeg struct {
	Start float64 `json:"start"`
	End   float64 `json:"end"`
	Text  string  `json:"text"`
}

func stamp(sec float64, comma bool) string {
	if sec < 0 {
		sec = 0
	}
	ms := int64(sec*1000 + 0.5)
	sep := "."
	if comma {
		sep = ","
	}
	return fmt.Sprintf("%02d:%02d:%02d%s%03d", ms/3600000, ms/60000%60, ms/1000%60, sep, ms%1000)
}

func toSRT(segs []whisperSeg) string {
	var b strings.Builder
	for i, s := range segs {
		fmt.Fprintf(&b, "%d\n%s --> %s\n%s\n\n", i+1, stamp(s.Start, true), stamp(s.End, true), s.Text)
	}
	return b.String()
}

func toVTT(segs []whisperSeg) string {
	var b strings.Builder
	b.WriteString("WEBVTT\n\n")
	for _, s := range segs {
		fmt.Fprintf(&b, "%s --> %s\n%s\n\n", stamp(s.Start, false), stamp(s.End, false), s.Text)
	}
	return b.String()
}

// runWhisper sends the audio to the Whisper engine and stores the transcript as text, subtitles and JSON.
func runWhisper(ctx context.Context, base, id string, in map[string]any, report func(string, float64)) (*jobResult, error) {
	name, _ := in["audio"].(string)
	lang, _ := in["language"].(string)
	task, _ := in["task"].(string)
	if lang == "" {
		lang = "auto"
	}
	if task == "" {
		task = "transcribe"
	}
	body, err := serverGet(ctx, id, name)
	if err != nil {
		return nil, err
	}
	defer body.Close()
	report("transcribing", -1)
	q := url.Values{"language": {lang}, "task": {task}}
	req, _ := http.NewRequestWithContext(ctx, "POST", base+"/transcribe?"+q.Encode(), body)
	req.Header.Set("Content-Type", "application/octet-stream")
	resp, err := (&http.Client{Timeout: 2 * time.Hour}).Do(req)
	if err != nil {
		return nil, err
	}
	defer resp.Body.Close()
	raw, _ := io.ReadAll(io.LimitReader(resp.Body, 64<<20))
	var r struct {
		Text     string       `json:"text"`
		Language string       `json:"language"`
		Segments []whisperSeg `json:"segments"`
		Error    string       `json:"error"`
	}
	json.Unmarshal(raw, &r)
	if resp.StatusCode != 200 {
		if r.Error == "" {
			r.Error = fmt.Sprintf("HTTP %d", resp.StatusCode)
		}
		return nil, fmt.Errorf("speech engine: %s", r.Error)
	}
	res := &jobResult{Text: r.Text}
	if len(res.Text) > 20000 {
		res.Text = res.Text[:20000]
	}
	for _, f := range []struct{ name, data string }{
		{"transcript.txt", r.Text + "\n"}, {"transcript.srt", toSRT(r.Segments)}, {"transcript.vtt", toVTT(r.Segments)}, {"transcript.json", string(raw)},
	} {
		if err := serverPut(ctx, id, f.name, strings.NewReader(f.data), int64(len(f.data))); err != nil {
			return nil, err
		}
		res.Outputs = append(res.Outputs, jobOutput{Name: f.name, Size: int64(len(f.data)), Mime: mimeOf(f.name)})
	}
	return res, nil
}