package main

import (
	"encoding/json"
	"fmt"
	"log"
	"net/http"
	"os"
	"path/filepath"
	"strconv"
	"strings"

	"github.com/gorilla/websocket"
)

var upgrader = websocket.Upgrader{
	ReadBufferSize:  1024,
	WriteBufferSize: 4096,
	CheckOrigin: func(r *http.Request) bool {
		return true
	},
}

const VIDEOS_DIR = "./videos"

type Video struct {
	Filename string `json:"filename"`
}

func notFound(w http.ResponseWriter, r *http.Request) {
	http.Error(w, "Endpoint not found", http.StatusNotFound)
}

func connectSearchSocket(w http.ResponseWriter, r *http.Request) {
	query := r.URL.Query().Get("q")
	if query == "" {
		http.Error(w, "Query is required.", http.StatusBadRequest)
		return
	}

	conn, err := upgrader.Upgrade(w, r, nil)
	if err != nil {
		fmt.Println(err)
		http.Error(w, "Could not upgrade to websocket.", http.StatusInternalServerError)
		return
	}

	defer func() {
		err := conn.WriteMessage(websocket.CloseMessage, websocket.FormatCloseMessage(websocket.CloseNormalClosure, ""))
		if err != nil {
			log.Println("Error during websocket close: ", err)
		}
		conn.Close()
	}()

	var msg = []byte{}
	err = filepath.Walk(VIDEOS_DIR, func(path string, info os.FileInfo, err error) error {
		fmt.Println("checking " + path)
		if strings.HasSuffix(path, ".mp4") && strings.Contains(path, query) {
			fmt.Println("found video " + path)
			jsonMsg, err := json.Marshal(Video{Filename: path})
			if err != nil {
				log.Println(err)
			}
			jsonMsg = append(jsonMsg, []byte("\n")...)

			if len(msg)+len(jsonMsg) > 4096 {
				fmt.Println("sending message")
				err = conn.WriteMessage(websocket.TextMessage, msg)
				if err != nil {
					log.Println(err)
					return err
				}
				msg = jsonMsg
			} else {
				msg = append(msg, jsonMsg...)
			}
		}
		return nil
	})

	if err != nil {
		log.Println(err)
	}

	// Send remaining messages
	if len(msg) > 0 {
		err = conn.WriteMessage(websocket.TextMessage, msg)
		if err != nil {
			log.Println(err)
		}
	}

	fmt.Println("done")

}

func streamVideo(w http.ResponseWriter, r *http.Request) {
	videoPath := r.URL.Query().Get("v")
	videoFile, err := os.Open(videoPath)
	if err != nil {
		http.Error(w, "Video file not found.", http.StatusNotFound)
		return
	}
	defer videoFile.Close()

	fileInfo, err := videoFile.Stat()
	if err != nil {
		http.Error(w, "Could not obtain file info.", http.StatusInternalServerError)
		return
	}

	fileSize := fileInfo.Size()
	rangeHeader := r.Header.Get("Range")
	if rangeHeader == "" {
		http.ServeFile(w, r, videoPath)
		return
	}

	start, end, err := getStartEndRange(rangeHeader, fileSize)
	if err != nil || start > end || end >= fileSize {
		http.Error(w, "Invalid range.", http.StatusBadRequest)
		return
	}

	w.Header().Set("Content-Type", "video/mp4")
	w.Header().Set("Content-Range", fmt.Sprintf("bytes %d-%d/%d", start, end, fileSize))
	w.Header().Set("Accept-Ranges", "bytes")
	w.WriteHeader(http.StatusPartialContent)

	videoFile.Seek(start, 0)
	buf := make([]byte, end-start+1)
	videoFile.Read(buf)
	w.Write(buf)
}

func getStartEndRange(rangeHeader string, fileSize int64) (int64, int64, error) {
	rangeParts := strings.Split(rangeHeader, "=")
	if len(rangeParts) != 2 || rangeParts[0] != "bytes" {
		return 0, 0, fmt.Errorf("invalid range header")
	}

	rangeSpec := strings.Split(rangeParts[1], "-")
	start, err := strconv.ParseInt(rangeSpec[0], 10, 64)
	if err != nil {
		return 0, 0, fmt.Errorf("invalid range start")
	}

	var end int64
	if len(rangeSpec) == 2 && rangeSpec[1] != "" {
		end, err = strconv.ParseInt(rangeSpec[1], 10, 64)
		if err != nil {
			return 0, 0, fmt.Errorf("invalid range end")
		}
	} else {
		end = fileSize - 1
	}

	return start, end, nil
}

func addCORSHeaders(next http.Handler) http.Handler {
	return http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
		w.Header().Set("Access-Control-Allow-Origin", "*")
		w.Header().Set("Access-Control-Allow-Methods", "GET, POST, OPTIONS")
		w.Header().Set("Access-Control-Allow-Headers", "Content-Type, Authorization")
		if r.Method == "OPTIONS" {
			w.WriteHeader(http.StatusOK)
			return
		}
		next.ServeHTTP(w, r)
	})
}

func createMux() *http.ServeMux {
	var mux *http.ServeMux = http.NewServeMux()
	mux.HandleFunc("/video", streamVideo)
	mux.HandleFunc("/search", connectSearchSocket)
	mux.HandleFunc("/", notFound)

	return mux
}

func main() {
	log.Fatal(http.ListenAndServe(":8080", addCORSHeaders(createMux())))
}