// steamcache/steamcache.go
package steamcache

import (
	"bufio"
	"context"
	"crypto/sha256"
	"encoding/hex"
	"io"
	"net"
	"net/http"
	"net/url"
	"os"
	"s1d3sw1ped/SteamCache2/steamcache/logger"
	"s1d3sw1ped/SteamCache2/vfs"
	"s1d3sw1ped/SteamCache2/vfs/cache"
	"s1d3sw1ped/SteamCache2/vfs/disk"
	"s1d3sw1ped/SteamCache2/vfs/gc"
	"s1d3sw1ped/SteamCache2/vfs/memory"
	"strings"
	"sync"
	"time"

	"github.com/docker/go-units"
)

// generateURLHash creates a SHA256 hash of the entire URL path for cache key
func generateURLHash(urlPath string) string {
	hash := sha256.Sum256([]byte(urlPath))
	return hex.EncodeToString(hash[:])
}

// generateSteamCacheKey creates a cache key from the URL path using SHA256
// Input: /depot/1684171/chunk/0016cfc5019b8baa6026aa1cce93e685d6e06c6e
// Output: steam/a1b2c3d4e5f678901234567890123456789012345678901234567890
func generateSteamCacheKey(urlPath string) string {
	// Handle Steam depot URLs by creating a SHA256 hash of the entire path
	if strings.HasPrefix(urlPath, "/depot/") {
		return "steam/" + generateURLHash(urlPath)
	}

	// For non-Steam URLs, return empty string (not cached)
	return ""
}

var hopByHopHeaders = map[string]struct{}{
	"Connection":          {},
	"Keep-Alive":          {},
	"Proxy-Authenticate":  {},
	"Proxy-Authorization": {},
	"TE":                  {},
	"Trailer":             {},
	"Transfer-Encoding":   {},
	"Upgrade":             {},
	"Date":                {},
	"Server":              {},
}

var (
	// Request coalescing structures
	coalescedRequests   = make(map[string]*coalescedRequest)
	coalescedRequestsMu sync.RWMutex
)

type coalescedRequest struct {
	responseChan chan *http.Response
	errorChan    chan error
	waitingCount int
	done         bool
	mu           sync.Mutex
}

func newCoalescedRequest() *coalescedRequest {
	return &coalescedRequest{
		responseChan: make(chan *http.Response, 1),
		errorChan:    make(chan error, 1),
		waitingCount: 1,
		done:         false,
	}
}

func (cr *coalescedRequest) addWaiter() {
	cr.mu.Lock()
	defer cr.mu.Unlock()
	cr.waitingCount++
}

func (cr *coalescedRequest) complete(resp *http.Response, err error) {
	cr.mu.Lock()
	defer cr.mu.Unlock()
	if cr.done {
		return
	}
	cr.done = true

	if err != nil {
		select {
		case cr.errorChan <- err:
		default:
		}
	} else {
		select {
		case cr.responseChan <- resp:
		default:
		}
	}
}

// getOrCreateCoalescedRequest gets an existing coalesced request or creates a new one
func getOrCreateCoalescedRequest(cacheKey string) (*coalescedRequest, bool) {
	coalescedRequestsMu.Lock()
	defer coalescedRequestsMu.Unlock()

	if cr, exists := coalescedRequests[cacheKey]; exists {
		cr.addWaiter()
		return cr, false
	}

	cr := newCoalescedRequest()
	coalescedRequests[cacheKey] = cr
	return cr, true
}

// removeCoalescedRequest removes a completed coalesced request
func removeCoalescedRequest(cacheKey string) {
	coalescedRequestsMu.Lock()
	defer coalescedRequestsMu.Unlock()
	delete(coalescedRequests, cacheKey)
}

type SteamCache struct {
	address  string
	upstream string

	vfs vfs.VFS

	memory *memory.MemoryFS
	disk   *disk.DiskFS

	memorygc *gc.GCFS
	diskgc   *gc.GCFS

	server *http.Server
	client *http.Client
	cancel context.CancelFunc
	wg     sync.WaitGroup
}

