squirrels.go

package main

import (
	"encoding/json"
	"errors"
	"flag"
	"fmt"
	"io"
	"math/rand"
	"net/http"
	"time"
)

type SearchResult struct {
	Image string `json:"image"`
	Thumb string `json:"thumbnail"`
}

type SearchResults struct {
	Next    string         `json:"next"`
	Results []SearchResult `json:"results"`
}

type Image struct {
	Url   string
	Thumb string
}

type ImageResult struct {
	Image Image
	Error error
}

type ImageCache struct {
	Images       []Image
	ImageRequest chan chan ImageResult
	MaxS         int
	Query        string
}

func NewImageCache(query string, capacity, maxs int) *ImageCache {
	c := new(ImageCache)
	c.Query = query
	c.MaxS = (maxs / capacity) * capacity
	c.Images = make([]Image, 0, capacity)
	c.ImageRequest = make(chan chan ImageResult)
	return c
}

func (c *ImageCache) fillCache() error {
	s := (int(rand.Int31n(int32(c.MaxS))) / cap(c.Images)) * cap(c.Images)
	u := fmt.Sprintf(
		"http://duckduckgo.com/i.js?o=json&q=%s&s=%d",
		c.Query,
		s)
	r, err := http.Get(u)
	if err != nil {
		return err
	}
	defer r.Body.Close()
	var sr SearchResults
	jd := json.NewDecoder(r.Body)
	if err := jd.Decode(&sr); err != nil {
		return err
	}
	if len(sr.Results) > cap(c.Images) {
		sr.Results = sr.Results[:cap(c.Images)]
	}
	for i := 0; i < len(sr.Results); i++ {
		c.Images = append(c.Images, Image{sr.Results[i].Image, sr.Results[i].Thumb})
	}
	return nil
}

func (c *ImageCache) take(index int) Image {
	img := c.Images[index]
	c.Images = append(c.Images[:index], c.Images[index+1:]...)
	return img
}

func (c *ImageCache) get() ImageResult {
	if len(c.Images) == 0 {
		if err := c.fillCache(); err != nil {
			return ImageResult{Image{}, err}
		}
		if len(c.Images) == 0 {
			return ImageResult{Image{}, errors.New("no images found")}
		}
	}
	s := int(rand.Int31n(int32(len(c.Images))))
	return ImageResult{c.take(s), nil}
}

func (c *ImageCache) Run() {
	for {
		select {
		case r := <-c.ImageRequest:
			r <- c.get()
		}
	}
}

func (c *ImageCache) Get() ImageResult {
	r := make(chan ImageResult)
	c.ImageRequest <- r
	return <-r
}

var (
	squirrelCache *ImageCache
)

func squirrel(w http.ResponseWriter, r *http.Request) {
	ir := squirrelCache.Get()
	if ir.Error != nil {
		http.Error(w, ir.Error.Error(), http.StatusInternalServerError)
		return
	}
	var s string
	if r.FormValue("huge") != "" {
		s = ir.Image.Url
	} else {
		s = ir.Image.Thumb
	}
	res, err := http.Get(s)
	if err != nil {
		http.Error(w, ir.Error.Error(), http.StatusInternalServerError)
		return
	}
	defer res.Body.Close()
	w.Header().Set("Original-Source", s)
	io.Copy(w, res.Body)
}

var (
	_listen string
)

func init() {
	flag.StringVar(&_listen, "l", "0.0.0.0:8080", "listen address")
	flag.Parse()
	rand.Seed(time.Now().Unix())
}

func main() {
	squirrelCache = NewImageCache("squirrel%20images", 35, 1000)
	go squirrelCache.Run()
	http.HandleFunc("/squirrel", squirrel)
	err := http.ListenAndServe(_listen, nil)
	if err != nil {
		panic(err)
	}
}