package main

import (
	"crypto/rand"
	"encoding/json"
	"fmt"
	"io"
	"log"
	"net/http"
	"os"
	"path/filepath"
	"strings"
	"time"
)

var uploadDir string

const idLength = 8
const idCharset = "abcdefghijklmnopqrstuvwxyz0123456789"

func main() {
	uploadDir = os.Getenv("UPLOAD_DIR")
	if uploadDir == "" {
		uploadDir = "/uploads"
	}

	port := os.Getenv("PORT")
	if port == "" {
		port = "8080"
	}

	if err := os.MkdirAll(uploadDir, 0755); err != nil {
		log.Fatalf("Failed to create upload directory: %v", err)
	}

	http.HandleFunc("/", handleRequest)

	log.Printf("Starting server on port %s, upload directory: %s", port, uploadDir)
	if err := http.ListenAndServe(":"+port, nil); err != nil {
		log.Fatalf("Server failed: %v", err)
	}
}

func handleRequest(w http.ResponseWriter, r *http.Request) {
	start := time.Now()

	if r.Method == "GET" && r.URL.Path == "/" {
		serveFrontpage(w, r)
		logRequest(r, 200, time.Since(start))
		return
	}

	if r.Method == "POST" && r.URL.Path == "/upload" {
		handleMultipartUpload(w, r, start)
		return
	}

	if r.Method == "PUT" {
		handleCurlUpload(w, r, start)
		return
	}

	if r.Method == "GET" {
		handleDownload(w, r, start)
		return
	}

	w.WriteHeader(http.StatusMethodNotAllowed)
	logRequest(r, 405, time.Since(start))
}

func handleCurlUpload(w http.ResponseWriter, r *http.Request, start time.Time) {
	filename := strings.TrimPrefix(r.URL.Path, "/")
	if filename == "" {
		http.Error(w, "Filename required", http.StatusBadRequest)
		logRequest(r, 400, time.Since(start))
		return
	}

	filename = filepath.Base(filename)

	id, err := generateUniqueID()
	if err != nil {
		http.Error(w, "Failed to generate unique ID", http.StatusInternalServerError)
		logRequest(r, 500, time.Since(start))
		return
	}

	targetPath := filepath.Join(uploadDir, id, filename)
	file, err := os.Create(targetPath)
	if err != nil {
		http.Error(w, "Failed to create file", http.StatusInternalServerError)
		logRequest(r, 500, time.Since(start))
		return
	}
	defer file.Close()

	written, err := io.Copy(file, r.Body)
	if err != nil {
		http.Error(w, "Failed to write file", http.StatusInternalServerError)
		logRequest(r, 500, time.Since(start))
		return
	}

	proto := getProtocol(r)
	host := getHost(r)
	downloadURL := fmt.Sprintf("%s://%s/%s/%s\n", proto, host, id, filename)

	w.WriteHeader(http.StatusOK)
	w.Write([]byte(downloadURL))

	log.Printf("Uploaded %s (%d bytes) -> %s", filename, written, downloadURL)
	logRequest(r, 200, time.Since(start))
}

func handleMultipartUpload(w http.ResponseWriter, r *http.Request, start time.Time) {
	if err := r.ParseMultipartForm(32 << 20); err != nil {
		http.Error(w, "Failed to parse multipart form", http.StatusBadRequest)
		logRequest(r, 400, time.Since(start))
		return
	}

	files := r.MultipartForm.File["files"]
	if len(files) == 0 {
		http.Error(w, "No files provided", http.StatusBadRequest)
		logRequest(r, 400, time.Since(start))
		return
	}

	proto := getProtocol(r)
	host := getHost(r)
	urls := make([]string, 0, len(files))

	for _, fileHeader := range files {
		src, err := fileHeader.Open()
		if err != nil {
			continue
		}

		id, err := generateUniqueID()
		if err != nil {
			src.Close()
			continue
		}

		filename := filepath.Base(fileHeader.Filename)
		targetPath := filepath.Join(uploadDir, id, filename)

		dst, err := os.Create(targetPath)
		if err != nil {
			src.Close()
			continue
		}

		written, err := io.Copy(dst, src)
		dst.Close()
		src.Close()

		if err != nil {
			continue
		}

		downloadURL := fmt.Sprintf("%s://%s/%s/%s", proto, host, id, filename)
		urls = append(urls, downloadURL)
		log.Printf("Uploaded %s (%d bytes) -> %s", filename, written, downloadURL)
	}

	w.Header().Set("Content-Type", "application/json")
	json.NewEncoder(w).Encode(map[string][]string{"urls": urls})
	logRequest(r, 200, time.Since(start))
}

