package main

import (
	"context"
	"encoding/json"
	"fmt"
	"log"
	"net"
	"net/http"
)

type AddressType string

const (
	IPv4 AddressType = "IPv4"
	IPv6 AddressType = "IPv6"
)

type NetAddress struct {
	Address string      `json:"address"`
	Type    AddressType `json:"type"`
}

type contextKey string

const netAddrKey = contextKey("netAddr")

func RealIpMiddleware(next http.Handler) http.Handler {
	return http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
		netAddr := extractNetAddress(r)
		if netAddr.Address == "" {
			http.Error(w, "Internal Server Error", http.StatusInternalServerError)
			return
		}
		ctx := context.WithValue(r.Context(), netAddrKey, netAddr)
		next.ServeHTTP(w, r.WithContext(ctx))
	})
}

func extractNetAddress(r *http.Request) NetAddress {
	realIpHeader := r.Header.Get("X-Real-Ip")
	if realIpHeader != "" {
		if ip := net.ParseIP(realIpHeader); ip != nil {
			return NetAddress{
				Address: realIpHeader,
				Type:    determineIPType(ip),
			}
		}
		log.Printf("Invalid X-Real-Ip header: %s", realIpHeader)
	}

	host, _, err := net.SplitHostPort(r.RemoteAddr)
	if err != nil {
		log.Printf("Failed to split host and port: %v", err)
		return NetAddress{}
	}

	if ip := net.ParseIP(host); ip != nil {
		return NetAddress{
			Address: host,
			Type:    determineIPType(ip),
		}
	}
	return NetAddress{}
}

func determineIPType(ip net.IP) AddressType {
	if ip.To4() != nil {
		return IPv4
	}
	return IPv6
}

func mainHandler(w http.ResponseWriter, r *http.Request) {
	netAddr, ok := r.Context().Value(netAddrKey).(NetAddress)
	if !ok {
		http.Error(w, "Internal Server Error", http.StatusInternalServerError)
		return
	}

	switch r.Method {
	case http.MethodGet:
		switch r.URL.Path {
		case "/":
			fmt.Fprintf(w, "%s\n%s", netAddr.Address, netAddr.Type)
		case "/api":
			responseJSON(w, netAddr)
		default:
			http.NotFound(w, r)
		}
	default:
		http.Error(w, "405 method not allowed", http.StatusMethodNotAllowed)
	}
}

func responseJSON(w http.ResponseWriter, data interface{}) {
	w.Header().Set("Content-Type", "application/json")
	err := json.NewEncoder(w).Encode(data)
	if err != nil {
		http.Error(w, "Failed to encode JSON", http.StatusInternalServerError)
	}
}

func main() {
	httpHandler := http.HandlerFunc(mainHandler)
	log.Println("Starting server on :8080")
	if err := http.ListenAndServe(":8080", RealIpMiddleware(httpHandler)); err != nil {
		log.Fatalf("Server failed: %v", err)
	}
}