/*
 * HTTP tunnel over WebRTC - Client implementation
 * Author: Konstantinos Drakontidis
 * Email: gedra100sh@gmail.com
 */
package main

import (
	"bufio"
	"bytes"
	"encoding/gob"
	"encoding/json"
	"fmt"
	"io"
	"net/http"
	"net/url"
	"os"
	"os/signal"
	"sync"
	"syscall"

	"github.com/go-yaml/yaml"
	"github.com/google/uuid"
	"github.com/gorilla/websocket"
	"github.com/pion/webrtc/v4"
)

var config_t Config
var http_port string

/*
 * This function gets:
 *   id: uuid
 *   p2p: pointer to the peer connection
 *   conn: pointer to the websocket connection
 * Gracefully quits the program by closing connections.
 */
func quit(id uuid.UUID, p2p *webrtc.PeerConnection, conn *websocket.Conn) {
	if p2p != nil {
		p2p_close_error := p2p.Close()
		if p2p_close_error != nil {
			fmt.Fprintf(os.Stderr, config_t.Messages.WEBRTC.CLOSE.ERROR, p2p_close_error)
		} else {
			fmt.Print(config_t.Messages.WEBRTC.CLOSE.SUCCESS)
		}
	}

	ws_close_msg := websocket.FormatCloseMessage(1000, id.String())
	ws_close_msg_error := conn.WriteMessage(websocket.CloseMessage, ws_close_msg)
	if ws_close_msg_error != nil {
		fmt.Fprintf(os.Stderr, config_t.Messages.WS.CLOSE_MESSAGE.ERROR, ws_close_msg_error)
	} else {
		fmt.Print(config_t.Messages.WS.CLOSE_MESSAGE.SUCCESS)
	}

	ws_close_error := conn.Close()
	if ws_close_error != nil {
		fmt.Fprintf(os.Stderr, config_t.Messages.WS.CLOSE.ERROR, ws_close_error)
	} else {
		fmt.Print(config_t.Messages.WS.CLOSE.SUCCESS)
	}

	os.Exit(0)
}

/*
 * This function gets:
 *   pending_mutex: pointer to mutex
 *   pending: map with response channels
 *   dc: pointer to data channel
 * Creates a HTTP server. Sends HTTP requests to data channel, waits and writes HTTP responses.
 * The function does not return anything.
 */
func httpHandler(pending_mutex *sync.Mutex, pending map[string]chan Response,
	dc *webrtc.DataChannel) {
	// Creating the HTTP server
	http.HandleFunc("/", func(w http.ResponseWriter, r *http.Request) {
		// Generating a HTTP request id
		id := uuid.New().String()

		// Reading HTTP request body
		req_body, read_error := io.ReadAll(r.Body)
		if read_error != nil {
			fmt.Fprintf(os.Stderr, config_t.Messages.HTTP.READ.ERROR, read_error)
			os.Exit(-1)
		} else {
			fmt.Print(config_t.Messages.HTTP.READ.SUCCESS)
		}

		// Constructing HTTP request struct
		req_t := Request{
			ID:      id,
			Method:  r.Method,
			Path:    r.URL.Path + "?" + r.URL.RawQuery,
			Headers: r.Header,
			Body:    req_body,
		}

		// Creating a response channel and adding it to the map
		res_chan := make(chan Response)
		(*pending_mutex).Lock()
		pending[id] = res_chan
		(*pending_mutex).Unlock()

		// Encoding HTTP request struct and sending to data channel
		var buf bytes.Buffer
		enc := gob.NewEncoder(&buf)

		encode_error := enc.Encode(req_t)
		if encode_error != nil {
			fmt.Fprintf(os.Stderr, config_t.Messages.GOB.ENCODE.ERROR, encode_error)
			os.Exit(-1)
		} else {
			fmt.Print(config_t.Messages.GOB.ENCODE.SUCCESS)
		}

		data_channel_send_error := dc.Send(buf.Bytes())
		if data_channel_send_error != nil {
			fmt.Fprintf(os.Stderr, config_t.Messages.WEBRTC.SEND_DATA_CHANNEL.ERROR,
				data_channel_send_error)
			os.Exit(-1)
		} else {
			fmt.Print(config_t.Messages.WEBRTC.SEND_DATA_CHANNEL.SUCCESS)
		}

		// Waiting until HTTP response arrives
		response := <-res_chan

		// Setting HTTP headers and writing response
		for key, vals := range response.Headers {
			for _, val := range vals {
				// Updating Location header - needed for HTTP redirects
				if key == "Location" {
					redirect_url, url_parse_error := url.Parse(val)
					if url_parse_error != nil {
						fmt.Fprintf(os.Stderr, config_t.Messages.HTTP.URL_PARSE.ERROR, url_parse_error)
						os.Exit(-1)
					} else {
						fmt.Print(config_t.Messages.HTTP.URL_PARSE.SUCCESS)
						redirect_url.Host = redirect_url.Hostname() + ":" + http_port
						val = redirect_url.String()
					}
				}
				w.Header().Add(key, val)
			}
		}
		w.WriteHeader(response.Status)
		_, http_write_error := w.Write(response.Body)
		if http_write_error != nil {
			fmt.Fprintf(os.Stderr, config_t.Messages.HTTP.WRITE.ERROR, http_write_error)
			os.Exit(-1)
		} else {
			fmt.Print(config_t.Messages.HTTP.WRITE.SUCCESS)
		}

		// Removing response channel for answered requests
		(*pending_mutex).Lock()
		delete(pending, id)
		(*pending_mutex).Unlock()
	})

	http_listen_error := http.ListenAndServe(":"+http_port, nil)
	if http_listen_error != nil {
		fmt.Fprintf(os.Stderr, config_t.Messages.HTTP.LISTEN.ERROR, http_listen_error)
		os.Exit(-1)
	} else {
		fmt.Print(config_t.Messages.HTTP.LISTEN.SUCCESS)
	}
}

