package main

import (
	"bytes"
	"fmt"
	"image"
	_ "image/gif"
	_ "image/jpeg"
	_ "image/png"
	"io"
	"log"
	"math/rand"
	"net/http"
	"os"
	"path/filepath"
	"strconv"
	"strings"
	"sync"
	"time"

	xdraw "golang.org/x/image/draw"
)

const (
	defaultListen  = ":3333"
	defaultDir     = "./images"
	defaultWidth   = 90
	cacheSuffixFmt = ".ansi.%d.txt"
	allowedUA      = "curl"
	maxImageDim    = 2000
)

var (
	imgDir      string
	redirectURL string
	cacheLock   sync.RWMutex
	cachedPaths []string
)

func main() {
	rand.Seed(time.Now().UnixNano())

	imgDir = getenv("IMAGES_DIR", defaultDir)
	redirectURL = getenv("REDIRECT_URL", "https://example.com")
	listen := getenv("LISTEN", defaultListen)

	if err := os.MkdirAll(imgDir, 0755); err != nil {
		log.Fatalf("failed to create images dir: %v", err)
	}

	log.Printf("images directory: %s", imgDir)
	log.Printf("pre-generating ascii caches at width=%d ...", defaultWidth)
	scanAndPreGen(defaultWidth)
	refreshCachedPaths(defaultWidth)

	http.HandleFunc("/", rootHandler)

	log.Printf("listening on %s — only curl clients accepted; non-curl will be redirected to %s", listen, redirectURL)
	if err := http.ListenAndServe(listen, nil); err != nil {
		log.Fatalf("server failed: %v", err)
	}
}

func getenv(k, fallback string) string {
	v := os.Getenv(k)
	if v == "" {
		return fallback
	}
	return v
}

func rootHandler(w http.ResponseWriter, r *http.Request) {
	ua := strings.ToLower(r.Header.Get("User-Agent"))
	if !strings.Contains(ua, allowedUA) {
		http.Redirect(w, r, redirectURL, http.StatusFound)
		return
	}

	qw := r.URL.Query().Get("w")
	width := defaultWidth
	if qw != "" {
		if n, err := strconv.Atoi(qw); err == nil && n > 10 && n <= 1000 {
			width = n
		}
	}

	cacheLock.RLock()
	candidates := filteredCachedPaths(width)
	cacheLock.RUnlock()

	if len(candidates) == 0 {
		images := listImageFiles()
		if len(images) == 0 {
			http.Error(w, "no images found in images directory\n", http.StatusInternalServerError)
			return
		}
		imgFile := images[rand.Intn(len(images))]
		cachePath := cacheFilePath(imgFile, width)

		log.Printf("generating cache for %s (width=%d) on demand", imgFile, width)
		if err := generateAndSaveASCII(imgFile, cachePath, width); err != nil {
			log.Printf("generation error: %v", err)
			http.Error(w, "failed to generate ascii\n", http.StatusInternalServerError)
			return
		}

		cacheLock.Lock()
		refreshCachedPaths(width)
		candidates = filteredCachedPaths(width)
		cacheLock.Unlock()
	}

	if len(candidates) == 0 {
		http.Error(w, "no cached ascii available\n", http.StatusInternalServerError)
		return
	}

	chosen := candidates[rand.Intn(len(candidates))]
	f, err := os.Open(chosen)
	if err != nil {
		http.Error(w, "failed to read ascii file\n", http.StatusInternalServerError)
		return
	}
	defer f.Close()

	w.Header().Set("Content-Type", "text/plain; charset=utf8")
	w.WriteHeader(http.StatusOK)
	if _, err := io.Copy(w, f); err != nil {
		log.Printf("write error: %v", err)
	}
}

func listImageFiles() []string {
	var files []string
	_ = filepath.Walk(imgDir, func(path string, info os.FileInfo, err error) error {
		if err != nil || info.IsDir() {
			return nil
		}
		ext := strings.ToLower(filepath.Ext(path))
		switch ext {
		case ".jpg", ".jpeg", ".png", ".gif", ".bmp":
			files = append(files, path)
		default:
		}
		return nil
	})
	return files
}

