package hosterrorscache

import (
	"errors"
	"net"
	"net/url"
	"regexp"
	"strings"
	"sync"
	"sync/atomic"

	"github.com/projectdiscovery/gcache"
	"github.com/projectdiscovery/gologger"
	"github.com/projectdiscovery/nuclei/v3/pkg/catalog/config"
	"github.com/projectdiscovery/nuclei/v3/pkg/protocols/common/contextargs"
	"github.com/projectdiscovery/nuclei/v3/pkg/types/nucleierr"
	"github.com/projectdiscovery/utils/errkit"
	stringsutil "github.com/projectdiscovery/utils/strings"
	"golang.org/x/net/idna"
)

// CacheInterface defines the signature of the hosterrorscache so that
// users of Nuclei as embedded lib may implement their own cache
type CacheInterface interface {
	SetVerbose(verbose bool)                                                  // log verbosely
	Close()                                                                   // close the cache
	Check(protoType string, ctx *contextargs.Context) bool                    // return true if the host should be skipped
	Remove(ctx *contextargs.Context)                                          // remove a host from the cache
	MarkFailed(protoType string, ctx *contextargs.Context, err error)         // record a failure (and cause) for the host
	MarkFailedOrRemove(protoType string, ctx *contextargs.Context, err error) // record a failure (and cause) for the host or remove it
	IsPermanentErr(ctx *contextargs.Context, err error) bool                  // return true if the error is permanent for the host
}

var (
	_ CacheInterface = (*Cache)(nil)
)

// Cache is a cache for host based errors. It allows skipping
// certain hosts based on an error threshold.
//
// It uses an LRU cache internally for skipping unresponsive hosts
// that remain so for a duration.
type Cache struct {
	MaxHostError    int
	verbose         bool
	maxHostsCount   int
	failedTargets   gcache.Cache[string, *cacheItem]
	failedTargetsMu sync.RWMutex
	TrackError      []string
	closed          atomic.Bool
}

type cacheItem struct {
	sync.Once
	errors         atomic.Int32
	isPermanentErr bool
	cause          error // optional cause
	mu             sync.Mutex
}

const DefaultMaxHostsCount = 10000

// New returns a new host max errors cache
func New(maxHostError, maxHostsCount int, trackError []string) *Cache {
	return &Cache{
		MaxHostError:  maxHostError,
		maxHostsCount: maxHostsCount,
		TrackError:    trackError,
	}
}

func (c *Cache) getFailedTargets(create bool) gcache.Cache[string, *cacheItem] {
	c.failedTargetsMu.RLock()
	cache := c.failedTargets
	c.failedTargetsMu.RUnlock()
	if cache != nil || !create || c.closed.Load() {
		return cache
	}

	c.failedTargetsMu.Lock()
	defer c.failedTargetsMu.Unlock()
	if c.failedTargets == nil && !c.closed.Load() {
		c.failedTargets = gcache.New[string, *cacheItem](c.maxHostsCount).ARC().Build()
	}
	return c.failedTargets
}

// SetVerbose sets the cache to log at verbose level
func (c *Cache) SetVerbose(verbose bool) {
	c.verbose = verbose
}

// Close closes the host errors cache
func (c *Cache) Close() {
	if !c.closed.CompareAndSwap(false, true) {
		return
	}
	c.failedTargetsMu.Lock()
	cache := c.failedTargets
	c.failedTargets = nil
	c.failedTargetsMu.Unlock()
	if cache != nil && config.DefaultConfig.IsDebugArgEnabled(config.DebugArgHostErrorStats) {
		items := cache.GetALL(false)
		for k, val := range items {
			gologger.Info().Label("MaxHostErrorStats").Msgf("Host: %s, Errors: %d", k, val.errors.Load())
		}
	}
	// The cache owns no background resource. Purge rebuilds an empty ARC with
	// the original capacity, which allocates several megabytes immediately
	// before callers discard a closed cache. Releasing the reference lets the
	// garbage collector reclaim both entries and storage without replacement.
}

// NormalizeCacheValue processes the input value and returns a normalized cache
// value.
func (c *Cache) NormalizeCacheValue(value string) string {
	var normalizedValue = value

	u, err := url.ParseRequestURI(value)
	if err != nil || u.Host == "" {
		if strings.Contains(value, ":") {
			return normalizedValue
		}
		u, err2 := url.ParseRequestURI("https://" + value)
		if err2 != nil {
			return normalizedValue
		}

		normalizedValue = u.Host
	} else {
		port := u.Port()
		if port == "" {
			switch u.Scheme {
			case "https":
				normalizedValue = net.JoinHostPort(u.Host, "443")
			case "http":
				normalizedValue = net.JoinHostPort(u.Host, "80")
			}
		} else {
			normalizedValue = u.Host
		}
	}

	return normalizedValue
}