func New(address string, memorySize string, diskSize string, diskPath, upstream, memoryGC, diskGC string) *SteamCache {
	memorysize, err := units.FromHumanSize(memorySize)
	if err != nil {
		panic(err)
	}

	disksize, err := units.FromHumanSize(diskSize)
	if err != nil {
		panic(err)
	}

	c := cache.New()

	var m *memory.MemoryFS
	var mgc *gc.GCFS
	if memorysize > 0 {
		m = memory.New(memorysize)
		memoryGCAlgo := gc.GCAlgorithm(memoryGC)
		if memoryGCAlgo == "" {
			memoryGCAlgo = gc.LRU // default to LRU
		}
		mgc = gc.New(m, memoryGCAlgo)
	}

	var d *disk.DiskFS
	var dgc *gc.GCFS
	if disksize > 0 {
		d = disk.New(diskPath, disksize)
		diskGCAlgo := gc.GCAlgorithm(diskGC)
		if diskGCAlgo == "" {
			diskGCAlgo = gc.LRU // default to LRU
		}
		dgc = gc.New(d, diskGCAlgo)
	}

	// configure the cache to match the specified mode (memory only, disk only, or memory and disk) based on the provided sizes
	if disksize == 0 && memorysize != 0 {
		//memory only mode - no disk

		c.SetSlow(mgc)
	} else if disksize != 0 && memorysize == 0 {
		// disk only mode

		c.SetSlow(dgc)
	} else if disksize != 0 && memorysize != 0 {
		// memory and disk mode

		c.SetFast(mgc)
		c.SetSlow(dgc)
	} else {
		// no memory or disk isn't a valid configuration
		logger.Logger.Error().Bool("memory", false).Bool("disk", false).Msg("configuration invalid :( exiting")
		os.Exit(1)
	}

	transport := &http.Transport{
		MaxIdleConns:        200,               // Increased from 100
		MaxIdleConnsPerHost: 50,                // Increased from 10
		IdleConnTimeout:     120 * time.Second, // Increased from 90s
		DialContext: (&net.Dialer{
			Timeout:   30 * time.Second,
			KeepAlive: 30 * time.Second,
		}).DialContext,
		TLSHandshakeTimeout:   15 * time.Second, // Increased from 10s
		ResponseHeaderTimeout: 30 * time.Second, // Increased from 10s
		ExpectContinueTimeout: 5 * time.Second,  // Increased from 1s
		DisableCompression:    true,             // Steam doesn't use compression
		ForceAttemptHTTP2:     true,             // Enable HTTP/2 if available
	}

	client := &http.Client{
		Transport: transport,
		Timeout:   120 * time.Second, // Increased from 60s
	}

	sc := &SteamCache{
		upstream: upstream,
		address:  address,
		vfs:      c,
		memory:   m,
		disk:     d,
		memorygc: mgc,
		diskgc:   dgc,
		client:   client,
		server: &http.Server{
			Addr:              address,
			ReadTimeout:       30 * time.Second,  // Increased
			WriteTimeout:      60 * time.Second,  // Increased
			IdleTimeout:       120 * time.Second, // Good for keep-alive
			ReadHeaderTimeout: 10 * time.Second,  // New, for header attacks
			MaxHeaderBytes:    1 << 20,           // 1MB, optional
		},
	}

	// Log GC algorithm configuration
	if m != nil {
		logger.Logger.Info().Str("memory_gc", memoryGC).Msg("Memory cache GC algorithm configured")
	}
	if d != nil {
		logger.Logger.Info().Str("disk_gc", diskGC).Msg("Disk cache GC algorithm configured")
	}

	if d != nil {
		if d.Size() > d.Capacity() {
			gcHandler := gc.GetGCAlgorithm(gc.GCAlgorithm(diskGC))
			gcHandler(d, uint(d.Size()-d.Capacity()))
		}
	}

	return sc
}