func cacheFilePath(imagePath string, width int) string {
	dir := filepath.Dir(imagePath)
	base := filepath.Base(imagePath)
	ext := filepath.Ext(base)
	name := strings.TrimSuffix(base, ext)
	return filepath.Join(dir, fmt.Sprintf("%s"+cacheSuffixFmt, name, width))
}

func scanAndPreGen(width int) {
	images := listImageFiles()
	for _, img := range images {
		cache := cacheFilePath(img, width)
		if _, err := os.Stat(cache); err == nil {
			continue
		}
		if err := generateAndSaveASCII(img, cache, width); err != nil {
			log.Printf("failed to generate for %s: %v", img, err)
		}
	}
}

func refreshCachedPaths(width int) {
	suffix := fmt.Sprintf(cacheSuffixFmt, width)
	found := []string{}
	_ = filepath.Walk(imgDir, func(path string, info os.FileInfo, err error) error {
		if err != nil || info.IsDir() {
			return nil
		}
		if strings.HasSuffix(path, suffix) {
			found = append(found, path)
		}
		return nil
	})

	cacheLock.Lock()
	cachedPaths = found
	cacheLock.Unlock()
}

func filteredCachedPaths(width int) []string {
	suffix := fmt.Sprintf(cacheSuffixFmt, width)
	out := make([]string, 0, len(cachedPaths))
	for _, p := range cachedPaths {
		if strings.HasSuffix(p, suffix) {
			out = append(out, p)
		}
	}
	return out
}

func generateAndSaveASCII(imagePath, cachePath string, width int) error {
	f, err := os.Open(imagePath)
	if err != nil {
		return err
	}
	defer f.Close()

	img, _, err := image.Decode(f)
	if err != nil {
		return err
	}

	w0 := img.Bounds().Dx()
	h0 := img.Bounds().Dy()
	if w0 > maxImageDim || h0 > maxImageDim {
		scale := float64(maxImageDim) / float64(max(w0, h0))
		newW := int(float64(w0) * scale)
		newH := int(float64(h0) * scale)
		if newW < 1 {
			newW = 1
		}
		if newH < 1 {
			newH = 1
		}
		dst := image.NewRGBA(image.Rect(0, 0, newW, newH))
		xdraw.ApproxBiLinear.Scale(dst, dst.Bounds(), img, img.Bounds(), xdraw.Over, nil)
		img = dst
	}

	ansi := ImageToANSI(img, width)

	if err := os.MkdirAll(filepath.Dir(cachePath), 0755); err != nil {
		return err
	}
	if err := os.WriteFile(cachePath, []byte(ansi), 0644); err != nil {
		return err
	}

	cacheLock.Lock()
	cachedPaths = append(cachedPaths, cachePath)
	cacheLock.Unlock()
	return nil
}

func ImageToANSI(src image.Image, cols int) string {
	b := src.Bounds()
	origW := b.Dx()
	origH := b.Dy()
	if origW == 0 || origH == 0 {
		return ""
	}

	charRowsFloat := (float64(origH) / float64(origW)) * float64(cols) * 0.5
	charRows := int(charRowsFloat)
	if charRows < 1 {
		charRows = 1
	}
	targetH := charRows * 2
	targetW := cols

	dst := image.NewRGBA(image.Rect(0, 0, targetW, targetH))
	xdraw.ApproxBiLinear.Scale(dst, dst.Bounds(), src, src.Bounds(), xdraw.Over, nil)

	var out bytes.Buffer
	for y := 0; y < targetH; y += 2 {
		for x := 0; x < targetW; x++ {
			tr := colorAt(dst, x, y)
			br := colorAt(dst, x, y+1)
			out.WriteString(fmt.Sprintf("\x1b[38;2;%d;%d;%dm\x1b[48;2;%d;%d;%dm▀",
				tr.R, tr.G, tr.B,
				br.R, br.G, br.B))
		}
		out.WriteString("\x1b[0m\n")
	}
	return out.String()
}

type c8 struct{ R, G, B uint8 }

func colorAt(img *image.RGBA, x, y int) c8 {
	r, g, b, _ := img.At(x, y).RGBA()
	return c8{uint8(r >> 8), uint8(g >> 8), uint8(b >> 8)}
}

func max(a, b int) int {
	if a > b {
		return a
	}
	return b
}