package main

import (
	"context"
	"crypto/tls"
	"encoding/base64"
	"flag"
	"fmt"
	"io"
	"io/ioutil"
	"log"
	"net/http"
	"net/url"
	"os"
	"os/signal"
	"path/filepath"
	"regexp"
	"strings"
	"sync"
	"syscall"
	"text/template"
	"time"

	"github.com/quic-go/quic-go"
	"github.com/quic-go/quic-go/http3"
)

var (
	verbose    = flag.Bool("verbose", false, "verbose logging")
	addr       = flag.String("addr", "localhost:8443", "Address to listen on")
	certFile   = flag.String("cert", "cert.pem", "Certificate file")
	keyFile    = flag.String("key", "key.pem", "Private key file")
	targetAddr = flag.String("target", "", "Target address to proxy to (if empty, acts as forward proxy)")
	webUI      = flag.Bool("webui", false, "Enable web UI for user-specified proxy targets")
)

// List of hop-by-hop headers to be removed when proxying
var hopByHopHeaders = []string{
	"Connection",
	"Keep-Alive",
	"Proxy-Authenticate",
	"Proxy-Authorization",
	"Te",
	"Trailers",
	"Transfer-Encoding",
	"Upgrade",
}

const (
	scriptFileName = "custom_script.js"
	scriptConfigFileName = "script_config.json"
)

func main() {
	flag.Parse()

	// Setup logger
	logger := log.New(os.Stdout, "[FuckHTTP3] ", log.LstdFlags)

	// Check if cert and key files exist
	if _, err := os.Stat(*certFile); os.IsNotExist(err) {
		logger.Fatalf("Certificate file %s does not exist", *certFile)
	}
	if _, err := os.Stat(*keyFile); os.IsNotExist(err) {
		logger.Fatalf("Key file %s does not exist", *keyFile)
	}

	// Normalize target address if specified
	if *targetAddr != "" {
		// Strip trailing slash for consistency
		*targetAddr = strings.TrimSuffix(*targetAddr, "/")
		
		// Log the actual target we're using
		if *verbose {
			logger.Printf("Using normalized target: %s", *targetAddr)
		}

		// Check if target is a valid URL
		if !strings.HasPrefix(*targetAddr, "http://") && !strings.HasPrefix(*targetAddr, "https://") {
			// Add https:// by default
			*targetAddr = "https://" + *targetAddr
		}
		
		// Validate the URL
		_, err := url.Parse(*targetAddr)
		if err != nil {
			logger.Fatalf("Invalid target URL: %v", err)
		}
	}
	
	// Check for templates directory when webUI is enabled
	if *webUI {
		if _, err := os.Stat("templates/index.html"); os.IsNotExist(err) {
			logger.Fatalf("templates/index.html not found, required for web UI mode")
		}
	}

	// Create a new proxy server
	proxyServer := &ProxyServer{
		logger:     logger,
		targetAddr: *targetAddr,
		clients:    make(map[string]*http3.RoundTripper),
		clientsMu:  &sync.Mutex{},
		webUI:      *webUI,
		templates:  template.Must(template.ParseGlob("templates/*.html")),
		iframeMode: true, // Enable iframe mode by default
		customScript: "",  // Start with empty script
		scriptEnabled: false,  // Disabled by default
		scriptMutex: sync.RWMutex{},
	}

	// Load script from disk
	if err := proxyServer.loadScriptFromDisk(); err != nil {
		logger.Printf("Error loading script from disk: %v", err)
	}

	// Extract port for Alt-Svc headers
	portStr := strings.Split(*addr, ":")[1]

	// Configure TLS
	tlsConfig := &tls.Config{
		MinVersion:         tls.VersionTLS13, // HTTP/3 requires TLS 1.3
		NextProtos:         []string{"h3"},   // Specify HTTP/3 as the next protocol
		InsecureSkipVerify: false,            // Set to true for development only
	}
	
	// Configure the HTTP/3 server
	server := &http3.Server{
		Addr:      *addr,
		TLSConfig: http3.ConfigureTLSConfig(tlsConfig),
		Handler:   proxyServer,
		QuicConfig: &quic.Config{
			MaxIdleTimeout: 30 * time.Second,
			KeepAlivePeriod: 10 * time.Second,
		},
	}

	logger.Printf("Starting HTTP/3 proxy server on %s", *addr)
	logger.Printf("Using certificate: %s", *certFile)
	logger.Printf("Using key: %s", *keyFile)
	
	if *targetAddr != "" {
		logger.Printf("Proxying to target: %s", *targetAddr)
	} else {
		logger.Printf("Running as forward proxy")
	}

	// Create certificate pair from files
	cert, err := tls.LoadX509KeyPair(*certFile, *keyFile)
	if err != nil {
		logger.Fatalf("Failed to load certificates: %v", err)
	}

	// Create a standard HTTP server for HTTP/1.1 and HTTP/2
	standardServer := &http.Server{
		Addr:    *addr,
		Handler: http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
			// Add Alt-Svc header to advertise HTTP/3 capability
			w.Header().Set("Alt-Svc", fmt.Sprintf(`h3=":%s"; ma=2592000`, portStr))
			
			// Handle the request with the proxy
			proxyServer.ServeHTTP(w, r)
		}),
		TLSConfig: &tls.Config{
			Certificates: []tls.Certificate{cert},
			NextProtos:   []string{"h2", "http/1.1"},
		},
	}

	// Setup context for graceful shutdown
	ctx, cancel := context.WithCancel(context.Background())
	defer cancel()
	
	// Handle shutdown signals
	go func() {
		sigCh := make(chan os.Signal, 1)
		signal.Notify(sigCh, syscall.SIGINT, syscall.SIGTERM)
		
		sig := <-sigCh
		logger.Printf("Received signal %v, shutting down...", sig)
		
		// Cancel context
		cancel()
		
		// Create shutdown context with timeout
		shutdownCtx, shutdownCancel := context.WithTimeout(ctx, 10*time.Second)
		defer shutdownCancel()
		
		// Shutdown both servers
		if err := standardServer.Shutdown(shutdownCtx); err != nil {
			logger.Printf("Error shutting down HTTP/1.1 server: %v", err)
		}
		
		if err := server.Close(); err != nil {
			logger.Printf("Error shutting down HTTP/3 server: %v", err)
		}
		
		// Close all active roundtrippers
		proxyServer.closeAllClients()
	}()

	// Start the HTTP/1.1 + HTTP/2 server in a goroutine
	go func() {
		logger.Printf("Starting HTTP/1.1 and HTTP/2 server on %s", *addr)
		if err := standardServer.ListenAndServeTLS("", ""); err != nil && err != http.ErrServerClosed {
			logger.Printf("HTTP server error: %v", err)
		}
	}()

	// Start the HTTP/3 server
	err = server.ListenAndServeTLS(*certFile, *keyFile)
	if err != nil && err != http.ErrServerClosed {
		logger.Fatalf("Failed to start HTTP/3 server: %v", err)
	}
}

