package main

import (
	"context"
	"encoding/json"
	"errors"
	"fmt"
	"io"
	"io/ioutil"
	"log"
	"net/http"
	"net/http/httputil"
	"os"
	"os/exec"
	"strconv"
	"strings"
)

var ErrNoStreams = errors.New("no streams, file probably doesn't exist")

type mediaFormat struct {
	Type    string `json:"codec_type"`
	Codec   string `json:"codec_name"`
	Bitrate int64
}

var shouldDumpMedia = os.Getenv("DUMP_FFMPEG") != ""

func getMediaFormats(ctx context.Context, path string) (video, audio *mediaFormat, err error) {
	cmd := exec.CommandContext(
		ctx,
		"ffprobe", "-i", path,
		"-v", "quiet",
		"-show_streams",
		"-print_format", "json",
	)
	stdout, err := cmd.StdoutPipe()
	if err != nil {
		return nil, nil, err
	}
	if err := cmd.Start(); err != nil {
		return nil, nil, err
	}
	var res struct {
		Streams []struct {
			Type    string `json:"codec_type"`
			Codec   string `json:"codec_name"`
			Bitrate string `json:"bit_rate"`
			Tags    map[string]string
		}
	}
	body, _ := ioutil.ReadAll(stdout)
	if shouldDumpMedia {
		println(string(body))
	}
	if err := json.Unmarshal(body, &res); err != nil {
		return nil, nil, err
	}
	if err := cmd.Wait(); err != nil {
		return nil, nil, ErrNoStreams
	}
	for _, stream := range res.Streams {
		if stream.Type == "video" {
			bitrate, _ := strconv.ParseInt(stream.Bitrate, 10, 64)
			if bitrate == 0 {
				bitrate, _ = strconv.ParseInt(stream.Tags["BPS"], 10, 64)
			}
			video = &mediaFormat{stream.Type, stream.Codec, bitrate}
		}
		if stream.Type == "audio" {
			bitrate, _ := strconv.ParseInt(stream.Bitrate, 10, 64)
			if bitrate == 0 {
				bitrate, _ = strconv.ParseInt(stream.Tags["BPS"], 10, 64)
			}
			audio = &mediaFormat{stream.Type, stream.Codec, bitrate}
		}
	}
	return video, audio, nil
}

var (
	clientSupportsHEVC = os.Getenv("HEVC_SUPPORTED") != ""
	clientSupportsAC3  = false
)

func ffmpeg(ctx context.Context, path string, offset int64) (*exec.Cmd, error) {
	passthru := []string{"-c:v", "copy", "-c:a", "copy"}
	args := passthru

	video, audio, err := getMediaFormats(ctx, path)
	if err != nil {
		switch err {
		case ErrNoStreams:
			if strings.HasSuffix(path, ".mp4") {
				path = strings.Replace(path, ".mp4", ".mkv", 1)
				video, audio, err = getMediaFormats(ctx, path)
				if err != nil {
					return nil, err
				}
			}
		default:
			return nil, err
		}
	}

	if !clientSupportsAC3 && (audio.Codec == "ac3" || audio.Codec == "eac3") {
		// convert ac3 to aac
		convertaudio := []string{"-c:v", "copy", "-c:a", "aac", "-ac", "2"}
		args = convertaudio
	}

	if !clientSupportsHEVC && video.Codec == "hevc" {
		// convert hevc to h264
		convertboth := []string{"-preset", "superfast", "-c:v", "libx264", "-c:a", "aac", "-ac", "2"}
		args = convertboth
	}

	args = append([]string{"-i", path}, args...)
	if offset > 0 {
		totalBitrate := video.Bitrate + audio.Bitrate
		args = append(args, "-ss", fmt.Sprintf("%d", offset*8/totalBitrate))
	}
	args = append(args, "-f", "matroska", "-movflags", "emptymoov", "-movflags", "faststart", "-fflags", "fastseek", "-")
	log.Printf("\n\tffmpeg %s", strings.Join(args, ` `))
	return exec.CommandContext(
		ctx,
		"ffmpeg",
		args...,
	), nil
}

var shouldDump = os.Getenv("DUMP_TRAFFIC") != ""

func writeError(w http.ResponseWriter, err string, code int) {
	log.Printf("ERROR %d: %s", code, err)
	http.Error(w, err, 500)
}

func transcodingHandler(w http.ResponseWriter, r *http.Request) {
	if !strings.HasSuffix(r.URL.Path, ".mkv") && !strings.HasSuffix(r.URL.Path, ".mp4") {
		http.ServeFile(w, r, "."+r.URL.Path)
		return
	}
	if shouldDump {
		dump, _ := httputil.DumpRequest(r, true)
		println(string(dump))
	}
	// Cancel command when client has gone away
	ctx, cancel := context.WithCancel(r.Context())
	defer cancel()
	if n, ok := w.(http.CloseNotifier); ok {
		go func(ctx context.Context) {
			defer cancel()
			defer println("client has gone away, cancelling ffmpeg...")
			<-n.CloseNotify()
		}(ctx)
	}

	var offset int64
	rng := r.Header.Get("Range")
	if rng != "" && rng != `bytes=0-` {
		if _, err := fmt.Sscanf(rng, `bytes=%d-`, &offset); err != nil {
			log.Printf("error parsing Range header %q: %s", rng, err)
		}
	}
	cmd, err := ffmpeg(ctx, strings.TrimPrefix(r.URL.Path, "/"), offset)
	if err != nil {
		writeError(w, err.Error(), 500)
		return
	}
	if shouldDumpMedia {
		cmd.Stderr = os.Stderr
	}
	stdout, err := cmd.StdoutPipe()
	if err != nil {
		writeError(w, err.Error(), 500)
		return
	}
	if err := cmd.Start(); err != nil {
		writeError(w, err.Error(), 500)
		return
	}
	f, err := os.Stat(strings.TrimPrefix(r.URL.Path, "/"))
	if err != nil {
		writeError(w, err.Error(), 500)
		return
	}
	w.Header().Set("Content-Type", "video/x-matroska")
	size := f.Size()
	w.Header().Set("Content-Length", fmt.Sprintf("%d", size))
	w.Header().Set("Content-Range", fmt.Sprintf("bytes %d-%d/%d", offset, size-offset-1, size))
	w.Header().Set("Accept-Ranges", "bytes")

	if shouldDump {
		log.Println("HTTP/1.1 206 Partial Content")
		for name := range w.Header() {
			log.Printf("%s: %s", name, w.Header().Get(name))
		}
	}

	w.WriteHeader(http.StatusPartialContent)
	io.Copy(w, stdout)

	if err := cmd.Wait(); err != nil {
		println(err.Error())
		return
	}
}

func main() {
	port := os.Getenv("PORT")
	if port == "" {
		port = "8000"
	}
	http.HandleFunc("/", transcodingHandler)
	log.Fatal(http.ListenAndServe("0.0.0.0:"+port, nil))
}