package main

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

	"github.com/google/uuid"
	"github.com/gorilla/websocket"
)

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

var tasks = []string{}

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 searchVideo(w http.ResponseWriter, r *http.Request) {
	query := r.URL.Query().Get("q")
	if query == "" {
		http.Error(w, "Query is required.", http.StatusBadRequest)
		return
	}

	task := uuid.New().String()
	tasks = append(tasks, task)

	w.Header().Set("Content-Type", "application/json")
	w.WriteHeader(http.StatusOK)
	w.Write([]byte(task))

	fmt.Printf("Task %s created\n", task)
}

func connectSearchSocket(w http.ResponseWriter, r *http.Request) {
	taskId := r.URL.Query().Get("task")
	if !slices.Contains(tasks, taskId) {
		http.Error(w, "Task not found.", http.StatusNotFound)
		return
	}

	fmt.Printf("Task %s started\n", taskId)

	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()
	}()

	files := []string{}
	err = filepath.Walk(VIDEOS_DIR, func(path string, info os.FileInfo, err error) error {
		if err != nil {
			return err
		}
		files = append(files, path)
		return nil
	})

	if err != nil {
		// http.Error(w, "Could not read video directory.", http.StatusInternalServerError)
		return
	}

	for _, file := range files {
		if strings.Contains(file, ".mp4") {
			jsonMsg, err := json.Marshal(Video{Filename: file})
			if err != nil {
				log.Println(err)
			}

			err = conn.WriteMessage(websocket.TextMessage, jsonMsg)
			if err != nil {
				log.Println(err)
			}
		}
	}

	tasks = slices.DeleteFunc(tasks, func(task string) bool { return task == taskId })
	fmt.Printf("Task %s finished\n", taskId)
}

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", searchVideo)
	mux.HandleFunc("/search_socket", connectSearchSocket)
	mux.HandleFunc("/", notFound)

	return mux
}

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