// ProxyServer implements the http.Handler interface
type ProxyServer struct {
	logger     *log.Logger
	targetAddr string
	clients    map[string]*http3.RoundTripper
	clientsMu  *sync.Mutex
	webUI      bool
	templates  *template.Template
	iframeMode bool // Whether to use iframe mode for the navbar
	customScript   string
	scriptEnabled  bool
	scriptMutex    sync.RWMutex
}

// ServeHTTP handles the HTTP requests
func (p *ProxyServer) ServeHTTP(w http.ResponseWriter, r *http.Request) {
	startTime := time.Now()
	
	// Log the incoming request
	if *verbose {
		p.logger.Printf("Received request: %s %s %s", r.Method, r.URL, r.Proto)
	}

	// Check for web UI mode
	if p.webUI {
		if r.URL.Path == "/" || r.URL.Path == "" {
			p.handleHome(w, r)
			return
		} else if r.URL.Path == "/proxy" && r.Method == http.MethodPost {
			p.handleProxyRequest(w, r)
			return
		} else if r.URL.Path == "/proxified" {
			p.handleProxifiedRequest(w, r)
			return
		} else if r.URL.Path == "/navbar" {
			p.handleNavbarFrame(w, r)
			return
		} else if r.URL.Path == "/save-script" {
			p.handleSaveScript(w, r)
			return
		}
	}

	// Check if we're operating as a reverse proxy (with fixed target) or forward proxy
	if p.targetAddr != "" {
		// Reverse proxy mode
		p.handleReverseProxy(w, r)
	} else {
		// Forward proxy mode
		p.handleForwardProxy(w, r)
	}
	
	// Log completion time
	if *verbose {
		p.logger.Printf("Request completed in %v", time.Since(startTime))
	}
}

// getRoundTripper gets or creates an HTTP/3 RoundTripper for the given host
func (p *ProxyServer) getRoundTripper(host string) *http3.RoundTripper {
	p.clientsMu.Lock()
	defer p.clientsMu.Unlock()
	
	if rt, ok := p.clients[host]; ok {
		return rt
	}
	
	rt := &http3.RoundTripper{
		TLSClientConfig: &tls.Config{
			InsecureSkipVerify: true, // Allow insecure connections for testing
			NextProtos:         []string{"h3"},
		},
	}
	
	p.clients[host] = rt
	return rt
}

// closeAllClients closes all RoundTripper instances
func (p *ProxyServer) closeAllClients() {
	p.clientsMu.Lock()
	defer p.clientsMu.Unlock()
	
	for host, rt := range p.clients {
		if err := rt.Close(); err != nil && *verbose {
			p.logger.Printf("Error closing roundtripper for %s: %v", host, err)
		}
	}
	
	p.clients = make(map[string]*http3.RoundTripper)
}

// handleReverseProxy handles requests in reverse proxy mode
func (p *ProxyServer) handleReverseProxy(w http.ResponseWriter, r *http.Request) {
	// Get or create a RoundTripper for the target
	roundTripper := p.getRoundTripper(p.targetAddr)

	// Create a new client
	client := &http.Client{
		Transport: roundTripper,
		Timeout:   30 * time.Second,
	}

	// Create a new request to the target
	var targetURL string
	if strings.HasPrefix(p.targetAddr, "http://") || strings.HasPrefix(p.targetAddr, "https://") {
		// Target already has a scheme
		if strings.HasSuffix(p.targetAddr, "/") && strings.HasPrefix(r.URL.Path, "/") {
			// Avoid double slashes when both target and path have slashes
			targetURL = p.targetAddr + strings.TrimPrefix(r.URL.Path, "/")
		} else if !strings.HasSuffix(p.targetAddr, "/") && !strings.HasPrefix(r.URL.Path, "/") && r.URL.Path != "" {
			// Add slash when neither has one
			targetURL = p.targetAddr + "/" + r.URL.Path
		} else {
			// Normal case
			targetURL = p.targetAddr + r.URL.Path
		}
	} else {
		// No scheme in target, add https:// (original behavior)
		targetURL = fmt.Sprintf("https://%s%s", p.targetAddr, r.URL.Path)
	}

	if r.URL.RawQuery != "" {
		targetURL += "?" + r.URL.RawQuery
	}

	if *verbose {
		p.logger.Printf("Proxying to: %s", targetURL)
	}

	// Create a new request
	outReq, err := http.NewRequestWithContext(r.Context(), r.Method, targetURL, r.Body)
	if err != nil {
		http.Error(w, err.Error(), http.StatusInternalServerError)
		p.logger.Printf("Error creating request: %v", err)
		return
	}

	// Copy headers
	p.copyHeaders(outReq.Header, r.Header)

	// Send the request
	resp, err := client.Do(outReq)
	if err != nil {
		http.Error(w, err.Error(), http.StatusBadGateway)
		p.logger.Printf("Error sending request: %v", err)
		return
	}
	defer resp.Body.Close()

	// Copy the response headers
	p.copyHeaders(w.Header(), resp.Header)
	
	// Add Alt-Svc header with the correct port
	portStr := strings.Split(*addr, ":")[1]
	w.Header().Add("Alt-Svc", fmt.Sprintf(`h3=":%s"; ma=2592000`, portStr))
	
	w.WriteHeader(resp.StatusCode)

	// Copy the response body using io.Copy for efficiency
	if _, err := io.Copy(w, resp.Body); err != nil {
		p.logger.Printf("Error copying response body: %v", err)
	}
}