func handleDownload(w http.ResponseWriter, r *http.Request, start time.Time) {
	parts := strings.Split(strings.TrimPrefix(r.URL.Path, "/"), "/")
	if len(parts) != 2 {
		http.NotFound(w, r)
		logRequest(r, 404, time.Since(start))
		return
	}

	id := parts[0]
	filename := parts[1]

	filePath := filepath.Join(uploadDir, id, filename)
	file, err := os.Open(filePath)
	if err != nil {
		http.NotFound(w, r)
		logRequest(r, 404, time.Since(start))
		return
	}
	defer file.Close()

	stat, err := file.Stat()
	if err != nil {
		http.Error(w, "Failed to stat file", http.StatusInternalServerError)
		logRequest(r, 500, time.Since(start))
		return
	}

	w.Header().Set("Content-Type", detectContentType(filename))
	w.Header().Set("Content-Disposition", fmt.Sprintf("inline; filename=\"%s\"", filename))
	w.Header().Set("Content-Length", fmt.Sprintf("%d", stat.Size()))

	io.Copy(w, file)
	logRequest(r, 200, time.Since(start))
}

func generateUniqueID() (string, error) {
	const maxRetries = 100

	for i := 0; i < maxRetries; i++ {
		id := generateID()
		targetDir := filepath.Join(uploadDir, id)

		if _, err := os.Stat(targetDir); os.IsNotExist(err) {
			if err := os.MkdirAll(targetDir, 0755); err != nil {
				return "", err
			}
			return id, nil
		}
	}

	return "", fmt.Errorf("failed to generate unique ID after %d retries", maxRetries)
}

func generateID() string {
	bytes := make([]byte, idLength)
	if _, err := rand.Read(bytes); err != nil {
		panic(err)
	}

	for i, b := range bytes {
		bytes[i] = idCharset[int(b)%len(idCharset)]
	}

	return string(bytes)
}

func getHost(r *http.Request) string {
	// Check X-Forwarded-Host (common proxy header)
	if host := r.Header.Get("X-Forwarded-Host"); host != "" {
		return host
	}

	// Check standard Forwarded header (RFC 7239)
	if forwarded := r.Header.Get("Forwarded"); forwarded != "" {
		// Parse "Forwarded: for=...; host=example.com; proto=https"
		parts := strings.Split(forwarded, ";")
		for _, part := range parts {
			part = strings.TrimSpace(part)
			if strings.HasPrefix(part, "host=") {
				return strings.TrimPrefix(part, "host=")
			}
		}
	}

	// Fallback to Host header
	if host := r.Header.Get("Host"); host != "" {
		return host
	}

	return "localhost:8080"
}

func getProtocol(r *http.Request) string {
	// Check X-Forwarded-Proto (common proxy header)
	if proto := r.Header.Get("X-Forwarded-Proto"); proto != "" {
		return proto
	}

	// Check standard Forwarded header (RFC 7239)
	if forwarded := r.Header.Get("Forwarded"); forwarded != "" {
		// Parse "Forwarded: for=...; proto=https; host=..."
		parts := strings.Split(forwarded, ";")
		for _, part := range parts {
			part = strings.TrimSpace(part)
			if strings.HasPrefix(part, "proto=") {
				return strings.TrimPrefix(part, "proto=")
			}
		}
	}

	// Check X-Forwarded-SSL (some proxies use this)
	if ssl := r.Header.Get("X-Forwarded-SSL"); ssl == "on" {
		return "https"
	}

	// Check if connection is TLS
	if r.TLS != nil {
		return "https"
	}

	return "http"
}

func detectContentType(filename string) string {
	ext := strings.ToLower(filepath.Ext(filename))
	types := map[string]string{
		".txt":  "text/plain",
		".html": "text/html",
		".css":  "text/css",
		".js":   "application/javascript",
		".json": "application/json",
		".xml":  "application/xml",
		".pdf":  "application/pdf",
		".zip":  "application/zip",
		".jpg":  "image/jpeg",
		".jpeg": "image/jpeg",
		".png":  "image/png",
		".gif":  "image/gif",
		".svg":  "image/svg+xml",
		".mp3":  "audio/mpeg",
		".mp4":  "video/mp4",
		".webm": "video/webm",
	}

	if ct, ok := types[ext]; ok {
		return ct
	}
	return "application/octet-stream"
}

func logRequest(r *http.Request, status int, duration time.Duration) {
	log.Printf("[%s] %s %s %d %v", time.Now().Format("2006-01-02 15:04:05"), r.Method, r.URL.Path, status, duration)
}