/*
 * This function gets:
 *   pending_mutex: pointer to mutex
 *   pending: map with response channels
 *   raw: slice of bytes
 * Sends the message to the corresponding channel from the map.
 * The function does not return anything.
 */
func msgHandler(pending_mutex *sync.Mutex, pending map[string]chan Response, raw []byte) {
	// Decoding bytes and storing them to struct
	buf := bytes.NewBuffer(raw)
	dec := gob.NewDecoder(buf)

	var response_t Response

	decode_error := dec.Decode(&response_t)
	if decode_error != nil {
		fmt.Fprintf(os.Stderr, config_t.Messages.GOB.DECODE.ERROR, decode_error)
		os.Exit(-1)
	} else {
		fmt.Print(config_t.Messages.GOB.DECODE.SUCCESS)
	}

	// Locking the map before indexing and unlocking
	pending_mutex.Lock()
	response_channel, exists := pending[response_t.ID]
	if exists {
		response_channel <- response_t
	}
	pending_mutex.Unlock()
}

/*
 * This function gets:
 *   id: pointer to uuid
 *   p2p: pointer to the pointer of the peer connection
 *   offer: pointer to a session description
 *   conn: pointer to the websocket connection
 *   raw: slice of bytes
 * Creates a new peer connection and a data channel and sets the listeners for them. Creates a
 * peer connection offer and sends it to the signaling server.
 * The function returns with pointers:
 *   id: set to the id received from the signaling server
 *   p2p: set to the peer connection that was created
 *   offer: set to the offer that was created
 */