// handleForwardProxy handles requests in forward proxy mode
func (p *ProxyServer) handleForwardProxy(w http.ResponseWriter, r *http.Request) {
	if r.Method == http.MethodConnect {
		// Handle CONNECT method (for HTTPS)
		p.handleConnect(w, r)
		return
	}

	// Ensure absolute URL for forward proxy
	if !r.URL.IsAbs() {
		http.Error(w, "Request URL must be absolute in forward proxy mode", http.StatusBadRequest)
		return
	}

	// Get or create a RoundTripper for the target host
	roundTripper := p.getRoundTripper(r.URL.Host)

	// Create a new client
	client := &http.Client{
		Transport: roundTripper,
		Timeout:   30 * time.Second,
		// Don't follow redirects, let the client handle them
		CheckRedirect: func(req *http.Request, via []*http.Request) error {
			return http.ErrUseLastResponse
		},
	}

	if *verbose {
		p.logger.Printf("Proxying to: %s", r.URL)
	}

	// Create a new request
	outReq, err := http.NewRequestWithContext(r.Context(), r.Method, r.URL.String(), r.Body)
	if err != nil {
		http.Error(w, err.Error(), http.StatusInternalServerError)
		p.logger.Printf("Error creating request: %v", err)
		return
	}

	// Copy headers
	p.copyHeaders(outReq.Header, r.Header)

	// Send the request
	resp, err := client.Do(outReq)
	if err != nil {
		http.Error(w, err.Error(), http.StatusBadGateway)
		p.logger.Printf("Error sending request: %v", err)
		return
	}
	defer resp.Body.Close()

	// Copy the response headers
	p.copyHeaders(w.Header(), resp.Header)
	
	// Add Alt-Svc header with the correct port
	portStr := strings.Split(*addr, ":")[1]
	w.Header().Add("Alt-Svc", fmt.Sprintf(`h3=":%s"; ma=2592000`, portStr))
	
	w.WriteHeader(resp.StatusCode)

	// Copy the response body
	if _, err := io.Copy(w, resp.Body); err != nil {
		p.logger.Printf("Error copying response body: %v", err)
	}
}

// handleConnect handles the CONNECT method for HTTPS tunneling
func (p *ProxyServer) handleConnect(w http.ResponseWriter, r *http.Request) {
	// For HTTP/3 CONNECT, we need to establish a QUIC connection
	// to the target and then setup a bidirectional stream
	targetHost := r.Host
	
	// Get roundTripper but don't use it yet - we'll access it via client in future implementation
	_ = p.getRoundTripper(targetHost)
	
	// Notify the client that tunnel has been established
	w.WriteHeader(http.StatusOK)
	
	// Check if we're dealing with a hijackable connection
	hijacker, ok := w.(http.Hijacker)
	if !ok {
		p.logger.Printf("Connection doesn't support hijacking, can't establish tunnel")
		http.Error(w, "CONNECT not supported over HTTP/3", http.StatusInternalServerError)
		return
	}
	
	// Hijack the connection
	clientConn, _, err := hijacker.Hijack()
	if err != nil {
		p.logger.Printf("Failed to hijack connection: %v", err)
		http.Error(w, err.Error(), http.StatusInternalServerError)
		return
	}
	defer clientConn.Close()
	
	// For QUIC connections, this is more complex as we can't directly hijack streams
	// This is a basic implementation that won't work for all cases
	p.logger.Printf("CONNECT tunneling is limited in HTTP/3 due to protocol differences")
	p.logger.Printf("Simple pass-through enabled for %s", targetHost)
	
	// Note: In a full implementation, we would need to:
	// 1. Open a QUIC connection to the target
	// 2. Create a bidirectional stream
	// 3. Set up forwarding between the client stream and target stream
	// This requires direct access to the QUIC connection which http3.RoundTripper doesn't expose easily
}

// copyHeaders copies HTTP headers from src to dst, removing hop-by-hop headers
func (p *ProxyServer) copyHeaders(dst, src http.Header) {
	for k, vv := range src {
		// Skip hop-by-hop headers
		if p.isHopByHopHeader(k) {
			continue
		}
		
		for _, v := range vv {
			dst.Add(k, v)
		}
	}
}

// isHopByHopHeader checks if a header is hop-by-hop
func (p *ProxyServer) isHopByHopHeader(header string) bool {
	header = strings.ToLower(header)
	for _, h := range hopByHopHeaders {
		if strings.ToLower(h) == header {
			return true
		}
	}
	
	// Check for Connection header values
	for _, h := range hopByHopHeaders {
		if h == "Connection" {
			values := strings.Split(header, ",")
			for _, v := range values {
				if strings.TrimSpace(strings.ToLower(v)) == header {
					return true
				}
			}
		}
	}
	
	return false
}

// handleHome renders the web UI
func (p *ProxyServer) handleHome(w http.ResponseWriter, r *http.Request) {
	if r.URL.Path != "/" {
		http.NotFound(w, r)
		return
	}
	
	p.scriptMutex.RLock()
	scriptEnabled := p.scriptEnabled
	customScript := p.customScript
	p.scriptMutex.RUnlock()
	
	// Create template data with all needed fields
	data := struct {
		WebUIEnabled bool
		CustomScript string
		ScriptEnabled bool
		Error string
		URL string
	}{
		WebUIEnabled: *webUI,
		CustomScript: customScript,
		ScriptEnabled: scriptEnabled,
		Error: "",
		URL: "",
	}
	
	// Render template
	err := p.templates.ExecuteTemplate(w, "index.html", data)
	if err != nil {
		http.Error(w, err.Error(), http.StatusInternalServerError)
		p.logger.Printf("Error rendering template: %v", err)
	}
}