func serveFrontpage(w http.ResponseWriter, r *http.Request) {
	proto := getProtocol(r)
	host := getHost(r)
	html := fmt.Sprintf(`<!DOCTYPE html>
<html>
<head>
    <meta charset="UTF-8">
    <meta name="viewport" content="width=device-width, initial-scale=1.0">
    <title>Beam - File Upload Service</title>
    <style>
        * { margin: 0; padding: 0; box-sizing: border-box; }
        body {
            font-family: monospace;
            max-width: 800px;
            margin: 40px auto;
            padding: 20px;
            line-height: 1.6;
        }
        h1 { margin-bottom: 30px; font-size: 2em; }
        h2 { margin: 30px 0 15px 0; font-size: 1.3em; }
        .section {
            background: #f5f5f5;
            padding: 20px;
            margin: 20px 0;
            border-radius: 4px;
        }
        code {
            background: #333;
            color: #0f0;
            padding: 15px;
            display: block;
            margin: 10px 0;
            border-radius: 4px;
            overflow-x: auto;
        }
        .upload-area {
            border: 3px dashed #ccc;
            padding: 40px;
            text-align: center;
            cursor: pointer;
            transition: all 0.3s;
        }
        .upload-area.dragover {
            border-color: #333;
            background: #f0f0f0;
        }
        .upload-area:hover {
            border-color: #666;
        }
        input[type="file"] { display: none; }
        button {
            background: #333;
            color: white;
            border: none;
            padding: 10px 20px;
            cursor: pointer;
            font-family: monospace;
            font-size: 1em;
            border-radius: 4px;
            margin-top: 10px;
        }
        button:hover { background: #555; }
        .progress-container {
            width: 100%%;
            background: #ddd;
            border-radius: 4px;
            margin: 10px 0;
            height: 20px;
            display: none;
        }
        .progress-bar {
            height: 100%%;
            background: #333;
            border-radius: 4px;
            width: 0;
            transition: width 0.3s;
        }
        .results {
            margin-top: 20px;
        }
        .result-item {
            background: white;
            padding: 10px;
            margin: 5px 0;
            border-radius: 4px;
            word-break: break-all;
        }
        .result-item a {
            color: #0066cc;
        }
    </style>
</head>
<body>
    <h1>🚀 Beam - File Upload Service</h1>

    <div class="section">
        <h2>Upload with curl (Recommended)</h2>
        <code>curl --upload-file myfile.txt %s://%s/myfile.txt</code>
        <p style="margin-top: 10px;">The server will respond with a download URL.</p>
    </div>

    <div class="section">
        <h2>Download a file</h2>
        <code>curl %s://%s/xxxxxxxx/myfile.txt -o myfile.txt</code>
        <p style="margin-top: 10px;">Or simply open the URL in your browser.</p>
    </div>

    <div class="section">
        <h2>Web Upload</h2>
        <div class="upload-area" id="uploadArea">
            <p>📁 Click to select files or drag and drop here</p>
            <input type="file" id="fileInput" multiple>
            <button onclick="document.getElementById('fileInput').click()">Select Files</button>
        </div>
        <div class="progress-container" id="progressContainer">
            <div class="progress-bar" id="progressBar"></div>
        </div>
        <div class="results" id="results"></div>
    </div>

    <script>
        const uploadArea = document.getElementById('uploadArea');
        const fileInput = document.getElementById('fileInput');
        const progressContainer = document.getElementById('progressContainer');
        const progressBar = document.getElementById('progressBar');
        const results = document.getElementById('results');

        uploadArea.addEventListener('click', () => fileInput.click());

        uploadArea.addEventListener('dragover', (e) => {
            e.preventDefault();
            uploadArea.classList.add('dragover');
        });

        uploadArea.addEventListener('dragleave', () => {
            uploadArea.classList.remove('dragover');
        });

        uploadArea.addEventListener('drop', (e) => {
            e.preventDefault();
            uploadArea.classList.remove('dragover');
            const files = e.dataTransfer.files;
            uploadFiles(files);
        });

        fileInput.addEventListener('change', (e) => {
            uploadFiles(e.target.files);
        });

        function uploadFiles(files) {
            if (files.length === 0) return;

            const formData = new FormData();
            for (let file of files) {
                formData.append('files', file);
            }

            progressContainer.style.display = 'block';
            progressBar.style.width = '0%%';
            results.innerHTML = '';

            const xhr = new XMLHttpRequest();

            xhr.upload.addEventListener('progress', (e) => {
                if (e.lengthComputable) {
                    const percent = (e.loaded / e.total) * 100;
                    progressBar.style.width = percent + '%%';
                }
            });

            xhr.addEventListener('load', () => {
                if (xhr.status === 200) {
                    const response = JSON.parse(xhr.responseText);
                    displayResults(response.urls);
                } else {
                    results.innerHTML = '<div class="result-item">Upload failed: ' + xhr.statusText + '</div>';
                }
                progressBar.style.width = '100%%';
            });

            xhr.addEventListener('error', () => {
                results.innerHTML = '<div class="result-item">Upload failed: Network error</div>';
            });

            xhr.open('POST', '/upload');
            xhr.send(formData);
        }

        function displayResults(urls) {
            results.innerHTML = '<h3>Download URLs:</h3>';
            urls.forEach(url => {
                const div = document.createElement('div');
                div.className = 'result-item';
                div.innerHTML = '<a href="' + url + '" target="_blank">' + url + '</a>';
                results.appendChild(div);
            });
        }
    </script>
</body>
</html>`, proto, host, proto, host)

	w.Header().Set("Content-Type", "text/html; charset=utf-8")
	w.Write([]byte(html))
}