package main

import (
    "flag"
    "io"
    "log"
    "net/http"
    "sync"

    "github.com/gorilla/mux"
)

type Client struct {
    w         http.ResponseWriter
    r         *http.Request
    connected bool
    closed    bool
    chbuf     chan []byte
}

func newClient(w http.ResponseWriter, r *http.Request) *Client {
    chbuf := make(chan []byte, 0)
    c := Client{w, r, true, false, chbuf}

    return &c
}

type Channel struct {
    value []*Client
    mu    sync.Mutex
}

func newChannel() *Channel {
    c := &Channel{}

    return c
}

func main() {
    key := flag.String("key", "", "Global stream key")
    flag.Parse()
    if *key == "" {
        panic("Supply a stream key")
    }

    rtr := mux.NewRouter()
    channelsmu := sync.Mutex{}
    channels := map[string]*Channel{}

    rtr.HandleFunc("/ingest/{name}", func(w http.ResponseWriter, r *http.Request) {
        vars := mux.Vars(r)
        log.Printf("[%s] Producer connected to channel '%s'.\n", remoteAddr(r), vars["name"])

        keys, ok := r.URL.Query()["key"]
        if !ok || len(keys[0]) < 1 {
            log.Printf("[%s] No stream key supplied.", remoteAddr(r))
            w.WriteHeader(http.StatusForbidden)
            return
        }
        suppliedkey := keys[0]
        if *key != suppliedkey {
            log.Printf("[%s] Invalid stream key supplied.", remoteAddr(r))
            w.WriteHeader(http.StatusForbidden)
            return
        }

        channelsmu.Lock()
        channels[vars["name"]] = newChannel()
        channel := channels[vars["name"]]
        channelsmu.Unlock()
        defer (func() {
            log.Printf("Removing channel '%s'", vars["name"])
            channelsmu.Lock()
            delete(channels, vars["name"])
            channelsmu.Unlock()
            channel.mu.Lock()
            for _, v := range channel.value {
                close(v.chbuf)
            }
            channel.mu.Unlock()
        })()

        if r.Body == nil {
            log.Printf("[%s] No body from connected producer.\n", remoteAddr(r))
            return
        }

        buf := make([]byte, 4096)

        for {
            n, err := r.Body.Read(buf)

            tmp := make([]byte, n)
            copy(tmp, buf[:n])

            if err == io.EOF {
                log.Printf("[%s] Producer EOF'd dropping.\n", remoteAddr(r))
                break
            }
            if err != nil && err != io.EOF {
                log.Printf("[%s] Producer error: %v.\n", remoteAddr(r), err)
                break
            }
            channel.mu.Lock()
            for _, v := range channel.value {
                if v.connected {
                    v.chbuf <- tmp
                } else if !v.connected && !v.closed {
                    close(v.chbuf)
                    v.closed = true
                }

            }
            channel.mu.Unlock()
        }

    })

    rtr.HandleFunc("/{name}", func(w http.ResponseWriter, r *http.Request) {
        vars := mux.Vars(r)
        log.Printf("[%s] Consumer connected to channel '%s'.\n", remoteAddr(r), vars["name"])
        channelsmu.Lock()
        channel := channels[vars["name"]]
        channelsmu.Unlock()
        if channel == nil {
            log.Printf("Channel '%s' does not exist.\n", vars["name"])
            w.WriteHeader(http.StatusNotFound)
            return
        }

        cl := newClient(w, r)

        channel.mu.Lock()
        channel.value = append(channel.value, cl)
        channel.mu.Unlock()
        for dat := range cl.chbuf {
            if !cl.connected {
                continue
            }
            _, err := w.Write(dat)
            if err != nil {
                log.Printf("[%s] Couldn't write to consumer. Error: %s. Flagging connection.\n", remoteAddr(r), err)
                cl.connected = false
                continue
            }
            if f, ok := w.(http.Flusher); ok {
                f.Flush()
            }
        }

        channel.mu.Lock()
        newClients := make([]*Client, 0)
        for _, v := range channel.value {
            if v != cl {
                newClients = append(newClients, v)
            } else {
                log.Printf("[%s] Dropped consumer.\n", remoteAddr(v.r))
            }
        }
        channel.value = newClients
        channel.mu.Unlock()
    })

    srv := &http.Server{
        Handler: rtr,
        Addr:    ":4242",
    }

    srv.ListenAndServe()
}

func remoteAddr(r *http.Request) string {
    ip := r.Header.Get("x-forwarded-for")
    if ip == "" {
        ip = r.RemoteAddr
    }
    return ip
}