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

import (
	"bytes"
	"encoding/gob"
	"encoding/json"
	"fmt"
	"io"
	"net/http"
	"os"

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

var config_t Config

/*
 * This function gets:
 *   dc: pointer to data channel
 *   raw: slice of bytes
 * Receives data from dc, decodes it, creates and sends a HTTP request,
 * encodes the HTTP response and sends it back to the dc.
 */
func channelHandler(dc *webrtc.DataChannel, raw []byte) {
	// Creating a buffer and store the received
	buf := bytes.NewBuffer(raw)
	// Creating a decoder and decode the data
	dec := gob.NewDecoder(buf)
	var request_t Request
	decode_error := dec.Decode(&request_t)
	if decode_error != nil {
		fmt.Fprintf(os.Stderr, config_t.Messages.GOB.DECODE.ERROR, decode_error)
		return
	} else {
		fmt.Print(config_t.Messages.GOB.DECODE.SUCCESS)
	}

	// Constructing the HTTP request
	req_body := bytes.NewBuffer(request_t.Body)
	req, req_error := http.NewRequest(request_t.Method, config_t.Config.PROXY_SERVER+request_t.Path, req_body)
	if req_error != nil {
		fmt.Fprintf(os.Stderr, config_t.Messages.HTTP.REQUEST.ERROR, req_error)
		return
	} else {
		fmt.Print(config_t.Messages.HTTP.REQUEST.SUCCESS)
	}
	req.Header = request_t.Headers

	// Creating a HTTP client - handles redirect requests as normal requests
	client := &http.Client{
		CheckRedirect: func(req *http.Request, via []*http.Request) error {
			return http.ErrUseLastResponse
		},
	}

	// Submitting the request
	res, send_req_error := client.Do(req)
	if send_req_error != nil {
		fmt.Fprintf(os.Stderr, config_t.Messages.HTTP.SEND_REQUEST.ERROR, send_req_error)
		return
	} else {
		fmt.Print(config_t.Messages.HTTP.SEND_REQUEST.SUCCESS)
	}

	// Constructing the response
	res_body, read_body_error := io.ReadAll(res.Body)
	if read_body_error != nil {
		fmt.Fprintf(os.Stderr, config_t.Messages.HTTP.READ.ERROR, read_body_error)
		return
	} else {
		fmt.Print(config_t.Messages.HTTP.READ.SUCCESS)
	}
	response_t := Response{
		ID:      request_t.ID,
		Status:  res.StatusCode,
		Headers: res.Header,
		Body:    res_body,
	}

	// Encoding and sending the response to the data channel
	var enc_buf bytes.Buffer
	enc := gob.NewEncoder(&enc_buf)
	encode_error := enc.Encode(response_t)
	if encode_error != nil {
		fmt.Fprintf(os.Stderr, config_t.Messages.GOB.ENCODE.ERROR, encode_error)
		return
	} else {
		fmt.Print(config_t.Messages.GOB.ENCODE.SUCCESS)
	}

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

/*
 * This function gets:
 *   conn: pointer to the websocket connection
 *   peer_connections: pointer to map of tracked peer connections
 *   raw: slice of bytes
 * Receives a peer connection offer, creates a new peer connection, sets the
 * listeners for it. Sets the remote descriptions, generates, sets and sends the answer.
 */
func offer(conn *websocket.Conn, peer_connections *map[uuid.UUID]*webrtc.PeerConnection, raw []byte) {
	// Constructing the offer
	var offer_t Offer
	message_unmarshal_error := json.Unmarshal(raw, &offer_t)
	if message_unmarshal_error != nil {
		fmt.Fprintf(os.Stderr, config_t.Messages.JSON.UNMARSHAL.ERROR, message_unmarshal_error)
		return
	} else {
		fmt.Print(config_t.Messages.JSON.UNMARSHAL.SUCCESS)
	}

	// 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)
		return
	} else {
		fmt.Print(config_t.Messages.WEBRTC.NEW.SUCCESS)
	}

	// Setting the peer listeners
	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:        offer_t.Id,
				Candidate: (*candidate).ToJSON(),
			}
			// Create JSON string for candidate and send it to the websocket
			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)
				return
			} 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)
				return
			} else {
				fmt.Print(config_t.Messages.WS.WRITE.SUCCESS)
			}
		}
	})
	p2p.OnDataChannel(func(dc *webrtc.DataChannel) {
		dc.OnOpen(func() {
			fmt.Print(config_t.Messages.WEBRTC.ON_DATA_CHANNEL)
		})
		dc.OnMessage(func(msg webrtc.DataChannelMessage) {
			fmt.Print(config_t.Messages.WEBRTC.ON_MESSAGE)
			channelHandler(dc, msg.Data)
		})
	})

	// Setting offer as remote description
	set_remote_description_error := p2p.SetRemoteDescription(offer_t.Offer)
	if set_remote_description_error != nil {
		fmt.Fprintf(os.Stderr, config_t.Messages.WEBRTC.REMOTE_DESCRIPTION.ERROR,
			set_remote_description_error)
		return
	} else {
		fmt.Print(config_t.Messages.WEBRTC.REMOTE_DESCRIPTION.SUCCESS)
	}

	// Generating answer
	answer_options := webrtc.AnswerOptions{
		OfferAnswerOptions: webrtc.OfferAnswerOptions{
			VoiceActivityDetection: false,
		},
	}
	answer, create_answer_error := p2p.CreateAnswer(&answer_options)
	if create_answer_error != nil {
		fmt.Fprintf(os.Stderr, config_t.Messages.WEBRTC.ANSWER.ERROR, create_answer_error)
		return
	} else {
		fmt.Print(config_t.Messages.WEBRTC.ANSWER.SUCCESS)
	}

	// Constructing answer, creating JSON from answer and sending it to the websocket
	answer_t := Answer{
		Type:   "answer",
		Id:     offer_t.Id,
		Answer: answer,
	}
	answer_msg, message_marshal_error := json.Marshal(answer_t)
	if message_marshal_error != nil {
		fmt.Fprintf(os.Stderr, config_t.Messages.JSON.MARSHAL.ERROR, message_marshal_error)
		return
	} else {
		fmt.Print(config_t.Messages.JSON.MARSHAL.SUCCESS)
	}
	ws_write_error := conn.WriteMessage(websocket.TextMessage, answer_msg)
	if ws_write_error != nil {
		fmt.Fprintf(os.Stderr, config_t.Messages.WS.WRITE.ERROR, ws_write_error)
		return
	} else {
		fmt.Print(config_t.Messages.WS.WRITE.SUCCESS)
	}
	// Setting answer as local description
	set_local_description_error := (*p2p).SetLocalDescription(answer)
	if set_local_description_error != nil {
		fmt.Fprintf(os.Stderr, config_t.Messages.WEBRTC.LOCAL_DESCRIPTION.ERROR,
			set_local_description_error)
		return
	} else {
		fmt.Print(config_t.Messages.WEBRTC.LOCAL_DESCRIPTION.SUCCESS)
	}

	// Storing the peer connection to the map
	(*peer_connections)[offer_t.Id] = p2p
}

