package sipguardian

import (
	"crypto/rand"
	"encoding/hex"
	"fmt"
	"net"
	"strconv"
	"strings"
	"time"

	"github.com/caddyserver/caddy/v2"
	"github.com/caddyserver/caddy/v2/caddyconfig/caddyfile"
	"github.com/mholt/caddy-l4/layer4"
	"go.uber.org/zap"
)

func init() {
	caddy.RegisterModule(TopologyHider{})
}

// TopologyHider is a Layer 4 handler that hides internal SIP topology
type TopologyHider struct {
	// Enabled toggles topology hiding
	Enabled bool `json:"enabled,omitempty"`

	// ProxyHost is the public IP address to use in rewritten headers
	ProxyHost string `json:"proxy_host,omitempty"`

	// ProxyPort is the public port to use in rewritten headers
	ProxyPort int `json:"proxy_port,omitempty"`

	// RewriteVia adds/modifies Via headers
	RewriteVia bool `json:"rewrite_via,omitempty"`

	// RewriteContact rewrites Contact headers to hide internal addresses
	RewriteContact bool `json:"rewrite_contact,omitempty"`

	// StripHeaders lists headers to remove (e.g., Server, P-Preferred-Identity)
	StripHeaders []string `json:"strip_headers,omitempty"`

	// AnonymizeCallID replaces Call-ID with proxy-generated value
	AnonymizeCallID bool `json:"anonymize_call_id,omitempty"`

	// HidePrivateIPs automatically detects and hides RFC 1918 addresses
	HidePrivateIPs bool `json:"hide_private_ips,omitempty"`

	// TransactionTimeout for cleanup of pending transactions
	TransactionTimeout caddy.Duration `json:"transaction_timeout,omitempty"`

	// Runtime state
	logger       *zap.Logger
	transactions *TransactionManager
	dialogs      *DialogManager
}

func (TopologyHider) CaddyModule() caddy.ModuleInfo {
	return caddy.ModuleInfo{
		ID:  "layer4.handlers.sip_topology_hider",
		New: func() caddy.Module { return new(TopologyHider) },
	}
}

func (h *TopologyHider) Provision(ctx caddy.Context) error {
	h.logger = ctx.Logger()

	// Set defaults
	if h.ProxyPort == 0 {
		h.ProxyPort = 5060
	}
	if h.TransactionTimeout == 0 {
		h.TransactionTimeout = caddy.Duration(32 * time.Second)
	}

	// Initialize transaction manager for response correlation
	h.transactions = NewTransactionManager(h.logger, time.Duration(h.TransactionTimeout))

	// Initialize dialog manager for stateful tracking
	h.dialogs = NewDialogManager(h.logger, 10*time.Minute)

	// Start cleanup goroutine
	go h.cleanupLoop(ctx)

	h.logger.Info("SIP Topology Hider initialized",
		zap.Bool("enabled", h.Enabled),
		zap.String("proxy_host", h.ProxyHost),
		zap.Int("proxy_port", h.ProxyPort),
		zap.Bool("rewrite_via", h.RewriteVia),
		zap.Bool("rewrite_contact", h.RewriteContact),
		zap.Strings("strip_headers", h.StripHeaders),
		zap.Bool("anonymize_call_id", h.AnonymizeCallID),
		zap.Bool("hide_private_ips", h.HidePrivateIPs),
	)

	return nil
}

