package main

import (
    "bufio"
    "crypto/md5"
    "encoding/json"
    "flag"
    "fmt"
    "html/template"
    "io"
    "io/ioutil"
    "log"
    "mime"
    "net/http"
    "os"
    "path"
    "strconv"
    "strings"
    "sync"

    "github.com/jmoiron/sqlx"
    "github.com/lib/pq"
)

type Config struct {
    Listen   string `json:"listen"`
    DBStr    string `json:"db"`
    TmplPath string `json:"tmpl"`
    Root     string `json:"root"`

    DataPath string `json:"data"`
    MaxSize  uint `json:"max_upload"`
    MaxPages uint `json:"max_pages"`
}

type cmdFunction func (w http.ResponseWriter, r *http.Request, args []string)

type Server struct {
    mux  *http.ServeMux
    conf *Config
    db   *sqlx.DB

    fslock sync.RWMutex
    filestore map[string][]string

    cmds map[string]cmdFunction
}

type note struct {
    Hash     string
    Name     string
    Location string
}

const scheme = `
    CREATE TABLE IF NOT EXISTS notes (
        hash     VARCHAR(4) PRIMARY KEY,
        name     TEXT NOT NULL,
        location TEXT NOT NULL,
        available BOOL DEFAULT True
    )
`

var validMIME = map[string]bool {
    "text/plain": true,
    "image/jpeg": true,
    "image/png": true,
    "application/pdf": true,
}

func main() {
    confFlag := flag.String("c", "config.json", "config file")
    flag.Parse()

    f, err := os.Open(*confFlag)
    if err != nil {
        log.Fatal(err)
    }

    var conf Config
    err = json.NewDecoder(f).Decode(&conf)
    if err != nil {
        log.Fatal(err)
    }

    s := NewServer(&conf)
    if err != nil {
        log.Fatal(err)
    }

    err = s.Connect()
    if err != nil {
        log.Fatal(err)
    }

    log.Fatal(s.Run())
}

func NewServer(conf *Config) *Server {
    s := &Server{}
    s.mux = http.NewServeMux()
    s.mux.HandleFunc("/", s.httpRoot)
    s.mux.HandleFunc("/upload", s.httpUpload)
    s.mux.Handle("/data/", http.StripPrefix("/data/",
        http.FileServer(http.Dir(conf.DataPath))))

    s.conf = conf

    s.filestore = map[string][]string{}

    s.cmds = map[string]cmdFunction{}
    s.addCmd("a", s.cmdCreateNode)
    s.addCmd("c", s.cmdChangeLoc)

    return s
}

func (s *Server) Connect() error {
    files, err := ioutil.ReadDir(s.conf.DataPath)
    if err != nil {
        return err
    }

    for _, file := range files {
        name := file.Name()

        hash := name[0:4]
        s.saveFile(hash, name)
    }

    s.db, err = sqlx.Connect("postgres", s.conf.DBStr)
    if err != nil {
        return err
    }

    s.db.MustExec(scheme)
    return nil
}

func (s *Server) Run() error {
    log.Printf("Listening on %s", s.conf.Listen)
    return http.ListenAndServe(s.conf.Listen, s.mux)
}

func (s *Server) renderTemplate(w http.ResponseWriter, data interface{}, pathname string) error {

    tmpl, err := template.ParseFiles(path.Join(s.conf.TmplPath, pathname))
    if err != nil {
        log.Print(err)
        return err
    }

    err = tmpl.Execute(w, data)
    if err != nil {
        log.Print(err)
    }
    return err
}

func (s *Server) httpLog(r *http.Request, format string, args... interface{}) {
    args = append([]interface{}{r.URL.Path, r.RemoteAddr}, args...)
    log.Printf("(url: %s, ip: %s) " + format, args...)
}

func (s *Server) Error(w http.ResponseWriter, r *http.Request, err string, code int) {
    // Do not leak internal errors
    if code < 500 {
        http.Error(w, err, code)
    } else {
        w.WriteHeader(code)
    }

    s.httpLog(r, "ERR: %s", err)
}

func (s *Server) cmdCreateNode(w http.ResponseWriter, r *http.Request, args []string) {
    if len(args) < 3 {
        s.Error(w, r, "Please specify both name and location", http.StatusBadRequest)
        return
    }
    var (
        name = args[1]
        location = args[2]
    )

    hash := fmt.Sprintf("%x", md5.Sum([]byte(name)))[0:4]

    _, err := s.db.ExecContext(r.Context(), `
        INSERT INTO notes (hash, name, location)
        VALUES ($1, $2, $3)`, hash, name, location)
    if err != nil {
        s.Error(w, r, err.Error(), http.StatusInternalServerError)
        return
    }

    s.httpLog(r, "Created node %s", hash)

    var ref strings.Builder
    fmt.Fprintf(&ref, "#%s", hash)
    if location != "" {
        fmt.Fprintf(&ref, "/%s", location)
    }

    fmt.Fprintf(w, `
    <html>
        <p><b>Note:</b> %s</p>
        <p><b>Location:</b> %s</p>
        <p><b>Hash:</b> %s"</p>
        <p><b>Ref:</b> %s</p>
    </html>`, name, location, hash, ref.String())
}