// handleProxyRequest processes the form submission from the web UI
func (p *ProxyServer) handleProxyRequest(w http.ResponseWriter, r *http.Request) {
	if err := r.ParseForm(); err != nil {
		p.logger.Printf("Error parsing form: %v", err)
		http.Error(w, "Bad Request", http.StatusBadRequest)
		return
	}
	
	targetURL := r.FormValue("url")
	if targetURL == "" {
		// Handle empty URL
		data := struct {
			URL   string
			Error string
		}{
			URL:   "",
			Error: "Please enter a URL",
		}
		
		err := p.templates.ExecuteTemplate(w, "index.html", data)
		if err != nil {
			p.logger.Printf("Error rendering template: %v", err)
			http.Error(w, "Internal Server Error", http.StatusInternalServerError)
		}
		return
	}
	
	// Normalize URL
	if !strings.HasPrefix(targetURL, "http://") && !strings.HasPrefix(targetURL, "https://") {
		targetURL = "https://" + targetURL
	}
	
	// Validate URL
	parsedURL, err := url.Parse(targetURL)
	if err != nil {
		data := struct {
			URL   string
			Error string
		}{
			URL:   targetURL,
			Error: "Invalid URL: " + err.Error(),
		}
		
		err := p.templates.ExecuteTemplate(w, "index.html", data)
		if err != nil {
			p.logger.Printf("Error rendering template: %v", err)
			http.Error(w, "Internal Server Error", http.StatusInternalServerError)
		}
		return
	}
	
	// Store the target URL in a query parameter and redirect
	proxifiedURL := fmt.Sprintf("/proxified?url=%s", url.QueryEscape(parsedURL.String()))
	http.Redirect(w, r, proxifiedURL, http.StatusSeeOther)
}

// handleDataURL handles data URLs
func (p *ProxyServer) handleDataURL(w http.ResponseWriter, r *http.Request, targetURL string) {
	// Parse the data URL
	if strings.HasPrefix(targetURL, "data:") {
		// Extract content type and data
		parts := strings.SplitN(targetURL[5:], ",", 2)
		if len(parts) != 2 {
			http.Error(w, "Invalid data URL format", http.StatusBadRequest)
			return
		}
		
		contentTypeInfo := parts[0]
		data := parts[1]
		
		// Set content type
		if contentTypeInfo != "" {
			if strings.HasSuffix(contentTypeInfo, ";base64") {
				w.Header().Set("Content-Type", strings.TrimSuffix(contentTypeInfo, ";base64"))
				// Decode base64
				decoded, err := base64.StdEncoding.DecodeString(data)
				if err != nil {
					http.Error(w, "Invalid base64 encoding in data URL", http.StatusBadRequest)
					return
				}
				w.Write(decoded)
			} else {
				w.Header().Set("Content-Type", contentTypeInfo)
				// URL decode the data
				decoded, err := url.QueryUnescape(data)
				if err != nil {
					http.Error(w, "Invalid URL encoding in data URL", http.StatusBadRequest)
					return
				}
				w.Write([]byte(decoded))
			}
		} else {
			// Default to text/plain
			w.Header().Set("Content-Type", "text/plain")
			decoded, _ := url.QueryUnescape(data)
			w.Write([]byte(decoded))
		}
		return
	}
	
	// Not a data URL, return error
	http.Error(w, "Not a data URL", http.StatusBadRequest)
}