// Handle processes the SIP message and applies topology hiding
func (h *TopologyHider) Handle(cx *layer4.Connection, next layer4.Handler) error {
	if !h.Enabled {
		return next.Handle(cx)
	}

	// Get client address
	remoteAddr := cx.RemoteAddr().String()
	clientHost, clientPortStr, err := net.SplitHostPort(remoteAddr)
	if err != nil {
		clientHost = remoteAddr
	}
	clientPort, _ := strconv.Atoi(clientPortStr)

	// Read the SIP message
	buf := make([]byte, 8192)
	n, err := cx.Read(buf)
	if err != nil || n == 0 {
		return next.Handle(cx)
	}
	buf = buf[:n]

	// Parse the SIP message
	msg, err := ParseSIPMessage(buf)
	if err != nil {
		h.logger.Debug("Failed to parse SIP message for topology hiding",
			zap.Error(err),
			zap.String("client", clientHost),
		)
		return next.Handle(cx)
	}

	// Apply topology hiding based on message type
	if msg.IsRequest {
		h.handleRequest(msg, clientHost, clientPort)
	} else {
		h.handleResponse(msg)
	}

	// Serialize modified message
	modifiedData := msg.Serialize()

	// Write modified data back to connection buffer
	// Note: This requires caddy-l4 support for modifying the connection data
	// For now, we'll pass through the next handler
	// In a full implementation, we'd use a wrapped connection

	h.logger.Debug("Topology hiding applied",
		zap.Bool("is_request", msg.IsRequest),
		zap.String("method", msg.Method),
		zap.Int("status_code", msg.StatusCode),
		zap.Int("original_size", n),
		zap.Int("modified_size", len(modifiedData)),
	)

	// Continue to next handler
	// TODO: Replace connection data with modified message
	return next.Handle(cx)
}

// handleRequest applies topology hiding to outgoing requests
func (h *TopologyHider) handleRequest(msg *SIPMessage, clientHost string, clientPort int) {
	callID := msg.GetCallID()
	fromTag := msg.GetFromTag()

	// Store original values for response correlation
	originalVia := ""
	if via := msg.GetHeader("Via"); via != nil {
		originalVia = via.Value
	}

	// Create transaction state
	branch := GenerateBranch()
	h.transactions.CreateTransaction(branch, msg.Method, clientHost, clientPort, originalVia)

	// For dialog-creating requests, create dialog state
	if msg.IsDialogCreating() {
		state := h.dialogs.CreateDialog(msg, clientHost, clientPort)
		if state != nil && originalVia != "" {
			h.dialogs.StoreOriginals(callID, fromTag, originalVia, "", "")
		}
	}

	// Add our Via header (prepend to top)
	if h.RewriteVia {
		newVia := h.buildViaHeader(branch)
		msg.PrependHeader("Via", newVia)

		h.logger.Debug("Added Via header",
			zap.String("via", newVia),
			zap.String("branch", branch),
		)
	}

	// Rewrite Contact header
	if h.RewriteContact {
		if contact := msg.GetHeader("Contact"); contact != nil {
			originalContact := contact.Value
			newContact := h.rewriteContactHeader(contact.Value)
			contact.Value = newContact

			// Store original for response rewriting
			if msg.IsDialogCreating() {
				h.dialogs.StoreOriginals(callID, fromTag, "", originalContact, "")
			}

			h.logger.Debug("Rewrote Contact header",
				zap.String("original", originalContact),
				zap.String("new", newContact),
			)
		}
	}

	// Anonymize Call-ID
	if h.AnonymizeCallID {
		if callIDHeader := msg.GetHeader("Call-ID"); callIDHeader != nil {
			originalCallID := callIDHeader.Value
			newCallID := h.generateAnonymousCallID()
			callIDHeader.Value = newCallID

			// Store mapping for response correlation
			h.dialogs.StoreOriginals(callID, fromTag, "", "", originalCallID)

			h.logger.Debug("Anonymized Call-ID",
				zap.String("original", originalCallID),
				zap.String("new", newCallID),
			)
		}
	}

	// Strip sensitive headers
	h.stripSensitiveHeaders(msg)

	// Hide private IPs in other headers
	if h.HidePrivateIPs {
		h.hidePrivateIPsInHeaders(msg)
	}
}

// handleResponse applies topology hiding to responses
func (h *TopologyHider) handleResponse(msg *SIPMessage) {
	// Find the transaction by branch
	branch := msg.GetViaBranch()
	transaction := h.transactions.GetTransaction(branch)

	if transaction == nil {
		h.logger.Debug("No transaction found for response",
			zap.String("branch", branch),
			zap.Int("status_code", msg.StatusCode),
		)
		return
	}

	// Remove our Via header (should be top Via)
	if h.RewriteVia {
		// The top Via should be ours - remove it
		topVia := msg.GetHeader("Via")
		if topVia != nil {
			viaHeader, _ := ParseViaHeader(topVia.Value)
			if viaHeader != nil && viaHeader.Branch == branch {
				msg.RemoveFirstHeader("Via")
				h.logger.Debug("Removed proxy Via header",
					zap.String("branch", branch),
				)
			}
		}
	}

	// Update dialog state based on response
	callID := msg.GetCallID()
	fromTag := msg.GetFromTag()
	toTag := msg.GetToTag()

	// Update dialog with To tag
	if toTag != "" {
		h.dialogs.UpdateDialog(callID, fromTag, toTag)
	}

	// Mark dialog as confirmed for 2xx responses to INVITE
	_, method := msg.GetCSeq()
	if msg.StatusCode >= 200 && msg.StatusCode < 300 && method == "INVITE" {
		h.dialogs.ConfirmDialog(callID, fromTag)
	}

	// Strip sensitive headers from response too
	h.stripSensitiveHeaders(msg)

	// Clean up transaction for final responses
	if msg.StatusCode >= 200 {
		h.transactions.RemoveTransaction(branch)
	}
}