func peerInit(id *uuid.UUID, p2p **webrtc.PeerConnection, offer *webrtc.SessionDescription,
	conn *websocket.Conn, raw []byte) {
	/*
	 * Creating a map with HTTP response channels and a mutex variable.
	 * These channels are used to return the HTTP response received from the data channel
	 * after a HTTP request has been made. Each request is identified by a uuid. The mutex
	 * locks the map so it is modified by one listener function at a time.
	 */
	var pending_mutex sync.Mutex
	pending := make(map[string]chan Response)

	// Decoding the received message
	var id_t ID
	message_unmarshal_error := json.Unmarshal(raw, &id_t)
	if message_unmarshal_error != nil {
		fmt.Fprintf(os.Stderr, config_t.Messages.JSON.UNMARSHAL.ERROR, message_unmarshal_error)
		os.Exit(-1)
	} else {
		fmt.Print(config_t.Messages.JSON.UNMARSHAL.SUCCESS)
	}
	*id = id_t.Id

	// Creating the peer configuration
	config := webrtc.Configuration{
		ICEServers: config_t.Config.ICE_SERVERS,
	}

	// Creating peer connection
	var new_peer_error error
	*p2p, new_peer_error = webrtc.NewPeerConnection(config)
	if new_peer_error != nil {
		fmt.Fprintf(os.Stderr, config_t.Messages.WEBRTC.NEW.ERROR, new_peer_error)
		os.Exit(-1)
	} else {
		fmt.Print(config_t.Messages.WEBRTC.NEW.SUCCESS)
	}

	// Setting peer connection event handlers
	(*p2p).OnICEConnectionStateChange(func(is webrtc.ICEConnectionState) {
		fmt.Printf("ICE: %s\n", is.String())
	})
	(*p2p).OnConnectionStateChange(func(pcs webrtc.PeerConnectionState) {
		fmt.Printf("WebRTC: %s\n", pcs.String())
	})
	(*p2p).OnICECandidate(func(candidate *webrtc.ICECandidate) {
		// Checking if ICE gathering has finished
		if candidate != nil {
			fmt.Print(config_t.Messages.WEBRTC.ON_CANDIDATE)
			// Constructing candidate, encoding to JSON and sending to websocket
			candidate_t := Candidate{
				Type:      "candidate",
				Id:        id_t.Id,
				Candidate: (*candidate).ToJSON(),
			}
			candidate_msg, message_marshal_error := json.Marshal(candidate_t)
			if message_marshal_error != nil {
				fmt.Fprintf(os.Stderr, config_t.Messages.JSON.MARSHAL.ERROR, message_marshal_error)
				os.Exit(-1)
			} else {
				fmt.Print(config_t.Messages.JSON.MARSHAL.SUCCESS)
			}
			ws_write_error := conn.WriteMessage(websocket.TextMessage, candidate_msg)
			if ws_write_error != nil {
				fmt.Fprintf(os.Stderr, config_t.Messages.WS.WRITE.ERROR, ws_write_error)
			} else {
				fmt.Print(config_t.Messages.WS.WRITE.SUCCESS)
			}
		}
	})
	// Creating data channel and setting event handlers
	dc, create_data_channel_error := (*p2p).CreateDataChannel(config_t.Config.DATA_CHANNEL,
		&webrtc.DataChannelInit{})
	if create_data_channel_error != nil {
		fmt.Fprintf(os.Stderr, config_t.Messages.WEBRTC.CREATE_DATA_CHANNEL.ERROR,
			create_data_channel_error)
		os.Exit(-1)
	} else {
		fmt.Print(config_t.Messages.WEBRTC.CREATE_DATA_CHANNEL.SUCCESS)
	}
	dc.OnOpen(func() {
		fmt.Print(config_t.Messages.WEBRTC.ON_DATA_CHANNEL)
		httpHandler(&pending_mutex, pending, dc)
	})
	dc.OnMessage(func(msg webrtc.DataChannelMessage) {
		fmt.Print(config_t.Messages.WEBRTC.ON_MESSAGE)
		msgHandler(&pending_mutex, pending, msg.Data)
	})

	offer_options := webrtc.OfferOptions{
		OfferAnswerOptions: webrtc.OfferAnswerOptions{
			VoiceActivityDetection: false,
		},
		ICERestart: false,
	}

	// Creating and sending offer to websocket
	var create_offer_error error
	*offer, create_offer_error = (*p2p).CreateOffer(&offer_options)
	if create_offer_error != nil {
		fmt.Fprintf(os.Stderr, config_t.Messages.WEBRTC.OFFER.ERROR, create_offer_error)
		os.Exit(-1)
	} else {
		fmt.Print(config_t.Messages.WEBRTC.OFFER.SUCCESS)
	}

	offer_t := Offer{
		Type:  "offer",
		Id:    id_t.Id,
		Offer: *offer,
	}

	offer_msg, message_marshal_error := json.Marshal(offer_t)
	if message_marshal_error != nil {
		fmt.Fprintf(os.Stderr, config_t.Messages.JSON.MARSHAL.ERROR, message_marshal_error)
		os.Exit(-1)
	} else {
		fmt.Print(config_t.Messages.JSON.MARSHAL.SUCCESS)
	}
	ws_write_error := conn.WriteMessage(websocket.TextMessage, offer_msg)
	if ws_write_error != nil {
		fmt.Fprintf(os.Stderr, config_t.Messages.WS.WRITE.ERROR, ws_write_error)
	} else {
		fmt.Print(config_t.Messages.WS.WRITE.SUCCESS)
	}
}

