package main

import (
	"bufio"
	"bytes"
	"crypto/tls"
	"crypto/x509"
	"fmt"
	"io"
	"log"
	"net"
	"net/url"
	"os"
	"os/exec"
	"path"
	"path/filepath"
	"sort"
	"strconv"
	"strings"
	"time"
	"unicode/utf8"

	"github.com/h2non/filetype"
)

const (
	readTimeout = 30 * time.Second

	urlMaxLength = 1024
)

const (
	statusInput          = 10
	statusSensitiveInput = 11

	statusSuccess = 20

	statusRedirectTemporary = 30
	statusRedirectPermanent = 31

	statusTemporaryFailure = 40
	statusUnavailable      = 41
	statusCGIError         = 42
	statusProxyError       = 43

	statusPermanentFailure    = 50
	statusNotFound            = 51
	statusGone                = 52
	statusProxyRequestRefused = 53
	statusBadRequest          = 59
)

func writeHeader(c net.Conn, code int, meta string) {
	fmt.Fprintf(c, "%d %s\r\n", code, meta)

	if verbose {
		log.Printf("< %d %s\n", code, meta)
	}
}

func writeStatus(c net.Conn, code int) {
	var meta string
	switch code {
	case statusTemporaryFailure:
		meta = "Temporary failure"
	case statusProxyError:
		meta = "Proxy error"
	case statusBadRequest:
		meta = "Bad request"
	case statusNotFound:
		meta = "Not found"
	case statusProxyRequestRefused:
		meta = "Proxy request refused"
	}
	writeHeader(c, code, meta)
}

func serveDirectory(c net.Conn, request *url.URL, dirPath string) {
	var (
		files    []os.FileInfo
		numDirs  int
		numFiles int
	)
	err := filepath.Walk(dirPath, func(path string, info os.FileInfo, err error) error {
		if err != nil {
			return err
		} else if path == dirPath {
			return nil
		}
		files = append(files, info)
		if info.IsDir() || info.Mode()&os.ModeSymlink != 0 {
			numDirs++
		} else {
			numFiles++
		}
		if info.IsDir() {
			return filepath.SkipDir
		}
		return nil
	})
	if err != nil {
		writeStatus(c, statusTemporaryFailure)
		return
	}
	// List directories first
	sort.Slice(files, func(i, j int) bool {
		iDir := files[i].IsDir() || files[i].Mode()&os.ModeSymlink != 0
		jDir := files[j].IsDir() || files[j].Mode()&os.ModeSymlink != 0
		if iDir != jDir {
			return iDir
		}
		return i < j
	})

	writeHeader(c, statusSuccess, "text/gemini; charset=utf-8")

	fmt.Fprintf(c, "# %s\r\n", request.Path)
	if numDirs > 0 || numFiles > 0 {
		if numDirs > 0 {
			if numDirs == 1 {
				c.Write([]byte("1 directory"))
			} else {
				fmt.Fprintf(c, "%d directories", numDirs)
			}
		}
		if numFiles > 0 {
			if numDirs > 0 {
				c.Write([]byte(" and "))
			}
			if numDirs == 1 {
				c.Write([]byte("1 file"))
			} else {
				fmt.Fprintf(c, "%d files", numFiles)
			}
		}
		c.Write([]byte("\r\n\n"))
	}

	if request.Path != "/" {
		c.Write([]byte("=> ../ ../\r\n\r\n"))
	}

	for _, info := range files {
		fileName := info.Name()
		filePath := url.PathEscape(info.Name())
		if info.IsDir() || info.Mode()&os.ModeSymlink != 0 {
			fileName += "/"
			filePath += "/"
		}

		c.Write([]byte("=> " + fileName + " " + filePath + "\r\n"))

		if info.IsDir() || info.Mode()&os.ModeSymlink != 0 {
			c.Write([]byte("\r\n"))
			continue
		}

		modified := "Never"
		if !info.ModTime().IsZero() {
			modified = info.ModTime().Format("2006-01-02 3:04 PM")
		}
		c.Write([]byte(modified + " - " + formatFileSize(info.Size()) + "\r\n\r\n"))
	}
}