// buildViaHeader creates a new Via header with proxy address
func (h *TopologyHider) buildViaHeader(branch string) string {
	via := &ViaHeader{
		Protocol:  "SIP/2.0",
		Transport: "UDP",
		Host:      h.ProxyHost,
		Port:      h.ProxyPort,
		Branch:    branch,
	}
	return via.Serialize()
}

// rewriteContactHeader replaces internal addresses with proxy address
func (h *TopologyHider) rewriteContactHeader(value string) string {
	contact, err := ParseContactHeader(value)
	if err != nil || contact.URI == "*" {
		return value
	}

	// Check if the Contact URI contains a private IP
	host := ExtractHostFromURI(contact.URI)
	if h.HidePrivateIPs && IsPrivateIP(host) {
		// Rewrite URI to use proxy address
		contact.URI = RewriteURIHost(contact.URI, h.ProxyHost, h.ProxyPort)
	} else if h.RewriteContact {
		// Always rewrite if RewriteContact is enabled
		contact.URI = RewriteURIHost(contact.URI, h.ProxyHost, h.ProxyPort)
	}

	return contact.Serialize()
}

// stripSensitiveHeaders removes headers that leak internal information
func (h *TopologyHider) stripSensitiveHeaders(msg *SIPMessage) {
	// Strip configured headers
	for _, header := range h.StripHeaders {
		msg.RemoveHeader(header)
	}

	// Always strip these if HidePrivateIPs is enabled
	if h.HidePrivateIPs {
		for _, header := range SensitiveHeaders {
			msg.RemoveHeader(header)
		}
	}
}

// hidePrivateIPsInHeaders scans all headers for private IPs
func (h *TopologyHider) hidePrivateIPsInHeaders(msg *SIPMessage) {
	// Headers that commonly contain IP addresses
	ipHeaders := []string{
		"Record-Route",
		"Route",
		"P-Visited-Network-ID",
	}

	for i := range msg.Headers {
		headerLower := strings.ToLower(msg.Headers[i].Name)

		for _, ipHeader := range ipHeaders {
			if headerLower == strings.ToLower(ipHeader) {
				// Check for private IPs in the value
				if containsPrivateIP(msg.Headers[i].Value) {
					// Rewrite or remove the header
					msg.Headers[i].Value = rewritePrivateIPs(msg.Headers[i].Value, h.ProxyHost)
				}
			}
		}
	}
}

// generateAnonymousCallID creates a new random Call-ID
func (h *TopologyHider) generateAnonymousCallID() string {
	bytes := make([]byte, 16)
	rand.Read(bytes)
	return hex.EncodeToString(bytes) + "@" + h.ProxyHost
}

// containsPrivateIP checks if a string contains any private IP addresses
func containsPrivateIP(value string) bool {
	// Simple pattern matching for common private IP formats
	privatePatterns := []string{
		"10.", "192.168.", "172.16.", "172.17.", "172.18.", "172.19.",
		"172.20.", "172.21.", "172.22.", "172.23.", "172.24.", "172.25.",
		"172.26.", "172.27.", "172.28.", "172.29.", "172.30.", "172.31.",
	}

	for _, pattern := range privatePatterns {
		if strings.Contains(value, pattern) {
			return true
		}
	}
	return false
}

// rewritePrivateIPs replaces private IP addresses in a header value
func rewritePrivateIPs(value, replacement string) string {
	// This is a simplified implementation
	// A full implementation would use regex to properly replace IPs in URIs

	// For Record-Route and Route, we may want to preserve the structure
	// but replace the host portion

	return value // TODO: Implement proper IP replacement
}

