// steamcache/steamcache.go
package steamcache

import (
	"bufio"
	"context"
	"crypto/sha1"
	"encoding/hex"
	"fmt"
	"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"
	"sort"
	"strings"
	"sync"
	"time"

	"bytes"

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

// extractHashFromSteamPath extracts a hash from Steam depot URLs
// Handles patterns like: /depot/123/chunk/abcdef... or /depot/123/manifest/456/789/hash
func extractHashFromSteamPath(path string) (string, bool) {
	path = strings.TrimPrefix(path, "/")

	parts := strings.Split(path, "/")
	if len(parts) < 3 {
		return "", false
	}

	// Handle chunk files: depot/{id}/chunk/{hash}
	if len(parts) >= 4 && parts[0] == "depot" && parts[2] == "chunk" {
		hash := parts[3]
		// Validate it's a 40-character hex hash
		if len(hash) == 40 && isHexString(hash) {
			return strings.ToLower(hash), true
		}
	}

	// Handle manifest files: depot/{id}/manifest/{manifest_id}/{version}/{hash}
	if len(parts) >= 6 && parts[0] == "depot" && parts[2] == "manifest" {
		hash := parts[5]
		// Note: Manifest hashes can be shorter than 40 characters
		if len(hash) >= 10 && isHexString(hash) {
			return strings.ToLower(hash), true
		}
	}

	return "", false
}

// isHexString checks if a string contains only hexadecimal characters
func isHexString(s string) bool {
	for _, r := range s {
		if !((r >= '0' && r <= '9') || (r >= 'a' && r <= 'f') || (r >= 'A' && r <= 'F')) {
			return false
		}
	}
	return true
}

// generateSteamCacheKey converts Steam depot paths to simplified cache keys
// Input: /depot/1684171/chunk/0016cfc5019b8baa6026aa1cce93e685d6e06c6e
// Output: steam/0016cfc5019b8baa6026aa1cce93e685d6e06c6e
func generateSteamCacheKey(urlPath string) string {
	if hash, ok := extractHashFromSteamPath(urlPath); ok {
		return "steam/" + hash
	}

	// Return empty string for unsupported depot URLs
	return ""
}

// calculateFileHash calculates the SHA1 hash of the given data
func calculateFileHash(data []byte) string {
	hash := sha1.Sum(data)
	return hex.EncodeToString(hash[:])
}

// calculateResponseHash calculates the SHA1 hash of the full HTTP response
func calculateResponseHash(resp *http.Response, bodyData []byte) string {
	hash := sha1.New()

	// Include status line
	statusLine := fmt.Sprintf("HTTP/1.1 %d %s\n", resp.StatusCode, resp.Status)
	hash.Write([]byte(statusLine))

	// Include headers (sorted for consistency)
	headers := make([]string, 0, len(resp.Header))
	for key, values := range resp.Header {
		for _, value := range values {
			headers = append(headers, fmt.Sprintf("%s: %s\n", key, value))
		}
	}
	sort.Strings(headers)
	for _, header := range headers {
		hash.Write([]byte(header))
	}

	// Include empty line between headers and body
	hash.Write([]byte("\n"))

	// Include body
	hash.Write(bodyData)

	return hex.EncodeToString(hash.Sum(nil))
}

// verifyFileHash verifies that the file content matches the expected hash
func verifyFileHash(data []byte, expectedHash string) bool {
	actualHash := calculateFileHash(data)
	return strings.EqualFold(actualHash, expectedHash)
}

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.Split(r.URL.String(), "?")[0]

		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()

		// Read the entire response body into memory for hash verification
		bodyData, err := io.ReadAll(resp.Body)
		if err != nil {
			logger.Logger.Error().Err(err).Str("url", req.URL.String()).Msg("Failed to read response body")

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

			http.Error(w, "Failed to read response body", http.StatusInternalServerError)
			return
		}

		// Extract hash from cache key for verification
		var expectedHash string
		var hasHash bool
		if strings.HasPrefix(cacheKey, "steam/") {
			expectedHash = cacheKey[6:]       // Remove "steam/" prefix
			hasHash = len(expectedHash) == 64 // SHA-256 hashes are 64 characters
		}

		// Hash verification using Steam's X-Content-Sha header and content length verification
		hashVerified := true
		if hasHash {
			// Get the hash from Steam's X-Content-Sha header
			steamHash := resp.Header.Get("X-Content-Sha")

			// Verify using Steam's hash
			if strings.EqualFold(steamHash, expectedHash) {
				// Hash verification succeeded
			} else {
				logger.Logger.Error().
					Str("key", cacheKey).
					Str("expected_hash", expectedHash).
					Str("steam_hash", steamHash).
					Int("content_length", len(bodyData)).
					Msg("Steam hash verification failed - Steam's hash doesn't match filename")
				hashVerified = false
			}
		} else {
			// No hash to verify
		}

		// Always verify content length as an additional safety check
		if resp.ContentLength > 0 && int64(len(bodyData)) != resp.ContentLength {
			logger.Logger.Error().
				Str("key", cacheKey).
				Int("actual_content_length", len(bodyData)).
				Int64("expected_content_length", resp.ContentLength).
				Msg("Content length verification failed")
			hashVerified = false
		}

		// Write to response (always serve the file)
		// 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")
		w.Write(bodyData)

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

		// Only cache the file if hash verification passed (or no hash was present)
		if hashVerified {
			writer, _ := sc.vfs.Create(cachePath, int64(0)) // size is not known in advance
			if writer != nil {
				defer writer.Close()
				// Write the full HTTP response to cache
				resp.Body = io.NopCloser(bytes.NewReader(bodyData)) // Reset body for writing
				resp.Write(writer)
			}
		} else {
			logger.Logger.Warn().
				Str("key", cacheKey).
				Msg("File served but not cached due to hash verification failure")
		}

		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)
}