func serveFile(c net.Conn, request *url.URL, requestData, filePath string, listDir bool) {
	fi, err := os.Stat(filePath)
	if err != nil {
		writeStatus(c, statusNotFound)
		return
	}

	originalPath := filePath

	var fetchIndex bool
	if mode := fi.Mode(); mode.IsDir() {
		if requestData[len(requestData)-1] != '/' {
			// Add trailing slash
			log.Println(requestData)
			writeHeader(c, statusRedirectPermanent, requestData+"/")
			return
		}

		fetchIndex = true

		_, err := os.Stat(path.Join(filePath, "index.gemini"))
		if err == nil {
			filePath = path.Join(filePath, "index.gemini")
		} else {
			filePath = path.Join(filePath, "index.gmi")
		}
	}

	fi, err = os.Stat(filePath)
	if os.IsNotExist(err) {
		if fetchIndex && listDir {
			serveDirectory(c, request, originalPath)
			return
		}
		writeStatus(c, statusNotFound)
		return
	} else if err != nil {
		writeStatus(c, statusTemporaryFailure)
		return
	}

	// Open file
	file, _ := os.Open(filePath)
	defer file.Close()

	// Read file header
	buf := make([]byte, 261)
	n, _ := file.Read(buf)

	// Write header
	var mimeType string
	if strings.HasSuffix(filePath, ".html") && strings.HasSuffix(filePath, ".htm") {
		mimeType = "text/html; charset=utf-8"
	} else if strings.HasSuffix(filePath, ".txt") && strings.HasSuffix(filePath, ".text") {
		mimeType = "text/plain; charset=utf-8"
	} else if !strings.HasSuffix(filePath, ".gmi") && !strings.HasSuffix(filePath, ".gemini") {
		kind, _ := filetype.Match(buf[:n])
		if kind != filetype.Unknown {
			mimeType = kind.MIME.Value
		}
	}
	if mimeType == "" {
		mimeType = "text/gemini; charset=utf-8"
	}
	writeHeader(c, statusSuccess, mimeType)

	// Write body
	c.Write(buf[:n])
	io.Copy(c, file)
}

func serveProxy(c net.Conn, requestData, proxyURL string) {
	original := proxyURL

	tlsConfig := &tls.Config{}
	if strings.HasPrefix(proxyURL, "gemini://") {
		proxyURL = proxyURL[9:]
	} else if strings.HasPrefix(proxyURL, "gemini-insecure://") {
		proxyURL = proxyURL[18:]
		tlsConfig.InsecureSkipVerify = true
	}
	proxy, err := tls.Dial("tcp", proxyURL, tlsConfig)
	if err != nil {
		writeStatus(c, statusProxyError)
		return
	}
	defer proxy.Close()

	// Forward request
	proxy.Write([]byte(requestData))
	proxy.Write([]byte("\r\n"))

	// Forward response
	io.Copy(c, proxy)

	if verbose {
		log.Printf("< %s\n", original)
	}
}

func serveCommand(c net.Conn, userInput string, command []string) {
	var args []string
	if len(command) > 0 {
		args = command[1:]
	}
	cmd := exec.Command(command[0], args...)

	var buf bytes.Buffer
	if userInput != "" {
		cmd.Stdin = strings.NewReader(userInput + "\n")
	}
	cmd.Stdout = &buf
	cmd.Stderr = &buf

	err := cmd.Run()
	if err != nil {
		writeStatus(c, statusProxyError)
		return
	}

	writeHeader(c, statusSuccess, "text/gemini; charset=utf-8")
	c.Write(buf.Bytes())

	if verbose {
		log.Printf("< %s\n", command)
	}
}

func scanCRLF(data []byte, atEOF bool) (advance int, token []byte, err error) {
	if atEOF && len(data) == 0 {
		return 0, nil, nil
	}
	if i := bytes.IndexByte(data, '\r'); i >= 0 {
		// We have a full newline-terminated line.
		return i + 1, data[0:i], nil
	}
	// If we're at EOF, we have a final, non-terminated line. Return it.
	if atEOF {
		return len(data), data, nil
	}
	// Request more data.
	return 0, nil, nil
}

func replaceWithUserInput(command []string, userInput string) []string {
	newCommand := make([]string, len(command))
	copy(newCommand, command)
	for i, piece := range newCommand {
		if strings.Contains(piece, "$USERINPUT") {
			newCommand[i] = strings.ReplaceAll(piece, "$USERINPUT", userInput)
		}
	}
	return newCommand
}