func (sc *SteamCache) Run() {
	if sc.upstream != "" {
		resp, err := sc.client.Get(sc.upstream)
		if err != nil || resp.StatusCode != http.StatusOK {
			logger.Logger.Error().Err(err).Int("status_code", resp.StatusCode).Str("upstream", sc.upstream).Msg("Failed to connect to upstream server")
			os.Exit(1)
		}
		resp.Body.Close()
	}

	sc.server.Handler = sc
	ctx, cancel := context.WithCancel(context.Background())
	sc.cancel = cancel

	sc.wg.Add(1)
	go func() {
		defer sc.wg.Done()
		err := sc.server.ListenAndServe()
		if err != nil && err != http.ErrServerClosed {
			logger.Logger.Error().Err(err).Msg("Failed to start SteamCache2")
			os.Exit(1)
		}
	}()

	<-ctx.Done()
	sc.server.Shutdown(ctx)
	sc.wg.Wait()
}

func (sc *SteamCache) Shutdown() {
	if sc.cancel != nil {
		sc.cancel()
	}
	sc.wg.Wait()
}

func (sc *SteamCache) ServeHTTP(w http.ResponseWriter, r *http.Request) {
	if r.Method != http.MethodGet {
		logger.Logger.Warn().Str("method", r.Method).Msg("Only GET method is supported")
		http.Error(w, "Only GET method is supported", http.StatusMethodNotAllowed)
		return
	}

	if r.URL.Path == "/" {
		w.WriteHeader(http.StatusOK) // this is used by steamcache2's upstream verification at startup
		return
	}

	if r.URL.String() == "/lancache-heartbeat" {
		w.Header().Add("X-LanCache-Processed-By", "SteamCache2")
		w.WriteHeader(http.StatusNoContent)
		w.Write(nil)
		return
	}

	if strings.HasPrefix(r.URL.String(), "/depot/") {
		// trim the query parameters from the URL path
		// this is necessary because the cache key should not include query parameters
		urlPath, _, _ := strings.Cut(r.URL.String(), "?")

		tstart := time.Now()

		// Generate simplified Steam cache key: steam/{hash}
		cacheKey := generateSteamCacheKey(urlPath)

		if cacheKey == "" {
			logger.Logger.Warn().Str("url", urlPath).Msg("Invalid URL")
			http.Error(w, "Invalid URL", http.StatusBadRequest)
			return
		}

		w.Header().Add("X-LanCache-Processed-By", "SteamCache2") // SteamPrefill uses this header to determine if the request was processed by the cache maybe steam uses it too

		cachePath := cacheKey // You may want to add a .http or .cache extension for clarity

		// Try to serve from cache
		file, err := sc.vfs.Open(cachePath)
		if err == nil {
			defer file.Close()
			buf := bufio.NewReader(file)
			resp, err := http.ReadResponse(buf, nil)
			if err == nil {
				// Remove hop-by-hop and server-specific headers
				for k, vv := range resp.Header {
					if _, skip := hopByHopHeaders[http.CanonicalHeaderKey(k)]; skip {
						continue
					}
					for _, v := range vv {
						w.Header().Add(k, v)
					}
				}
				// Add our own headers
				w.Header().Set("X-LanCache-Status", "HIT")
				w.Header().Set("X-LanCache-Processed-By", "SteamCache2")
				w.WriteHeader(resp.StatusCode)
				io.Copy(w, resp.Body)
				resp.Body.Close()
				logger.Logger.Info().
					Str("key", cacheKey).
					Str("host", r.Host).
					Str("status", "HIT").
					Dur("duration", time.Since(tstart)).
					Msg("request")
				return
			}
		}

		// Check for coalesced request (another client already downloading this)
		coalescedReq, isNew := getOrCreateCoalescedRequest(cacheKey)
		if !isNew {
			// Wait for the existing download to complete
			logger.Logger.Debug().
				Str("key", cacheKey).
				Int("waiting_clients", coalescedReq.waitingCount).
				Msg("Joining coalesced request")

			select {
			case resp := <-coalescedReq.responseChan:
				// Use the downloaded response
				defer resp.Body.Close()
				bodyData, err := io.ReadAll(resp.Body)
				if err != nil {
					logger.Logger.Error().Err(err).Str("key", cacheKey).Msg("Failed to read coalesced response body")
					http.Error(w, "Failed to read response body", http.StatusInternalServerError)
					return
				}

				// Serve the response
				for k, vv := range resp.Header {
					if _, skip := hopByHopHeaders[http.CanonicalHeaderKey(k)]; skip {
						continue
					}
					for _, v := range vv {
						w.Header().Add(k, v)
					}
				}
				w.Header().Set("X-LanCache-Status", "HIT-COALESCED")
				w.Header().Set("X-LanCache-Processed-By", "SteamCache2")
				w.WriteHeader(resp.StatusCode)
				w.Write(bodyData)

				logger.Logger.Info().
					Str("key", cacheKey).
					Str("host", r.Host).
					Str("status", "HIT-COALESCED").
					Dur("duration", time.Since(tstart)).
					Msg("request")

				return

			case err := <-coalescedReq.errorChan:
				logger.Logger.Error().Err(err).Str("key", cacheKey).Msg("Coalesced request failed")
				http.Error(w, "Upstream request failed", http.StatusInternalServerError)
				return
			}
		}

		// Remove coalesced request when done
		defer removeCoalescedRequest(cacheKey)

		var req *http.Request
		if sc.upstream != "" { // if an upstream server is configured, proxy the request to the upstream server
			ur, err := url.JoinPath(sc.upstream, urlPath)
			if err != nil {
				logger.Logger.Error().Err(err).Str("upstream", sc.upstream).Msg("Failed to join URL path")
				http.Error(w, "Failed to join URL path", http.StatusInternalServerError)
				return
			}

			req, err = http.NewRequest(http.MethodGet, ur, nil)
			if err != nil {
				logger.Logger.Error().Err(err).Str("upstream", sc.upstream).Msg("Failed to create request")
				http.Error(w, "Failed to create request", http.StatusInternalServerError)
				return
			}
			req.Host = r.Host
		} else { // if no upstream server is configured, proxy the request to the host specified in the request
			host := r.Host
			if r.Header.Get("X-Sls-Https") == "enable" {
				host = "https://" + host
			} else {
				host = "http://" + host
			}

			ur, err := url.JoinPath(host, urlPath)
			if err != nil {
				logger.Logger.Error().Err(err).Str("host", host).Msg("Failed to join URL path")
				http.Error(w, "Failed to join URL path", http.StatusInternalServerError)
				return
			}

			req, err = http.NewRequest(http.MethodGet, ur, nil)
			if err != nil {
				logger.Logger.Error().Err(err).Str("host", host).Msg("Failed to create request")
				http.Error(w, "Failed to create request", http.StatusInternalServerError)
				return
			}
		}

		// Copy headers from the original request to the new request
		for key, values := range r.Header {
			for _, value := range values {
				req.Header.Add(key, value)
			}
		}

		// Retry logic
		backoffSchedule := []time.Duration{1 * time.Second, 3 * time.Second, 10 * time.Second}
		var resp *http.Response
		for i, backoff := range backoffSchedule {
			resp, err = sc.client.Do(req)
			if err == nil && resp.StatusCode == http.StatusOK {
				break
			}
			if i < len(backoffSchedule)-1 {
				time.Sleep(backoff)
			}
		}
		if err != nil || resp.StatusCode != http.StatusOK {
			logger.Logger.Error().Err(err).Str("url", req.URL.String()).Msg("Failed to fetch the requested URL")

			// Complete coalesced request with error
			if isNew {
				coalescedReq.complete(nil, err)
			}

			http.Error(w, "Failed to fetch the requested URL", http.StatusInternalServerError)
			return
		}
		defer resp.Body.Close()

		// Fast path: Flexible lightweight validation for all files
		// Multiple validation layers ensure data integrity without blocking legitimate Steam content

		// Method 1: HTTP Status Validation
		if resp.StatusCode != http.StatusOK {
			logger.Logger.Error().
				Str("url", req.URL.String()).
				Int("status_code", resp.StatusCode).
				Msg("Steam returned non-OK status")
			http.Error(w, "Upstream server error", http.StatusBadGateway)
			return
		}

		// Method 2: Content-Type Validation (Steam files should be application/x-steam-chunk)
		contentType := resp.Header.Get("Content-Type")
		if contentType != "" && !strings.Contains(contentType, "application/x-steam-chunk") {
			logger.Logger.Warn().
				Str("url", req.URL.String()).
				Str("content_type", contentType).
				Msg("Unexpected content type from Steam - expected application/x-steam-chunk")
		}

		// Method 3: Content-Length Validation
		expectedSize := resp.ContentLength

		// Reject only truly invalid content lengths (zero or negative)
		if expectedSize <= 0 {
			logger.Logger.Error().
				Str("url", req.URL.String()).
				Int64("content_length", expectedSize).
				Msg("Invalid content length, rejecting file")
			http.Error(w, "Invalid content length", http.StatusBadGateway)
			return
		}

		// Content length is valid - no size restrictions to keep logs clean

		// Lightweight validation passed - trust the Content-Length and HTTP status
		// This provides good integrity with minimal performance overhead
		validationPassed := true

		// Write to response (stream the file directly)
		// Remove hop-by-hop and server-specific headers
		for k, vv := range resp.Header {
			if _, skip := hopByHopHeaders[http.CanonicalHeaderKey(k)]; skip {
				continue
			}
			for _, v := range vv {
				w.Header().Add(k, v)
			}
		}
		// Add our own headers
		w.Header().Set("X-LanCache-Status", "MISS")
		w.Header().Set("X-LanCache-Processed-By", "SteamCache2")

		// Stream the response body directly to client (no memory buffering)
		io.Copy(w, resp.Body)

		// Complete coalesced request for waiting clients
		if isNew {
			// Create a new response for coalesced clients with a fresh body
			coalescedResp := &http.Response{
				StatusCode: resp.StatusCode,
				Status:     resp.Status,
				Header:     make(http.Header),
				Body:       io.NopCloser(strings.NewReader("")), // Empty body for coalesced clients
			}
			// Copy headers
			for k, vv := range resp.Header {
				coalescedResp.Header[k] = vv
			}
			coalescedReq.complete(coalescedResp, nil)
		}

		// Cache the file if validation passed
		if validationPassed {
			// Create a new request to fetch the file again for caching
			cacheReq, err := http.NewRequest(http.MethodGet, req.URL.String(), nil)
			if err == nil {
				// Copy original headers
				for k, vv := range req.Header {
					cacheReq.Header[k] = vv
				}

				// Fetch fresh copy for caching
				cacheResp, err := sc.client.Do(cacheReq)
				if err == nil {
					defer cacheResp.Body.Close()
					// Use the validated size from the original response
					writer, _ := sc.vfs.Create(cachePath, expectedSize)
					if writer != nil {
						defer writer.Close()
						io.Copy(writer, cacheResp.Body)
					}
				}
			}
		}

		logger.Logger.Info().
			Str("key", cacheKey).
			Str("host", r.Host).
			Str("status", "MISS").
			Dur("duration", time.Since(tstart)).
			Msg("request")

		return
	}

	if r.URL.Path == "/favicon.ico" {
		w.WriteHeader(http.StatusNoContent)
		return
	}

	if r.URL.Path == "/robots.txt" {
		w.Header().Set("Content-Type", "text/plain")
		w.WriteHeader(http.StatusOK)
		w.Write([]byte("User-agent: *\nDisallow: /\n"))
		return
	}

	logger.Logger.Warn().Str("url", r.URL.String()).Msg("Not found")
	http.Error(w, "Not found", http.StatusNotFound)
}