package main

import (
	"io"
	"log"
	"net/http"
	"net/url"

	"strings"
	"time"

	"github.com/gin-contrib/cors"
	"github.com/gin-gonic/gin"
)

// SecureHeaders sets security-related HTTP headers
func SecureHeaders() gin.HandlerFunc {
	return func(c *gin.Context) {
		c.Writer.Header().Set("X-Content-Type-Options", "nosniff")
		c.Next()
	}
}

// ValidReqPath checks the X-Req-Path is safe for proxying
func ValidReqPath(path string) bool {
	if !strings.HasPrefix(path, "https://") {
		return false
	}
	if strings.Contains(path, "..") {
		return false
	}
	if strings.TrimSpace(path) == "" {
		return false
	}
	return true
}

// Extract the root domain (e.g., "example.com" from "sub.example.com")
func extractRootDomain(urlStr string) (string, error) {
	parsedURL, err := url.Parse(urlStr)
	if err != nil {
		return "", err
	}

	hostname := parsedURL.Hostname()
	parts := strings.Split(hostname, ".")

	// Handle special cases with fewer than 2 parts
	if len(parts) < 2 {
		return hostname, nil
	}

	// Get the last two parts as the root domain
	domain := parts[len(parts)-2] + "." + parts[len(parts)-1]
	return domain, nil
}

// IsValidBackendURL validates that the constructed URL belongs to the expected backend
func IsValidBackendURL(targetURL, backend string) bool {
	targetApexDomain, err := extractRootDomain(targetURL)
	if err != nil {
		return false
	}

	backendApexDomain, err := extractRootDomain(targetURL)
	if err != nil {
		return false
	}

	return targetApexDomain == backendApexDomain
}

func main() {
	r := gin.Default()
	r.Use(SecureHeaders())

	// CORS config for /api/req
	api := r.Group("/api")
	api.Use(cors.New(cors.Config{
		AllowOrigins:     []string{"*"},
		AllowMethods:     []string{"GET", "POST", "PUT", "DELETE", "OPTIONS"},
		AllowHeaders:     []string{"Origin", "Content-Type", "Authorization", "X-Req-Path"},
		ExposeHeaders:    []string{"Content-Length"},
		AllowCredentials: true,
		MaxAge:           12 * time.Hour,
	}))

	// Secure proxy forward
	api.Any("/req", func(c *gin.Context) {
		const backend = "https://deeploy.ml"

		reqPath := c.GetHeader("X-Req-Path")
		if !ValidReqPath(reqPath) {
			c.JSON(http.StatusBadRequest, gin.H{"error": "Invalid X-Req-Path"})
			return
		}

		targetURL := reqPath
		if !IsValidBackendURL(targetURL, backend) {
			c.JSON(http.StatusForbidden, gin.H{"error": "Target not allowed"})
			return
		}

		// Prepare the proxied request
		req, err := http.NewRequest(c.Request.Method, targetURL, c.Request.Body)
		if err != nil {
			c.JSON(http.StatusInternalServerError, gin.H{"error": "Request creation failed"})
			return
		}
		req.Header = c.Request.Header.Clone()
		req.Header.Del("X-Req-Path") // Don't forward custom header

		client := &http.Client{
			Timeout: 60 * time.Second,
		}
		resp, err := client.Do(req)
		if err != nil {
			log.Printf("Error proxying request: Method=%s URL=%s Headers=%v Error=%v",
				req.Method, req.URL, req.Header, err)
			c.JSON(http.StatusBadGateway, gin.H{"error": "Backend request failed."})
			return
		}

		defer resp.Body.Close()

		// Copy headers & status from the backend
		for k, v := range resp.Header {
			for _, vv := range v {
				c.Writer.Header().Add(k, vv)
			}
		}
		c.Status(resp.StatusCode)
		io.Copy(c.Writer, resp.Body)
	})

	// Serve static files from ./frontend, no directory listing
	r.StaticFS("/assets", gin.Dir("./assets", false))

	r.GET("/", func(c *gin.Context) {
		log.Println("root route")
		// Load HTML templates from ./frontend folder
		// Render the template by name
		r.LoadHTMLFiles("index.html")
		// Render the template by name

		c.HTML(http.StatusOK, "index.html", gin.H{})
	})

	// Route to serve parsed HTML template partials
	r.GET("/partials", func(c *gin.Context) {
		// Load HTML templates from ./partials folder
		r.LoadHTMLGlob("./frontend/partials/*")
		// Render the template by name
		c.HTML(http.StatusOK, "index_partial.html", gin.H{})
	})

	log.Println("Serving on http://localhost:7860")
	r.Run(":7860")
}