package main

import (
	"bytes"
	"compress/gzip"
	"compress/zlib"
	"encoding/json"
	"fmt"
	"html"
	"io"
	"log"
	"net/http"
	"net/url"
	"os"
	"strings"

	"github.com/andybalholm/brotli"
)

// FlareSolverrRequest represents the JSON payload to send to FlareSolverr.
type FlareSolverrRequest struct {
	Cmd        string `json:"cmd"`
	URL        string `json:"url"`
	MaxTimeout int    `json:"maxTimeout,omitempty"`
}

// FlareSolverrResponse represents the structure to parse FlareSolverr's response.
type FlareSolverrResponse struct {
	Solution struct {
		URL       string            `json:"url"`
		Status    int               `json:"status"`
		Cookies   []interface{}     `json:"cookies"`
		UserAgent string            `json:"userAgent"`
		Headers   map[string]string `json:"headers"`
		Response  string            `json:"response"`
	} `json:"solution"`
	Status  string `json:"status"`
	Message string `json:"message"`
}

func handler(flsEndpoint string) http.HandlerFunc {
	return func(w http.ResponseWriter, r *http.Request) {
		// Check compression support in order of preference
		acceptEncoding := r.Header.Get("Accept-Encoding")
		var writer io.Writer = w
		var closer io.Closer

		switch {
		case strings.Contains(acceptEncoding, "br"):
			// Brotli compression
			writer = brotli.NewWriter(w)
			closer = writer.(io.Closer)
			w.Header().Set("Content-Encoding", "br")
		case strings.Contains(acceptEncoding, "gzip"):
			// Gzip compression
			writer = gzip.NewWriter(w)
			closer = writer.(io.Closer)
			w.Header().Set("Content-Encoding", "gzip")
		case strings.Contains(acceptEncoding, "deflate"):
			// Deflate compression
			writer = zlib.NewWriter(w)
			closer = writer.(io.Closer)
			w.Header().Set("Content-Encoding", "deflate")
		}

		if closer != nil {
			defer closer.Close()
		}

		// Ignore common browser requests
		if r.URL.Path == "/favicon.ico" || r.URL.Path == "/robots.txt" {
			log.Printf("Ignoring browser request: %s", r.URL.Path)
			http.NotFound(w, r)
			return
		}

		// Expect the request path to be "/<urlencoded URI>".
		encoded := r.URL.Path[1:]
		if encoded == "" {
			log.Printf("Error: No URL provided in path")
			http.Error(w, "No URL provided in path", http.StatusBadRequest)
			return
		}

		// Decode the URL-encoded URI.
		decodedURL, err := url.QueryUnescape(encoded)
		if err != nil {
			log.Printf("Error decoding URL '%s': %v", encoded, err)
			http.Error(w, "Invalid URL encoding", http.StatusBadRequest)
			return
		}
		log.Printf("Processing request for URL: %s", decodedURL)

		// Create the JSON payload for FlareSolverr.
		reqPayload := FlareSolverrRequest{
			Cmd:        "request.get",
			URL:        decodedURL,
			MaxTimeout: 60000,
		}
		payloadBytes, err := json.Marshal(reqPayload)
		if err != nil {
			log.Printf("Error creating JSON payload: %v", err)
			http.Error(w, "Error creating JSON payload", http.StatusInternalServerError)
			return
		}

		// Send the POST request to the FlareSolverr endpoint.
		log.Printf("Sending request to FlareSolverr: %s", flsEndpoint)
		resp, err := http.Post(flsEndpoint, "application/json", bytes.NewBuffer(payloadBytes))
		if err != nil {
			log.Printf("Error querying FlareSolverr: %v", err)
			http.Error(w, fmt.Sprintf("Error querying FlareSolverr: %v", err), http.StatusInternalServerError)
			return
		}
		defer resp.Body.Close()

		// Read the response body.
		body, err := io.ReadAll(resp.Body)
		if err != nil {
			log.Printf("Error reading response from FlareSolverr: %v", err)
			http.Error(w, "Error reading response from FlareSolverr", http.StatusInternalServerError)
			return
		}

		// Parse the FlareSolverr response
		var flsResp FlareSolverrResponse
		if err := json.Unmarshal(body, &flsResp); err != nil {
			log.Printf("Error parsing FlareSolverr response: %v", err)
			http.Error(w, "Error parsing FlareSolverr response", http.StatusInternalServerError)
			return
		}

		// Check if we have a valid response
		if flsResp.Solution.Response == "" {
			log.Printf("No response content received from FlareSolverr for URL: %s", decodedURL)
			http.Error(w, "No response content available", http.StatusServiceUnavailable)
			return
		}

		// Log the response status
		log.Printf("Received response from FlareSolverr with status %d for URL: %s", flsResp.Solution.Status, decodedURL)

		// Check if the response is XML/RSS
		response := flsResp.Solution.Response
		if strings.Contains(response, "&lt;?xml") || strings.Contains(response, "&lt;rss") {
			log.Printf("Detected XML/RSS response for URL: %s", decodedURL)
			// Remove the HTML wrapper if present
			if strings.Contains(response, "<html>") {
				log.Printf("Found HTML wrapper, attempting to extract XML content")
				start := strings.Index(response, "&lt;")
				end := strings.LastIndex(response, "&gt;") + 4
				if start >= 0 && end > start {
					log.Printf("Extracting XML content from position %d to %d", start, end)
					response = response[start:end]
				} else {
					log.Printf("Could not find XML content boundaries in HTML wrapper")
				}
			} else {
				log.Printf("No HTML wrapper found, processing raw response")
			}
			// Decode HTML entities
			log.Printf("Decoding HTML entities in response")
			response = html.UnescapeString(response)
			w.Header().Set("Content-Type", "application/xml")
			log.Printf("Response processed and ready to send as XML")
		} else {
			log.Printf("Response is not XML/RSS, sending as HTML")
			w.Header().Set("Content-Type", "text/html")
		}

		// Write the response using the appropriate writer (compressed or not)
		if _, err := writer.Write([]byte(response)); err != nil {
			log.Printf("Error writing response: %v", err)
			http.Error(w, "Error writing response", http.StatusInternalServerError)
			return
		}
	}
}

func main() {
	// Get configuration from environment variables with defaults
	port := "38080"
	if p := os.Getenv("PORT"); p != "" {
		port = p
	}

	flsEndpoint := "http://localhost:8191/v1"
	if e := os.Getenv("FLS_ENDPOINT"); e != "" {
		flsEndpoint = e
	}

	http.HandleFunc("/", handler(flsEndpoint))
	log.Printf("Starting proxy server on :%s (FlareSolverr endpoint: %s)\n", port, flsEndpoint)
	if err := http.ListenAndServe(":"+port, nil); err != nil {
		log.Fatal(err)
	}
}