// cleanupLoop periodically cleans up expired state
func (h *TopologyHider) cleanupLoop(ctx caddy.Context) {
	ticker := time.NewTicker(30 * time.Second)
	defer ticker.Stop()

	for {
		select {
		case <-ctx.Done():
			return
		case <-ticker.C:
			transRemoved := h.transactions.Cleanup()
			dialogRemoved := h.dialogs.Cleanup()

			if transRemoved > 0 || dialogRemoved > 0 {
				h.logger.Debug("Topology hider cleanup",
					zap.Int("transactions_removed", transRemoved),
					zap.Int("dialogs_removed", dialogRemoved),
				)
			}
		}
	}
}

// GetStats returns statistics about the topology hider
func (h *TopologyHider) GetStats() map[string]interface{} {
	stats := map[string]interface{}{
		"enabled":          h.Enabled,
		"proxy_host":       h.ProxyHost,
		"proxy_port":       h.ProxyPort,
		"rewrite_via":      h.RewriteVia,
		"rewrite_contact":  h.RewriteContact,
		"anonymize_callid": h.AnonymizeCallID,
		"hide_private_ips": h.HidePrivateIPs,
	}

	if h.transactions != nil {
		for k, v := range h.transactions.GetStats() {
			stats["transactions_"+k] = v
		}
	}

	if h.dialogs != nil {
		for k, v := range h.dialogs.GetStats() {
			stats["dialogs_"+k] = v
		}
	}

	return stats
}

// UnmarshalCaddyfile implements caddyfile.Unmarshaler for TopologyHider.
// Usage in Caddyfile:
//
//	sip_topology_hider {
//	    enabled true
//	    proxy_host 203.0.113.1
//	    proxy_port 5060
//	    rewrite_via
//	    rewrite_contact
//	    strip_headers P-Preferred-Identity P-Asserted-Identity Server
//	    anonymize_call_id
//	    hide_private_ips
//	    transaction_timeout 32s
//	}
func (h *TopologyHider) UnmarshalCaddyfile(d *caddyfile.Dispenser) error {
	// Move past "sip_topology_hider" token
	d.Next()

	for nesting := d.Nesting(); d.NextBlock(nesting); {
		switch d.Val() {
		case "enabled":
			if !d.NextArg() {
				// No argument means enabled
				h.Enabled = true
			} else {
				val := d.Val()
				h.Enabled = val == "true" || val == "yes" || val == "on"
			}

		case "proxy_host":
			if !d.NextArg() {
				return d.ArgErr()
			}
			h.ProxyHost = d.Val()

		case "proxy_port":
			if !d.NextArg() {
				return d.ArgErr()
			}
			port, err := strconv.Atoi(d.Val())
			if err != nil {
				return d.Errf("invalid proxy_port: %v", err)
			}
			h.ProxyPort = port

		case "rewrite_via":
			h.RewriteVia = true

		case "rewrite_contact":
			h.RewriteContact = true

		case "strip_headers":
			h.StripHeaders = d.RemainingArgs()
			if len(h.StripHeaders) == 0 {
				return d.ArgErr()
			}

		case "anonymize_call_id":
			h.AnonymizeCallID = true

		case "hide_private_ips":
			h.HidePrivateIPs = true

		case "transaction_timeout":
			if !d.NextArg() {
				return d.ArgErr()
			}
			dur, err := caddy.ParseDuration(d.Val())
			if err != nil {
				return d.Errf("invalid transaction_timeout: %v", err)
			}
			h.TransactionTimeout = caddy.Duration(dur)

		default:
			return d.Errf("unknown sip_topology_hider directive: %s", d.Val())
		}
	}

	// Validate required fields
	if h.Enabled && h.ProxyHost == "" {
		return fmt.Errorf("proxy_host is required when topology hiding is enabled")
	}

	return nil
}

// Interface guards
var (
	_ layer4.NextHandler    = (*TopologyHider)(nil)
	_ caddy.Module          = (*TopologyHider)(nil)
	_ caddy.Provisioner     = (*TopologyHider)(nil)
	_ caddyfile.Unmarshaler = (*TopologyHider)(nil)
)