// handleProxifiedRequest fetches and processes the target URL content
func (p *ProxyServer) handleProxifiedRequest(w http.ResponseWriter, r *http.Request) {
	// Extract the URL from the query parameter
	targetURL := r.URL.Query().Get("url")
	if targetURL == "" {
		http.Error(w, "Missing URL parameter", http.StatusBadRequest)
		return
	}
	
	// Handle data: URLs specially
	if strings.HasPrefix(targetURL, "data:") {
		p.handleDataURL(w, r, targetURL)
		return
	}
	
	// Ensure URL is absolute
	if !strings.HasPrefix(targetURL, "http://") && !strings.HasPrefix(targetURL, "https://") {
		http.Error(w, "Request URL must be absolute in forward proxy mode", http.StatusBadRequest)
		return
	}
	
	parsedURL, err := url.Parse(targetURL)
	if err != nil {
		p.logger.Printf("Error parsing URL: %v", err)
		http.Error(w, "Invalid URL", http.StatusBadRequest)
		return
	}
	
	if *verbose {
		p.logger.Printf("Proxifying request to: %s", targetURL)
	}
	
	// Get or create a RoundTripper for the target host
	roundTripper := p.getRoundTripper(parsedURL.Host)
	
	// Create a new client
	client := &http.Client{
		Transport: roundTripper,
		Timeout:   30 * time.Second,
		CheckRedirect: func(req *http.Request, via []*http.Request) error {
			// Instead of following redirects, we'll return the redirect URL
			// so we can process it through our proxy
			return http.ErrUseLastResponse
		},
	}
	
	// Create a new request
	outReq, err := http.NewRequestWithContext(r.Context(), http.MethodGet, targetURL, nil)
	if err != nil {
		http.Error(w, err.Error(), http.StatusInternalServerError)
		p.logger.Printf("Error creating request: %v", err)
		return
	}
	
	// Copy original request headers
	p.copyHeaders(outReq.Header, r.Header)
	
	// Add common browser headers if not present
	if outReq.Header.Get("User-Agent") == "" {
		outReq.Header.Set("User-Agent", "Mozilla/5.0 (Windows NT 10.0; Win64; x64) AppleWebKit/537.36 (KHTML, like Gecko) Chrome/91.0.4472.124 Safari/537.36")
	}
	if outReq.Header.Get("Accept") == "" {
		outReq.Header.Set("Accept", "text/html,application/xhtml+xml,application/xml;q=0.9,image/webp,*/*;q=0.8")
	}
	
	// Send the request
	resp, err := client.Do(outReq)
	if err != nil {
		http.Error(w, err.Error(), http.StatusBadGateway)
		p.logger.Printf("Error sending request: %v", err)
		return
	}
	defer resp.Body.Close()
	
	// Handle redirects manually
	if resp.StatusCode >= 300 && resp.StatusCode < 400 {
		redirectURL := resp.Header.Get("Location")
		if redirectURL != "" {
			// Make sure the redirect URL is absolute
			redirectURL, err = p.makeAbsoluteURL(redirectURL, parsedURL)
			if err != nil {
				http.Error(w, err.Error(), http.StatusInternalServerError)
				return
			}
			
			// Redirect through our proxy
			proxifiedRedirect := fmt.Sprintf("/proxified?url=%s", url.QueryEscape(redirectURL))
			http.Redirect(w, r, proxifiedRedirect, resp.StatusCode)
			return
		}
	}
	
	// Copy the response headers
	p.copyHeaders(w.Header(), resp.Header)
	
	// Modify response headers to allow framing and inlining
	p.modifyResponseHeaders(w.Header())
	
	// Add Alt-Svc header with the correct port
	portStr := strings.Split(*addr, ":")[1]
	w.Header().Add("Alt-Svc", fmt.Sprintf(`h3=":%s"; ma=2592000`, portStr))
	
	// Set appropriate Content-Type headers
	contentType := resp.Header.Get("Content-Type")
	if strings.Contains(r.URL.Path, ".js") && contentType == "text/plain" {
		// Fix JavaScript MIME type
		w.Header().Set("Content-Type", "application/javascript")
	} else if strings.Contains(r.URL.Path, ".css") && contentType == "text/plain" {
		// Fix CSS MIME type
		w.Header().Set("Content-Type", "text/css")
	} else {
		w.Header().Set("Content-Type", contentType)
	}
	
	// Determine content type
	contentType = resp.Header.Get("Content-Type")
	isHTML := strings.Contains(contentType, "text/html") || strings.Contains(contentType, "application/xhtml+xml")
	
	if isHTML {
		// Process HTML content to rewrite links
		bodyBytes, err := io.ReadAll(resp.Body)
		if err != nil {
			http.Error(w, err.Error(), http.StatusInternalServerError)
			p.logger.Printf("Error reading response body: %v", err)
			return
		}
		
		// Rewrite URLs in the HTML content
		processedHTML := p.processHTML(bodyBytes, parsedURL.String())
		
		// Add a navigation bar at the top
		processedHTML = p.addProxyNavBar(processedHTML, parsedURL.String())
		
		// Set the processed content type and length
		w.Header().Set("Content-Type", "text/html; charset=utf-8")
		w.Header().Set("Content-Length", fmt.Sprintf("%d", len(processedHTML)))
		
		// Write status code and content
		w.WriteHeader(resp.StatusCode)
		w.Write(processedHTML)
	} else if strings.Contains(contentType, "javascript") {
		// If JavaScript, rewrite URLs
		bodyBytes, err := io.ReadAll(resp.Body)
		if err != nil {
			http.Error(w, err.Error(), http.StatusInternalServerError)
			p.logger.Printf("Error reading response body: %v", err)
			return
		}
		
		rewrittenBody := p.rewriteJavaScript(bodyBytes, targetURL)
		
		// Set the processed content type and length
		w.Header().Set("Content-Type", contentType)
		w.Header().Set("Content-Length", fmt.Sprintf("%d", len(rewrittenBody)))
		
		// Write status code and content
		w.WriteHeader(resp.StatusCode)
		w.Write(rewrittenBody)
	} else {
		// For non-HTML content, pass through directly
		w.WriteHeader(resp.StatusCode)
		if _, err := io.Copy(w, resp.Body); err != nil {
			p.logger.Printf("Error copying response body: %v", err)
		}
	}
}