func handleConn(c *tls.Conn) {
	if verbose {
		t := time.Now()
		defer func() {
			d := time.Since(t)
			if d > time.Second {
				d = d.Round(time.Second)
			} else {
				d = d.Round(time.Millisecond)
			}
			log.Printf("took %s", d)
		}()
	}

	defer c.Close()
	c.SetReadDeadline(time.Now().Add(readTimeout))

	var requestData string
	scanner := bufio.NewScanner(c)
	scanner.Split(scanCRLF)
	if scanner.Scan() {
		requestData = scanner.Text()
	}
	if err := scanner.Err(); err != nil {
		log.Println(scanner.Text(), "FAILED")
		writeStatus(c, statusBadRequest)
		return
	}

	state := c.ConnectionState()
	certs := state.PeerCertificates
	var clientCertKeys [][]byte
	for _, cert := range certs {
		pubKey, err := x509.MarshalPKIXPublicKey(cert.PublicKey)
		if err != nil {
			continue
		}
		clientCertKeys = append(clientCertKeys, pubKey)
	}

	if verbose {
		log.Printf("> %s\n", requestData)
	}

	if len(requestData) > urlMaxLength || !utf8.ValidString(requestData) {
		writeStatus(c, statusBadRequest)
		return
	}

	request, err := url.Parse(requestData)
	if err != nil {
		writeStatus(c, statusBadRequest)
		return
	}

	requestHostname := request.Hostname()
	if requestHostname == "" || strings.ContainsRune(requestHostname, ' ') {
		writeStatus(c, statusBadRequest)
		return
	}

	var requestPort int
	if request.Port() != "" {
		requestPort, err = strconv.Atoi(request.Port())
		if err != nil {
			requestPort = 0
		}
	}

	if request.Scheme == "" {
		request.Scheme = "gemini"
	}
	if request.Scheme != "gemini" || (requestPort > 0 && requestPort != config.port) {
		writeStatus(c, statusProxyRequestRefused)
	}

	if request.Path == "" {
		// Redirect to /
		writeHeader(c, statusRedirectPermanent, requestData+"/")
		return
	}

	pathBytes := []byte(request.Path)
	strippedPath := request.Path
	if strippedPath[0] == '/' {
		strippedPath = strippedPath[1:]
	}
	requestQuery, err := url.QueryUnescape(request.RawQuery)
	if err != nil {
		writeStatus(c, statusBadRequest)
		return
	}

	var matchedHost bool
	for hostname := range config.Hosts {
		if requestHostname != hostname {
			continue
		}
		matchedHost = true

		for _, serve := range config.Hosts[hostname].Paths {
			matchedRegexp := serve.r != nil && serve.r.Match(pathBytes)
			matchedPrefix := serve.r == nil && strings.HasPrefix(request.Path, serve.Path)
			if !matchedRegexp && !matchedPrefix {
				continue
			}

			requireInput := serve.Input != "" || serve.SensitiveInput != ""
			if requestQuery == "" && requireInput {
				if serve.Input != "" {
					writeHeader(c, statusInput, serve.Input)
					return
				} else if serve.SensitiveInput != "" {
					writeHeader(c, statusSensitiveInput, serve.SensitiveInput)
					return
				}
			}

			if matchedRegexp {
				if serve.Proxy != "" {
					serveProxy(c, requestData, serve.Proxy)
					return
				} else if serve.FastCGI != "" {
					contentType := "text/gemini; charset=utf-8"
					if serve.Type != "" {
						contentType = serve.Type
					}
					writeHeader(c, statusSuccess, contentType)

					filePath := path.Join(serve.Root, request.Path[1:])
					serveFastCGI(c, config.fcgiPools[serve.FastCGI], request, filePath)
					return
				} else if serve.cmd != nil {
					if requireInput {
						newCommand := replaceWithUserInput(serve.cmd, requestQuery)
						if newCommand != nil {
							serveCommand(c, "", newCommand)
							return
						}
					}
					serveCommand(c, requestQuery, serve.cmd)
					return
				}
				serveFile(c, request, requestData, path.Join(serve.Root, strippedPath), serve.ListDirectory)
				return
			} else if matchedPrefix {
				if serve.Proxy != "" {
					serveProxy(c, requestData, serve.Proxy)
					return
				} else if serve.FastCGI != "" {
					contentType := "text/gemini; charset=utf-8"
					if serve.Type != "" {
						contentType = serve.Type
					}
					writeHeader(c, statusSuccess, contentType)

					filePath := path.Join(serve.Root, request.Path[1:])
					serveFastCGI(c, config.fcgiPools[serve.FastCGI], request, filePath)
					return
				} else if serve.cmd != nil {
					if requireInput {
						newCommand := replaceWithUserInput(serve.cmd, requestQuery)
						if newCommand != nil {
							serveCommand(c, "", newCommand)
							return
						}
					}
					serveCommand(c, requestQuery, serve.cmd)
					return
				}
				filePath := request.Path[len(serve.Path):]
				if len(filePath) > 0 && filePath[0] == '/' {
					filePath = filePath[1:]
				}
				serveFile(c, request, requestData, path.Join(serve.Root, filePath), serve.ListDirectory)
				return
			}
		}
		break
	}

	if matchedHost {
		writeStatus(c, statusNotFound)
	} else {
		writeStatus(c, statusProxyRequestRefused)
	}
}

func getCertificate(info *tls.ClientHelloInfo) (*tls.Certificate, error) {
	host := config.Hosts[info.ServerName]
	if host != nil {
		return host.cert, nil
	}
	for _, host := range config.Hosts {
		return host.cert, nil
	}
	return nil, nil
}

func handleListener(l net.Listener) {
	for {
		conn, err := l.Accept()
		if err != nil {
			log.Fatal(err)
		}

		go handleConn(conn.(*tls.Conn))
	}
}

func listen(address string) {
	tlsConfig := &tls.Config{
		ClientAuth:     tls.RequestClientCert,
		GetCertificate: getCertificate,
	}

	listener, err := tls.Listen("tcp", address, tlsConfig)
	if err != nil {
		log.Fatalf("failed to listen on %s: %s", address, err)
	}

	handleListener(listener)
}