/*
 * This function gets:
 *   id: pointer to uuid
 *   p2p: pointer to the pointer of the peer connection
 *   offer: pointer to a session description
 *   raw: slice of bytes
 * Sets both local and remote sessions descriptions.
 * The function does not return anything.
 */
func answer(id *uuid.UUID, p2p **webrtc.PeerConnection, offer *webrtc.SessionDescription,
	raw []byte) {
	// Decoding JSON with answer in it
	var answer_t Answer
	message_unmarshal_error := json.Unmarshal(raw, &answer_t)
	if message_unmarshal_error != nil {
		fmt.Fprintf(os.Stderr, config_t.Messages.JSON.UNMARSHAL.ERROR, message_unmarshal_error)
		os.Exit(-1)
	} else {
		fmt.Print(config_t.Messages.JSON.UNMARSHAL.SUCCESS)
	}

	/*
	 * Checking if id matches and setting the session descriptions
	 * Once the local description is set, ICE gathering starts.
	 */
	if *id == answer_t.Id {
		set_local_description_error := (*p2p).SetLocalDescription(*offer)
		if set_local_description_error != nil {
			fmt.Fprintf(os.Stderr, config_t.Messages.WEBRTC.LOCAL_DESCRIPTION.ERROR,
				set_local_description_error)
			os.Exit(-1)
		} else {
			fmt.Print(config_t.Messages.WEBRTC.LOCAL_DESCRIPTION.SUCCESS)
		}
		set_remote_description_error := (*p2p).SetRemoteDescription(answer_t.Answer)
		if set_remote_description_error != nil {
			fmt.Fprintf(os.Stderr, config_t.Messages.WEBRTC.REMOTE_DESCRIPTION.ERROR,
				set_remote_description_error)
			os.Exit(-1)
		} else {
			fmt.Print(config_t.Messages.WEBRTC.REMOTE_DESCRIPTION.SUCCESS)
		}
	}
}

/*
 * This function gets:
 *   id: pointer to uuid
 *   p2p: pointer to the pointer of the peer connection
 *   raw: slice of bytes
 * Adds ICE candidate received from the signaling server.
 * The function does not return anything.
 */
func candidate(id *uuid.UUID, p2p **webrtc.PeerConnection, raw []byte) {
	// Decoding JSON with candidate in it
	var candidate_t Candidate
	message_unmarshal_error := json.Unmarshal(raw, &candidate_t)
	if message_unmarshal_error != nil {
		fmt.Fprintf(os.Stderr, config_t.Messages.JSON.UNMARSHAL.ERROR, message_unmarshal_error)
		os.Exit(-1)
	} else {
		fmt.Print(config_t.Messages.JSON.UNMARSHAL.SUCCESS)
	}

	// Checking if id matches and adding candidate
	if *id == candidate_t.Id {
		add_ice_error := (*p2p).AddICECandidate(candidate_t.Candidate)
		if add_ice_error != nil {
			fmt.Fprint(os.Stderr, config_t.Messages.WEBRTC.ADD_CANDIDATE.ERROR, add_ice_error)
		} else {
			fmt.Print(config_t.Messages.WEBRTC.ADD_CANDIDATE.SUCCESS)
		}
	}
}