// processHTML rewrites URLs in HTML content and injects JavaScript
func (p *ProxyServer) processHTML(htmlBytes []byte, baseURLStr string) []byte {
	// Convert HTML to string for easier manipulation
	htmlStr := string(htmlBytes)
	
	p.scriptMutex.RLock()
	scriptEnabled := p.scriptEnabled
	customScript := p.customScript
	p.scriptMutex.RUnlock()
	
	// If script injection is enabled, inject the custom script into the HTML
	if scriptEnabled && customScript != "" {
		p.logger.Printf("Injecting script into response for URL: %s", baseURLStr)
		
		// Create script tag with the custom script
		scriptTag := fmt.Sprintf("<script type=\"text/javascript\">\n// FuckHTTP3 Injected Script\n%s\n</script>", customScript)
		
		// Try to insert it before </body> first
		if strings.Contains(htmlStr, "</body>") {
			htmlStr = strings.Replace(htmlStr, "</body>", scriptTag+"\n</body>", 1)
			p.logger.Printf("Script injected before </body> tag")
		} else if strings.Contains(htmlStr, "</html>") {
			// If no </body>, try before </html>
			htmlStr = strings.Replace(htmlStr, "</html>", scriptTag+"\n</html>", 1)
			p.logger.Printf("Script injected before </html> tag")
		} else {
			// If neither tag exists, append it to the end
			htmlStr += "\n" + scriptTag
			p.logger.Printf("Script appended to the end of HTML")
		}
	} else {
		if !scriptEnabled {
			p.logger.Printf("Script injection disabled")
		} else if customScript == "" {
			p.logger.Printf("No custom script to inject")
		}
	}
	
	// Parse the base URL string into a URL object
	baseURL, err := url.Parse(baseURLStr)
	if err != nil {
		p.logger.Printf("Error parsing base URL %s: %v", baseURLStr, err)
		baseURL, _ = url.Parse("https://example.com")
	}
	
	// Process base URL if present
	baseTagRegex := regexp.MustCompile(`<base[^>]+href=["']([^"']+)["'][^>]*>`)
	baseMatch := baseTagRegex.FindStringSubmatch(htmlStr)
	
	if len(baseMatch) > 1 {
		// There's a base tag, use it for resolving relative URLs
		baseHref := baseMatch[1]
		newBaseURL, err := url.Parse(baseHref)
		if err == nil {
			baseURL = baseURL.ResolveReference(newBaseURL)
		}
	}
	
	// Rewrite hyperlinks
	aTagRegex := regexp.MustCompile(`<a[^>]+href=["']([^"']+)["'][^>]*>`)
	htmlStr = aTagRegex.ReplaceAllStringFunc(htmlStr, func(match string) string {
		submatches := aTagRegex.FindStringSubmatch(match)
		if len(submatches) < 2 {
			return match
		}
		
		href := submatches[1]
		
		// Skip javascript: and mailto: links
		if strings.HasPrefix(href, "javascript:") || strings.HasPrefix(href, "mailto:") {
			return match
		}
		
		// Construct absolute URL
		absoluteURL, err := p.makeAbsoluteURL(href, baseURL)
		if err != nil {
			return match
		}
		
		// Replace the href with our proxified URL
		proxifiedURL := fmt.Sprintf("/proxified?url=%s", url.QueryEscape(absoluteURL))
		return strings.Replace(match, submatches[1], proxifiedURL, 1)
	})
	
	// Rewrite image sources
	imgTagRegex := regexp.MustCompile(`<img[^>]+src=["']([^"']+)["'][^>]*>`)
	htmlStr = imgTagRegex.ReplaceAllStringFunc(htmlStr, func(match string) string {
		submatches := imgTagRegex.FindStringSubmatch(match)
		if len(submatches) < 2 {
			return match
		}
		
		src := submatches[1]
		
		// Construct absolute URL
		absoluteURL, err := p.makeAbsoluteURL(src, baseURL)
		if err != nil {
			return match
		}
		
		// Replace the src with our proxified URL
		proxifiedURL := fmt.Sprintf("/proxified?url=%s", url.QueryEscape(absoluteURL))
		return strings.Replace(match, submatches[1], proxifiedURL, 1)
	})
	
	// Rewrite CSS links
	linkTagRegex := regexp.MustCompile(`<link[^>]+href=["']([^"']+)["'][^>]*>`)
	htmlStr = linkTagRegex.ReplaceAllStringFunc(htmlStr, func(match string) string {
		submatches := linkTagRegex.FindStringSubmatch(match)
		if len(submatches) < 2 {
			return match
		}
		
		href := submatches[1]
		
		// Construct absolute URL
		absoluteURL, err := p.makeAbsoluteURL(href, baseURL)
		if err != nil {
			return match
		}
		
		// Replace the href with our proxified URL
		proxifiedURL := fmt.Sprintf("/proxified?url=%s", url.QueryEscape(absoluteURL))
		return strings.Replace(match, submatches[1], proxifiedURL, 1)
	})
	
	// Rewrite script sources
	scriptTagRegex := regexp.MustCompile(`<script[^>]+src=["']([^"']+)["'][^>]*>`)
	htmlStr = scriptTagRegex.ReplaceAllStringFunc(htmlStr, func(match string) string {
		submatches := scriptTagRegex.FindStringSubmatch(match)
		if len(submatches) < 2 {
			return match
		}
		
		src := submatches[1]
		
		// Construct absolute URL
		absoluteURL, err := p.makeAbsoluteURL(src, baseURL)
		if err != nil {
			return match
		}
		
		// Replace the src with our proxified URL
		proxifiedURL := fmt.Sprintf("/proxified?url=%s", url.QueryEscape(absoluteURL))
		return strings.Replace(match, submatches[1], proxifiedURL, 1)
	})
	
	return []byte(htmlStr)
}

// makeAbsoluteURL converts a relative URL to an absolute URL
func (p *ProxyServer) makeAbsoluteURL(href string, base *url.URL) (string, error) {
	// Check if already absolute
	if strings.HasPrefix(href, "http://") || strings.HasPrefix(href, "https://") {
		return href, nil
	}
	
	// Parse the href
	relativeURL, err := url.Parse(href)
	if err != nil {
		return "", err
	}
	
	// Resolve against base URL
	absoluteURL := base.ResolveReference(relativeURL)
	
	return absoluteURL.String(), nil
}