// ErrUnresponsiveHost is returned when a host is unresponsive
// var ErrUnresponsiveHost = errors.New("skipping as host is unresponsive")

// Check returns true if a host should be skipped as it has been
// unresponsive for a certain number of times.
//
// The value can be many formats -
//   - URL: https?:// type
//   - Host:port type
//   - host type
func (c *Cache) Check(protoType string, ctx *contextargs.Context) bool {
	failedTargets := c.getFailedTargets(false)
	if failedTargets == nil {
		return false
	}
	finalValue := c.GetKeyFromContext(ctx, nil)

	cache, err := failedTargets.GetIFPresent(finalValue)
	if err != nil {
		return false
	}

	cache.mu.Lock()
	defer cache.mu.Unlock()

	if cache.isPermanentErr {
		cache.Do(func() {
			gologger.Info().Msgf("Skipped %s from target list as found unresponsive permanently: %s", finalValue, cache.cause)
		})
		return true
	}

	if cache.errors.Load() >= int32(c.MaxHostError) {
		cache.Do(func() {
			gologger.Info().Msgf("Skipped %s from target list as found unresponsive %d times", finalValue, cache.errors.Load())
		})
		return true
	}

	return false
}

// Remove removes a host from the cache
func (c *Cache) Remove(ctx *contextargs.Context) {
	failedTargets := c.getFailedTargets(false)
	if failedTargets == nil {
		return
	}
	key := c.GetKeyFromContext(ctx, nil)
	_ = failedTargets.Remove(key) // remove even the cache is not present
}

// MarkFailed marks a host as failed previously
//
// Deprecated: Use MarkFailedOrRemove instead.
func (c *Cache) MarkFailed(protoType string, ctx *contextargs.Context, err error) {
	if err == nil {
		return
	}

	c.MarkFailedOrRemove(protoType, ctx, err)
}

// MarkFailedOrRemove marks a host as failed previously or removes it
func (c *Cache) MarkFailedOrRemove(protoType string, ctx *contextargs.Context, err error) {
	// A failure that occurs because the caller's context was cancelled or hit its
	// deadline is not the host's fault (e.g. the parent scan was cancelled). Such
	// errors classify as temporary/deadline and would otherwise be counted, so
	// ignore them. A nil error (success) still resets the host below.
	if err != nil && ctx != nil && ctx.Context() != nil && ctx.Context().Err() != nil {
		return
	}

	if err != nil && !c.checkError(protoType, err) {
		return
	}

	if err == nil {
		// Remove the host from cache
		//
		// NOTE(dwisiswant0): The decision was made to completely remove the
		// cached entry for the host instead of simply decrementing the error
		// count (using `(atomic.Int32).Swap` to update the value to `N-1`).
		// This approach was chosen because the error handling logic operates
		// concurrently, and decrementing the count could lead to UB (unexpected
		// behavior) even when the error is `nil`.
		//
		// To clarify, consider the following scenario where the error
		// encountered does NOT belong to the permanent network error category
		// (`errkit.ErrKindNetworkPermanent`):
		//
		// 1. Iteration 1: A timeout error occurs, and the error count for the
		//    host is incremented.
		// 2. Iteration 2: Another timeout error is encountered, leading to
		//    another increment in the host's error count.
		// 3. Iteration 3: A third timeout error happens, which increments the
		//    error count further. At this point, the host is flagged as
		//    unresponsive.
		// 4. Iteration 4: The host becomes reachable (no error or a transient
		//    issue resolved). Instead of performing a no-op and leaving the
		//    host in the cache, the host entry is removed entirely to reset its
		//    state.
		// 5. Iteration 5: A subsequent timeout error occurs after the host was
		//    removed and re-added to the cache. The error count is reset and
		//    starts from 1 again.
		//
		// This removal strategy ensures the cache is updated dynamically to
		// reflect the current state of the host without persisting stale or
		// irrelevant error counts that could interfere with future error
		// handling and tracking logic.
		c.Remove(ctx)

		return
	}

	cacheKey := c.GetKeyFromContext(ctx, err)
	failedTargets := c.getFailedTargets(true)
	if failedTargets == nil {
		return
	}
	cache, cacheErr := failedTargets.GetIFPresent(cacheKey)
	if errors.Is(cacheErr, gcache.KeyNotFoundError) {
		cache = &cacheItem{errors: atomic.Int32{}}
	}

	cache.mu.Lock()
	defer cache.mu.Unlock()

	if errkit.IsKind(err, errkit.ErrKindNetworkPermanent) {
		cache.isPermanentErr = true
	}

	cache.cause = err
	cache.errors.Add(1)

	_ = failedTargets.Set(cacheKey, cache)
}