func main() {
	// Attempting to read the config file
	raw_config, config_error := os.ReadFile("config.yml")
	if config_error != nil {
		fmt.Fprintf(os.Stderr, "Could not read config file.\n%v\n", config_error)
		os.Exit(-1)
	}

	// Loading config file to config struct
	config_error = yaml.Unmarshal(raw_config, &config_t)
	if config_error != nil {
		fmt.Fprintf(os.Stderr, "Could not load config file.\n%v\n", config_error)
		os.Exit(-1)
	}

	// Starting program
	fmt.Print(config_t.Messages.START)
	var (
		id    uuid.UUID
		p2p   *webrtc.PeerConnection
		offer webrtc.SessionDescription
	)

	if len(os.Args) != 2 {
		fmt.Fprint(os.Stderr, config_t.Messages.MISSING_PORT)
		os.Exit(-1)
	} else {
		http_port = os.Args[1]
	}

	// Attempting to connect to the signaling server via WebSocket
	conn, _, dial_error := websocket.DefaultDialer.Dial(config_t.Config.SIGNALING_SERVER, nil)
	if dial_error != nil {
		fmt.Fprintf(os.Stderr, config_t.Messages.WS.CONNECT.ERROR, dial_error)
		os.Exit(-1)
	} else {
		fmt.Print(config_t.Messages.WS.CONNECT.SUCCESS)
	}

	// Attempt to close WebSocket connection on quit
	defer func() {
		quit(id, p2p, conn)
	}()

	// Handling SIGINT and SIGTERM signals
	signal_channel := make(chan os.Signal, 1)
	signal.Notify(signal_channel, syscall.SIGINT, syscall.SIGTERM)
	go func() {
		<-signal_channel
		quit(id, p2p, conn)
	}()

	// Handling exit command
	go func() {
		scanner := bufio.NewScanner(os.Stdin)
		for scanner.Scan() {
			if scanner.Text() == "exit" {
				quit(id, p2p, conn)
			}
		}
	}()

	// Constructing the register message
	register_t := Register{
		Type:     "register",
		Username: config_t.Config.AUTHENTICATION.USERNAME,
		Password: config_t.Config.AUTHENTICATION.PASSWORD,
	}
	// Encoding register struct to JSON
	register_msg, register_marshal_error := json.Marshal(register_t)
	if register_marshal_error != nil {
		fmt.Fprintf(os.Stderr, config_t.Messages.JSON.MARSHAL.ERROR, register_marshal_error)
		return
	} else {
		fmt.Print(config_t.Messages.JSON.MARSHAL.SUCCESS)
	}

	// Registering to the signaling server as normal user
	ws_write_error := conn.WriteMessage(websocket.TextMessage, register_msg)
	if ws_write_error != nil {
		fmt.Fprintf(os.Stderr, config_t.Messages.WS.REGISTER.ERROR, ws_write_error)
		return
	} else {
		fmt.Print(config_t.Messages.WS.REGISTER.SUCCESS)
	}

	// WebSocket loop for incomming messages
	for {
		// Reading message
		_, ws_msg, ws_read_error := conn.ReadMessage()
		if ws_read_error != nil {
			fmt.Fprintf(os.Stderr, config_t.Messages.WS.READ.ERROR, ws_read_error)
			return
		} else {
			fmt.Print(config_t.Messages.WS.READ.SUCCESS)
		}

		// Decoding as any-format JSON
		var parsed map[string]any
		message_unmarshal_error := json.Unmarshal(ws_msg, &parsed)
		if message_unmarshal_error != nil {
			fmt.Fprintf(os.Stderr, config_t.Messages.JSON.UNMARSHAL.ERROR, message_unmarshal_error)
			continue
		} else {
			fmt.Print(config_t.Messages.JSON.UNMARSHAL.SUCCESS)
		}

		// Parsing message type
		msg_type, exists := parsed["type"].(string)
		if !exists {
			fmt.Fprint(os.Stderr, config_t.Messages.JSON.FORMAT.ERROR)
			continue
		} else {
			fmt.Printf(config_t.Messages.JSON.FORMAT.SUCCESS, msg_type)
		}

		// Setting handler function for each message type
		switch msg_type {
		case "id":
			fmt.Printf(config_t.Messages.JSON.TYPE.SUCCESS, msg_type)
			peerInit(&id, &p2p, &offer, conn, ws_msg)
		case "answer":
			fmt.Printf(config_t.Messages.JSON.TYPE.SUCCESS, msg_type)
			answer(&id, &p2p, &offer, ws_msg)
		case "candidate":
			fmt.Printf(config_t.Messages.JSON.TYPE.SUCCESS, msg_type)
			candidate(&id, &p2p, ws_msg)
		default:
			fmt.Fprintf(os.Stderr, config_t.Messages.JSON.TYPE.ERROR, msg_type)
			continue
		}
	}
}