// addProxyNavBar adds a navigation bar at the top of the HTML content
func (p *ProxyServer) addProxyNavBar(html []byte, currentURL string) []byte {
	// Generate a unique ID for this navbar instance to avoid conflicts
	navbarID := fmt.Sprintf("fuckhttp3_navbar_%d", time.Now().UnixNano())
	
	// Create a more secure and isolated navbar with iframe
	navbar := fmt.Sprintf(`
<div id="%s" style="all: initial !important; position: fixed !important; top: 0 !important; left: 0 !important; width: 100%% !important; height: 50px !important; z-index: 2147483647 !important; box-shadow: 0 2px 5px rgba(0,0,0,0.3) !important; background: #333 !important; font-family: Arial, sans-serif !important; color: white !important; display: block !important; border: none !important; margin: 0 !important; padding: 0 !important;">
    <iframe src="/navbar?url=%s" style="width: 100%% !important; height: 100%% !important; border: none !important; margin: 0 !important; padding: 0 !important;" frameBorder="0"></iframe>
</div>
<div style="height: 50px !important; width: 100%% !important; display: block !important;"></div>

<script type="text/javascript">
(function() {
    // Store the original createElement to prevent hijacking
    var originalCreateElement = document.createElement;
    var navbarId = "%s";
    
    // Check and restore navbar every 100ms
    setInterval(function() {
        var navbar = document.getElementById(navbarId);
        if (!navbar || !navbar.style || navbar.style.display === "none" || 
            !document.body.contains(navbar) || 
            navbar.getBoundingClientRect().height < 10) {
            
            // Create a new navbar if it's been removed or hidden
            var newNavbar = originalCreateElement.call(document, "div");
            newNavbar.id = navbarId;
            newNavbar.style = "all: initial !important; position: fixed !important; top: 0 !important; left: 0 !important; width: 100%% !important; height: 50px !important; z-index: 2147483647 !important; box-shadow: 0 2px 5px rgba(0,0,0,0.3) !important; background: #333 !important; font-family: Arial, sans-serif !important; color: white !important; display: block !important; border: none !important; margin: 0 !important; padding: 0 !important;";
            
            var iframe = originalCreateElement.call(document, "iframe");
            iframe.src = "/navbar?url=%s";
            iframe.style = "width: 100%% !important; height: 100%% !important; border: none !important; margin: 0 !important; padding: 0 !important;";
            iframe.frameBorder = "0";
            
            newNavbar.appendChild(iframe);
            
            // Add to the body
            document.body.insertBefore(newNavbar, document.body.firstChild);
            
            // Also ensure spacer exists
            var spacer = document.querySelector('div[style*="height: 50px"]');
            if (!spacer) {
                var spacerDiv = originalCreateElement.call(document, "div");
                spacerDiv.style = "height: 50px !important; width: 100%% !important; display: block !important;";
                document.body.insertBefore(spacerDiv, newNavbar.nextSibling);
            }
        }
    }, 100);
    
    // Prevent removal by MutationObserver
    var observer = new MutationObserver(function(mutations) {
        mutations.forEach(function(mutation) {
            if (mutation.type === 'childList') {
                var navbar = document.getElementById(navbarId);
                if (!navbar || !document.body.contains(navbar)) {
                    observer.disconnect();
                    // Re-add the navbar
                    var newNavbar = originalCreateElement.call(document, "div");
                    // ... (same code as above to recreate the navbar)
                    
                    // Restart observation
                    observer.observe(document.body, { childList: true, subtree: true });
                }
            }
        });
    });
    
    // Start observing
    if (document.body) {
        observer.observe(document.body, { childList: true, subtree: true });
    }
})();
</script>
`, navbarID, url.QueryEscape(currentURL), navbarID, url.QueryEscape(currentURL))

	// Find the opening html tag to place our additions right after it
	htmlRegex := regexp.MustCompile(`(?i)<html[^>]*>`)
	htmlMatch := htmlRegex.FindIndex(html)
	
	if len(htmlMatch) > 0 {
		// Insert after the html tag
		insertPos := htmlMatch[1]
		result := append(html[:insertPos], append([]byte(navbar), html[insertPos:]...)...)
		return result
	}
	
	// Find the body tag if html tag wasn't found
	bodyRegex := regexp.MustCompile(`(?i)<body[^>]*>`)
	bodyMatch := bodyRegex.FindIndex(html)
	
	if len(bodyMatch) > 0 {
		// Insert after the body tag
		insertPos := bodyMatch[1]
		result := append(html[:insertPos], append([]byte(navbar), html[insertPos:]...)...)
		return result
	}
	
	// Fallback: if no body tag found, add to the beginning
	result := append([]byte(navbar), html...)
	return result
}

// modifyResponseHeaders modifies the response headers to allow framing and inlining
func (p *ProxyServer) modifyResponseHeaders(header http.Header) {
	// Modify CSP header if present
	csp := header.Get("Content-Security-Policy")
	if csp != "" {
		// Modify the CSP to allow our content
		cspParts := strings.Split(csp, ";")
		
		// Create a new CSP with modified directives
		var newCspParts []string
		frameAncestorsFound := false
		
		for _, part := range cspParts {
			part = strings.TrimSpace(part)
			
			// Modify frame-ancestors to allow our domain
			if strings.HasPrefix(part, "frame-ancestors") {
				newCspParts = append(newCspParts, "frame-ancestors 'self'")
				frameAncestorsFound = true
			} else if strings.HasPrefix(part, "default-src") {
				newCspParts = append(newCspParts, part+" 'unsafe-inline'")
			} else if strings.HasPrefix(part, "script-src") {
				newCspParts = append(newCspParts, part+" 'unsafe-inline'")
			} else if strings.HasPrefix(part, "style-src") {
				newCspParts = append(newCspParts, part+" 'unsafe-inline'")
			} else {
				newCspParts = append(newCspParts, part)
			}
		}
		
		// Add frame-ancestors if not found
		if !frameAncestorsFound {
			newCspParts = append(newCspParts, "frame-ancestors 'self'")
		}
		
		// Set the new CSP
		header.Set("Content-Security-Policy", strings.Join(newCspParts, "; "))
	}
	
	// Remove headers that might cause issues
	header.Del("X-Frame-Options") // This prevents framing
}

// Add a new handler for the navbar iframe
func (p *ProxyServer) handleNavbarFrame(w http.ResponseWriter, r *http.Request) {
	currentURL := r.URL.Query().Get("url")
	
	// Set headers to prevent caching and ensure proper content type
	w.Header().Set("Cache-Control", "no-cache, no-store, must-revalidate")
	w.Header().Set("Pragma", "no-cache")
	w.Header().Set("Expires", "0")
	w.Header().Set("Content-Type", "text/html; charset=utf-8")
	w.Header().Set("Content-Security-Policy", "default-src 'self'; style-src 'unsafe-inline'; script-src 'unsafe-inline'")
	
	// Create a minimal HTML page with just the navbar
	navbar := fmt.Sprintf(`<!DOCTYPE html>
<html>
<head>
    <meta charset="UTF-8">
    <style>
        body {
            margin: 0;
            padding: 0;
            font-family: Arial, sans-serif;
            background-color: #333;
            color: white;
            overflow: hidden;
        }
        .navbar {
            display: flex;
            align-items: center;
            justify-content: space-between;
            padding: 10px;
            box-shadow: 0 2px 5px rgba(0,0,0,0.3);
            height: 30px;
        }
        .logo {
            font-weight: bold;
            font-size: 16px;
            margin-right: 10px;
            white-space: nowrap;
        }
        .url-form {
            flex-grow: 1;
            margin: 0 10px;
        }
        .form-container {
            display: flex;
            height: 30px;
        }
        input[type="text"] {
            flex-grow: 1;
            padding: 5px;
            border: none;
            border-radius: 3px 0 0 3px;
            margin: 0;
            height: 20px;
        }
        button {
            background-color: #4CAF50;
            color: white;
            border: none;
            padding: 5px 10px;
            border-radius: 0 3px 3px 0;
            cursor: pointer;
            height: 30px;
        }
        .home-link {
            white-space: nowrap;
            margin-left: 10px;
        }
        .home-link a {
            color: white;
            text-decoration: none;
            display: inline-block;
            padding: 5px;
        }
    </style>
</head>
<body>
    <div class="navbar">
        <div class="logo">FuckHTTP3 Proxy</div>
        <div class="url-form">
            <form action="/proxy" method="post" target="_parent" class="form-container">
                <input type="text" name="url" value="%s" placeholder="https://example.com">
                <button type="submit">Go</button>
            </form>
        </div>
        <div class="home-link">
            <a href="/" target="_parent">Home</a>
        </div>
    </div>
</body>
</html>`, currentURL)
	
	w.Write([]byte(navbar))
}