// IsPermanentErr returns true if the error is permanent for the host.
func (c *Cache) IsPermanentErr(ctx *contextargs.Context, err error) bool {
	if err == nil {
		return false
	}

	if errkit.IsKind(err, errkit.ErrKindNetworkPermanent) {
		return true
	}

	failedTargets := c.getFailedTargets(false)
	if failedTargets == nil {
		return false
	}
	cacheKey := c.GetKeyFromContext(ctx, err)
	cache, cacheErr := failedTargets.GetIFPresent(cacheKey)
	if cacheErr != nil {
		return false
	}

	cache.mu.Lock()
	defer cache.mu.Unlock()

	return cache.isPermanentErr
}

// GetKeyFromContext returns the key for the cache from the context
func (c *Cache) GetKeyFromContext(ctx *contextargs.Context, err error) string {
	// Note:
	// ideally any changes made to remote addr in template like {{Hostname}}:81 etc
	// should be reflected in contextargs but it is not yet reflected in some cases
	// and needs refactor of ScanContext + ContextArgs to achieve that
	// i.e why we use real address from error if present
	var address string

	// 1. the address carried inside the error (if the transport sets it)
	if err != nil {
		if v := errkit.GetAttrValue(err, "address"); v.Any() != nil {
			address = v.String()
		}
	}

	if address == "" {
		address = ctx.MetaInput.Address()
	} else if ctx.MetaInput.CustomIP != "" {
		address = pinnedAddress(ctx.MetaInput, address)
	}

	finalValue := c.NormalizeCacheValue(address)
	return finalValue
}

// pinnedAddress maps the address of a failed dial to a pinned target, which
// fastdialer reports with the hostname in ASCII form, to the pinned ip that
// Check keys the target by. Addresses of other hosts are returned as is.
func pinnedAddress(input *contextargs.MetaInput, address string) string {
	host, port, err := net.SplitHostPort(address)
	if err != nil || !strings.EqualFold(asciiHost(host), asciiHost(input.Hostname())) {
		return address
	}
	return net.JoinHostPort(input.CustomIP, port)
}

// asciiHost returns the ASCII form of an internationalized hostname.
func asciiHost(host string) string {
	if ascii, err := idna.ToASCII(host); err == nil {
		return ascii
	}
	return host
}

var reCheckError = regexp.MustCompile(`(no address found for host|could not resolve host|connection refused|connection reset by peer|could not connect to any address found for host|timeout awaiting response headers|i/o timeout)`)

// checkError checks if an error represents a type that should be
// added to the host skipping table.
// it first parses error and extracts the cause and checks for blacklisted
// or common errors that should be skipped
func (c *Cache) checkError(protoType string, err error) bool {
	if err == nil {
		return false
	}
	if protoType != "http" {
		return false
	}
	kind := errkit.GetErrorKind(err, nucleierr.ErrTemplateLogic)
	switch kind {
	case nucleierr.ErrTemplateLogic:
		// these are errors that are not related to the target
		// and are due to template logic
		return false
	case errkit.ErrKindNetworkTemporary:
		// a single temporary error (timeout, i/o reset) is transient, but a host
		// that produces them on every request with no success in between is
		// unresponsive. Count it; MarkFailedOrRemove resets the host on the next
		// successful response, so only consecutive failures reach MaxHostError.
		return true
	case errkit.ErrKindNetworkPermanent:
		// these should be counted as host errors
		return true
	case errkit.ErrKindDeadline:
		// these should not be counted as host errors
		return false
	default:
		// parse error for further processing
		errX := errkit.FromError(err)
		tmp := errX.Cause()
		cause := tmp.Error()
		if stringsutil.ContainsAll(cause, "ReadStatusLine:", "read: connection reset by peer") {
			// this is a FP and should not be counted as a host error
			// because server closes connection when it reads corrupted bytes which we send via rawhttp
			return false
		}
		if strings.HasPrefix(cause, "ReadStatusLine:") {
			// error is present in last part when using rawhttp
			// this will be fixed once errkit is used everywhere
			lastIndex := strings.LastIndex(cause, ":")
			if lastIndex == -1 {
				lastIndex = 0
			}
			if lastIndex >= len(cause)-1 {
				lastIndex = 0
			}
			cause = cause[lastIndex+1:]
		}
		for _, msg := range c.TrackError {
			if strings.Contains(cause, msg) {
				return true
			}
		}
		return reCheckError.MatchString(cause)
	}
}