func (s *Server) cmdChangeLoc(w http.ResponseWriter, r *http.Request, args []string) {
    argc := len(args)
    if argc < 3 {
        s.Error(w, r, "Please specify both id and new location", http.StatusBadRequest)
        return
    }

    var (
        hashes = args[1:argc-1]
        loc = args[argc-1]
    )

    _, err := s.db.ExecContext(r.Context(), `
        UPDATE notes SET location = $1 WHERE hash = ANY($2)`,
        loc, pq.Array(hashes))
    if err != nil {
        s.Error(w, r, err.Error(), http.StatusInternalServerError)
        return
    }

    http.Redirect(w, r, "/", http.StatusSeeOther)
}

func (s *Server) addCmd(name string, f cmdFunction) {
    s.cmds[name] = f
}

func (s *Server) httpRunCmd(w http.ResponseWriter, r *http.Request) {
    cmd := r.FormValue("cmd")

    scanner := bufio.NewScanner(strings.NewReader(cmd))
    scanner.Split(bufio.ScanWords)
    var words []string
    for scanner.Scan() {
        words = append(words, scanner.Text())
    }
    if err := scanner.Err(); err != nil {
        s.Error(w, r, err.Error(), http.StatusInternalServerError)
        return
    }

    cmd = strings.ToLower(words[0])
    f, ok := s.cmds[cmd]
    if !ok {
        s.Error(w, r, fmt.Sprintf("No such command '%s'", cmd), http.StatusBadRequest)
        return
    }
    f(w, r, words)
}

func (s *Server) httpRoot(w http.ResponseWriter, r *http.Request) {
    if r.URL.Path != "/" {
        w.WriteHeader(http.StatusNotFound)
        return
    }

    if r.Method == "POST" {
        s.httpRunCmd(w, r)
        return
    }

    var page struct {
        Notes []note
        Msg   string
        Files *map[string][]string
        Root  string
    }
    page.Files = &s.filestore
    page.Root = s.conf.Root

    err := s.db.SelectContext(r.Context(), &page.Notes, `
        SELECT hash, name, location FROM notes WHERE available = True ORDER BY name`)
    if err != nil {
        s.Error(w, r, err.Error(), http.StatusInternalServerError)
        return
    }

    s.fslock.RLock()
    s.renderTemplate(w, &page, "root.template")
    s.fslock.RUnlock()
}

func (s *Server) httpUpload(w http.ResponseWriter, r *http.Request) {
    if r.Method != "POST" {
        w.WriteHeader(http.StatusMethodNotAllowed)
        return
    }

    r.Body = http.MaxBytesReader(w, r.Body, int64(s.conf.MaxSize) * 1000000)

    // Head hash
    hash := r.FormValue("h")
    _, err := strconv.ParseUint(hash, 16, 16)
    if err != nil || len(hash) != 4 {
        s.Error(w, r, "Invalid hash", http.StatusBadRequest)
        return
    }

    // Load file
    inf, header, err := r.FormFile("f")
    if err != nil {
        s.Error(w, r, err.Error(), http.StatusInternalServerError)
        return
    }
    defer inf.Close()

    // Check if valid type
    mtype := header.Header.Get("Content-Type")
    ok := validMIME[mtype]
    if !ok {
        s.Error(w, r, "Invalid file type", http.StatusBadRequest)
        return
    }

    // Create file
    fname, err := s.allocFile(hash, mtype)
    if err != nil {
        s.Error(w, r, err.Error(), http.StatusBadRequest)
        return
    }

    f, err := os.Create(path.Join(s.conf.DataPath, fname))
    if err != nil {
        s.Error(w, r, err.Error(), http.StatusInternalServerError)
        return
    }
    defer f.Close()

    _, err = io.Copy(f, inf)
    if err != nil {
        s.Error(w, r, err.Error(), http.StatusInternalServerError)
        return
    }

    s.saveFile(hash, fname)
    s.httpLog(r, "Uploaded file %s", fname)

    http.Redirect(w, r, s.conf.Root + "/", http.StatusSeeOther)
}

func (s *Server) allocFile(hash string, t string) (string, error) {
    s.fslock.RLock()
    existing := s.filestore[hash]
    s.fslock.RUnlock()

    if len(existing) >= int(s.conf.MaxPages) {
        return "", fmt.Errorf("No more than %d pages is allowed", s.conf.MaxPages)
    }

    var ext string
    extarr, _ := mime.ExtensionsByType(t)
    if len(extarr) > 0 {
        ext = extarr[0]
    }

    return fmt.Sprintf("%s.%d%s", hash, len(existing), ext), nil
}

func (s *Server) saveFile(hash string, fname string) {
    s.fslock.Lock()
    s.filestore[hash] = append(s.filestore[hash], fname)

    s.fslock.Unlock()
}