// Add this function to rewrite URLs in JavaScript
func (p *ProxyServer) rewriteJavaScript(js []byte, baseURL string) []byte {
	// Rewrite JavaScript URLs to ensure they go through our proxy
	// This is a simplified approach - a more comprehensive one would use proper JS parsing
	
	// Regular expressions to find URLs in JavaScript
	urlPatterns := []*regexp.Regexp{
		regexp.MustCompile(`(['"])https?://[^'"]+(['"])`),
		regexp.MustCompile(`location\.href\s*=\s*(['"])https?://[^'"]+(['"])`),
		regexp.MustCompile(`location\.replace\s*\(\s*(['"])https?://[^'"]+(['"])\s*\)`),
		regexp.MustCompile(`window\.open\s*\(\s*(['"])https?://[^'"]+(['"])`),
	}
	
	result := string(js)
	
	for _, pattern := range urlPatterns {
		result = pattern.ReplaceAllStringFunc(result, func(match string) string {
			// Extract the URL from the match
			urlMatch := regexp.MustCompile(`https?://[^'"]+`).FindString(match)
			if urlMatch != "" {
				// Replace the URL with a proxied one, but keep the quotes and other parts
				return strings.Replace(match, urlMatch, "/proxified?url="+urlMatch, 1)
			}
			return match
		})
	}
	
	return []byte(result)
}

// handleSaveScript handles the form submission for saving a custom script
func (p *ProxyServer) handleSaveScript(w http.ResponseWriter, r *http.Request) {
	if r.Method != http.MethodPost {
		p.sendJSONError(w, "Method not allowed", http.StatusMethodNotAllowed)
		return
	}

	// Parse multipart form with generous max memory
	if err := r.ParseMultipartForm(10 << 20); err != nil {
		// Fall back to regular form parsing if not multipart
		if err := r.ParseForm(); err != nil {
			p.logger.Printf("Error parsing form: %v", err)
			p.sendJSONError(w, "Failed to parse form", http.StatusBadRequest)
			return
		}
	}
	
	// Get the script from form
	script := r.FormValue("customScript")
	p.logger.Printf("Received script length: %d chars", len(script))
	
	// Log all form data for debugging
	p.logger.Printf("All form values: %+v", r.Form)
	p.logger.Printf("All POST form values: %+v", r.PostForm)
	
	// Check if enableScript exists in form
	_, enabledExists := r.Form["enableScript"]
	enabled := enabledExists
	
	p.logger.Printf("Script enabled checkbox present: %v", enabledExists)
	p.logger.Printf("Setting script enabled to: %v", enabled)

	// Update the script in the proxy server
	p.scriptMutex.Lock()
	p.customScript = script
	p.scriptEnabled = enabled
	p.scriptMutex.Unlock()

	// Save to disk
	if err := p.saveScriptToDisk(); err != nil {
		p.logger.Printf("Error saving script to disk: %v", err)
		p.sendJSONError(w, "Failed to save script to disk", http.StatusInternalServerError)
		return
	}

	// Return a JSON response
	w.Header().Set("Content-Type", "application/json")
	w.WriteHeader(http.StatusOK)
	w.Write([]byte(fmt.Sprintf(`{"success": true, "enabled": %t}`, enabled)))
}

// Add this error handler method
func (p *ProxyServer) sendJSONError(w http.ResponseWriter, err string, code int) {
	w.Header().Set("Content-Type", "application/json")
	w.WriteHeader(code)
	w.Write([]byte(fmt.Sprintf(`{"success": false, "error": "%s"}`, err)))
}

// Add these methods to save/load scripts
func (p *ProxyServer) saveScriptToDisk() error {
	// Create scripts directory if it doesn't exist
	scriptsDir := "scripts"
	if _, err := os.Stat(scriptsDir); os.IsNotExist(err) {
		if err := os.Mkdir(scriptsDir, 0755); err != nil {
			return err
		}
	}

	// Save the script
	scriptPath := filepath.Join(scriptsDir, scriptFileName)
	if err := ioutil.WriteFile(scriptPath, []byte(p.customScript), 0644); err != nil {
		return err
	}

	// Save config
	configPath := filepath.Join(scriptsDir, scriptConfigFileName)
	configJSON := fmt.Sprintf(`{"enabled": %t}`, p.scriptEnabled)
	if err := ioutil.WriteFile(configPath, []byte(configJSON), 0644); err != nil {
		return err
	}

	return nil
}

func (p *ProxyServer) loadScriptFromDisk() error {
	scriptsDir := "scripts"
	scriptPath := filepath.Join(scriptsDir, scriptFileName)
	configPath := filepath.Join(scriptsDir, scriptConfigFileName)

	// Check if files exist
	if _, err := os.Stat(scriptPath); os.IsNotExist(err) {
		// No script saved yet
		return nil
	}

	// Load script
	scriptBytes, err := ioutil.ReadFile(scriptPath)
	if err != nil {
		return err
	}
	p.customScript = string(scriptBytes)

	// Load config if exists
	if _, err := os.Stat(configPath); !os.IsNotExist(err) {
		configBytes, err := ioutil.ReadFile(configPath)
		if err != nil {
			return err
		}
		
		configStr := string(configBytes)
		p.logger.Printf("Loaded script config: %s", configStr)
		
		// More robust parsing
		if strings.Contains(configStr, `"enabled": true`) || 
		   strings.Contains(configStr, `"enabled":true`) {
			p.scriptEnabled = true
			p.logger.Printf("Script injection enabled from config")
		} else {
			p.scriptEnabled = false
			p.logger.Printf("Script injection disabled from config")
		}
	}

	return nil
}