/*
 * This function gets:
 *   peer_connections: pointer to map of tracked peer connections
 *   raw: slice of bytes
 * Adds ICE candidate received from the signaling server.
 * The function does not return anything.
 */
func candidate(peer_connections *map[uuid.UUID]*webrtc.PeerConnection, raw []byte) {
	// Constructing candidate from received data
	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)
		return
	} else {
		fmt.Print(config_t.Messages.JSON.UNMARSHAL.SUCCESS)
	}

	// Getting peer instance from map
	p2p := (*peer_connections)[candidate_t.Id]
	// Adding candidate to peer connection
	add_ice_error := p2p.AddICECandidate(candidate_t.Candidate)
	if add_ice_error != nil {
		fmt.Fprintf(os.Stderr, config_t.Messages.WEBRTC.ADD_CANDIDATE.ERROR, add_ice_error)
		return
	} else {
		fmt.Print(config_t.Messages.WEBRTC.ADD_CANDIDATE.SUCCESS)
	}
}

/*
 * This function gets:
 *   peer_connections: pointer to map of tracked peer connections
 *   raw: slice of bytes
 * Terminates a peer connection.
 */
func terminate(peer_connections *map[uuid.UUID]*webrtc.PeerConnection, raw []byte) {
	var terminate_t Terminate
	message_unmarshal_error := json.Unmarshal(raw, &terminate_t)
	if message_unmarshal_error != nil {
		fmt.Fprintf(os.Stderr, config_t.Messages.JSON.UNMARSHAL.ERROR, message_unmarshal_error)
		return
	} else {
		fmt.Print(config_t.Messages.JSON.UNMARSHAL.SUCCESS)
	}

	// Getting peer instance from map and closing the connection
	p2p := (*peer_connections)[terminate_t.Id]
	p2p_close_error := p2p.Close()
	if p2p_close_error != nil {
		fmt.Fprintf(os.Stderr, config_t.Messages.WEBRTC.CLOSE.ERROR, p2p_close_error)
		return
	} else {
		fmt.Print(config_t.Messages.WEBRTC.CLOSE.SUCCESS)
	}
	// Remove peer connection entry from map
	delete(*peer_connections, terminate_t.Id)
}

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)
	peer_connections := make(map[uuid.UUID]*webrtc.PeerConnection)

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

	defer func() {
		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)
		}
	}()

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

	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)
			os.Exit(-1)
		} 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 "offer":
			fmt.Printf(config_t.Messages.JSON.TYPE.SUCCESS, msg_type)
			offer(conn, &peer_connections, ws_msg)
		case "candidate":
			fmt.Printf(config_t.Messages.JSON.TYPE.SUCCESS, msg_type)
			candidate(&peer_connections, ws_msg)
		case "terminate":
			fmt.Printf(config_t.Messages.JSON.TYPE.SUCCESS, msg_type)
			terminate(&peer_connections, ws_msg)
		default:
			fmt.Fprintf(os.Stderr, config_t.Messages.JSON.TYPE.ERROR, msg_type)
			continue
		}
	}
}