package main

import (
	"context"
	"encoding/json"
	"errors"
	"fmt"
	"io/fs"
	"log"
	"net"
	"net/http"
	"os"
	"os/signal"
	"strconv"
	"strings"
	"sync"

	"github.com/oschwald/geoip2-golang"
)

var (
	asnDB  *geoip2.Reader
	cityDB *geoip2.Reader
	config serverConfig
	cache  sync.Map
)

const (
	devMode  = "dev"
	prodMode = "prod"
)

type serverConfig struct {
	ASNDB            string
	CityDB           string
	Listen           string
	Mode             string
	SocketPermission fs.FileMode
}

type body struct {
	IP           net.IP `json:"ip"`
	Country      string `json:"country"`
	Region       string `json:"region"`
	City         string `json:"city"`
	ASN          uint   `json:"asn"`
	Organization string `json:"organization"`
}

type bodyError struct {
	Error string `json:"error"`
}

func main() {
	perm, err := strconv.ParseUint(getEnvOr("IP_CHECKER_SOCKET_PERMISSION", "0600"), 8, 32)
	if err != nil {
		log.Fatal(err)
	}

	config = serverConfig{
		ASNDB:            os.Getenv("IP_CHECKER_ASN_DB"),
		CityDB:           os.Getenv("IP_CHECKER_CITY_DB"),
		Listen:           getEnvOr("IP_CHECKER_LISTEN", ":8080"),
		SocketPermission: fs.FileMode(perm),
		Mode:             os.Getenv("IP_CHECKER_MODE"),
	}

	if config.CityDB != "" {
		var err error
		cityDB, err = geoip2.Open(config.CityDB)
		if err != nil {
			log.Fatalf("Error opening GeoIP City database: %v", err)
		}
		defer cityDB.Close()
	} else {
		log.Fatalln("'IP_CHECKER_CITY_DB' not set")
	}
	if config.ASNDB != "" {
		var err error
		asnDB, err = geoip2.Open(config.ASNDB)
		if err != nil {
			log.Fatalf("Error opening GeoIP ASN database: %v", err)
		}
		defer asnDB.Close()
	} else {
		log.Fatalln("'IP_CHECKER_ASN_DB' not set")
	}

	http.HandleFunc("/api/v1/ip", handleRequest)

	log.Printf("Starting server on %s", config.Listen)

	var srv http.Server

	// graceful shutdown
	go func() {
		sigint := make(chan os.Signal, 1)
		signal.Notify(sigint, os.Interrupt)
		<-sigint

		srv.Shutdown(context.Background())
	}()

	if strings.HasPrefix(config.Listen, "unix/") {
		socketPath := strings.TrimPrefix(config.Listen, "unix/")
		unixListener, err := net.Listen("unix", socketPath)
		if err != nil {
			log.Fatal(err)
		}

		if err := os.Chmod(socketPath, config.SocketPermission); err != nil {
			log.Fatal(err)
		}

		if err := srv.Serve(unixListener); err != http.ErrServerClosed {
			log.Fatal(err)
		}
	} else {
		srv.Addr = config.Listen
		if err := srv.ListenAndServe(); err != http.ErrServerClosed {
			log.Fatal(err)
		}
	}
}

func generateJSON(ip string) (body, error) {
	var data body

	parsedIP := net.ParseIP(ip)

	// check for reserved ip
	if parsedIP == nil {
		return body{}, errors.New("invalid IP")
	} else if parsedIP.IsPrivate() {
		return body{}, errors.New("private IP")
	} else if parsedIP.IsLoopback() {
		return body{}, errors.New("loopback IP")
	} else if parsedIP.IsUnspecified() {
		return body{}, errors.New("unspecified IP")
	}

	// cache
	if cached, found := cache.Load(ip); found && config.Mode == prodMode {
		return cached.(body), nil
	}

	// query
	cityRecord, err := cityDB.City(parsedIP)
	if err != nil {
		log.Printf("GeoIP City lookup failed for IP: %s, error: %v", ip, err)
		return body{}, errors.New("City record lookup failed")
	}

	ASNRecord, err := asnDB.ASN(parsedIP)
	if err != nil {
		log.Printf("GeoIP ASN lookup failed for IP: %s, error: %v", ip, err)
		return body{}, errors.New("ASN record lookup failed")
	}

	data.IP = parsedIP
	data.Country = cityRecord.Country.Names["en"]
	data.Region = cityRecord.Subdivisions[0].Names["en"]
	data.City = cityRecord.City.Names["en"]
	data.ASN = ASNRecord.AutonomousSystemNumber
	data.Organization = ASNRecord.AutonomousSystemOrganization

	if config.Mode == prodMode {
		cache.Store(ip, data)
	}
	return data, nil
}

func handleRequest(w http.ResponseWriter, r *http.Request) {
	w.Header().Add("Access-Control-Allow-Origin", "*")
	w.Header().Add("Content-Type", "application/json")

	sourceIP := func() string {
		if q := r.URL.Query().Get("q"); q != "" {
			return q
		}

		if config.Mode == "dev" {
			return "223.5.5.5"
		}

		if ip := r.Header.Get("X-Forwarded-For"); ip != "" {
			return strings.Split(ip, ",")[0]
		}

		host, _, err := net.SplitHostPort(r.RemoteAddr)
		if err != nil {
			errorText := fmt.Sprintf("Unable to parse remote address: %v", err)
			log.Printf(errorText)
			return errorText
		}

		return host
	}()

	bodyData, err := generateJSON(sourceIP)
	if err != nil {
		errorBody := bodyError{
			Error: err.Error(),
		}
		errorData, _ := json.MarshalIndent(errorBody, "", "  ")
		http.Error(w, string(errorData)+"\n", http.StatusInternalServerError)
		return
	}
	bodyJson, _ := json.MarshalIndent(bodyData, "", "  ")

	fmt.Fprintf(w, "%s\n", bodyJson)
}

func getEnvOr(key string, defaultValue string) string {
	if value, exists := os.LookupEnv(key); exists {
		return value
	}
	return defaultValue
}