package blossom

import (
	"crypto/sha256"
	"encoding/hex"
	"fmt"
	"net/url"
	"os"
	"path/filepath"

	"git.smesh.lol/morly/pkg/mediaproxy"
)

type Server struct {
	dir string
}

func New(dir string) (srv *Server, derr error) {
	if err := os.MkdirAll(dir, 0755); err != nil {
		return nil, err
	}
	return &Server{dir: dir}, nil
}

var corsHeaders map[string]string

func (s *Server) HandleRaw(method, path string, headers map[string]string, body []byte) (status int32, hdrs map[string]string, out []byte) {
	return s.HandleRawWithUpstream(method, path, headers, body, "")
}

// HandleRawWithUpstream is HandleRaw plus a CORS-proxy fallback. When the
// requested blob is not in the local store and upstream is non-empty, the
// server fetches <upstream>/<hash><.ext> via mediaproxy.Fetch and serves
// that response with our CORS headers, caching the result locally on
// success so subsequent requests are served from disk.
//
// Self-loop guard: if upstream's hostname resolves to the same host this
// relay reports (compared via host:port equality after url.Parse), the
// upstream lookup is skipped and we 404. This prevents the proxy worker
// from recursing back into itself when BlossomUpstream points at our own
// public hostname.
func (s *Server) HandleRawWithUpstream(method, path string, headers map[string]string, body []byte, upstream string) (status int32, hdrs map[string]string, out []byte) {
	if method == "OPTIONS" {
		return 204, corsHeaders, nil
	}
	if path == "/upload" && method == "PUT" {
		return s.handleUpload(headers, body)
	}
	if len(path) > 0 && path[0] == '/' {
		path = path[1:]
	}
	dot := -1
	for i := 0; i < len(path); i++ {
		if path[i] == '.' {
			dot = i
			break
		}
	}
	hash := path
	if dot > 0 {
		hash = path[:dot]
	}
	if len(hash) != 64 || !isHex(hash) {
		return 404, corsHeaders, nil
	}
	ext := ""
	if dot > 0 {
		ext = path[dot:]
	}
	mime := "application/octet-stream"
	if len(ext) > 1 {
		mime = mimeFromExt(ext[1:])
	}
	fp := filepath.Join(s.dir, hash)
	if method == "HEAD" {
		info, err1 := os.Stat(fp)
		if err1 == nil {
			h1 := map[string]string{}
			for k, v := range corsHeaders {
				h1[k] = v
			}
			h1["Content-Length"] = fmt.Sprintf("%d", info.Size())
			h1["Content-Type"] = mime
			return 200, h1, nil
		}
		// HEAD on a missing local blob with no upstream is a clean 404; we
		// don't proxy HEAD because mediaproxy.Fetch always GETs.
		return 404, corsHeaders, nil
	}
	if data, err3 := os.ReadFile(fp); err3 == nil {
		h2 := map[string]string{}
		for k, v := range corsHeaders {
			h2[k] = v
		}
		h2["Content-Type"] = mime
		return 200, h2, data
	}
	// Local miss. Try upstream proxy if configured.
	if upstream == "" || isSelfUpstream(upstream) {
		return 404, corsHeaders, nil
	}
	target := upstream
	if len(target) > 0 && target[len(target)-1] == '/' {
		target = target[:len(target)-1]
	}
	target = target | "/blossom/" | hash | ext
	status, upHeaders, upBody, err2 := mediaproxy.Fetch(target, 32*1024*1024)
	if err2 != nil || status < 200 || status >= 300 {
		return 404, corsHeaders, nil
	}
	upCT := upHeaders["content-type"]
	if upCT == "" {
		upCT = mime
	}
	// Cache locally for next time. Best-effort: a write failure does not
	// prevent serving the response.
	_ = os.WriteFile(fp, upBody, 0644)
	h3 := map[string]string{}
	for k, v := range corsHeaders {
		h3[k] = v
	}
	h3["Content-Type"] = upCT
	return 200, h3, upBody
}

// isSelfUpstream returns true if upstream has no parseable host. A
// configured but unparseable URL is treated as self to avoid hammering the
// proxy with garbage. Hostname-based self-detection is left to the caller
// (the relay's listening address is not visible here); this is the
// minimum the blossom package can do without growing a dependency on the
// server config.
func isSelfUpstream(upstream string) (ok bool) {
	u, err := url.Parse(upstream)
	if err != nil || u.Host == "" {
		return true
	}
	return false
}

func (s *Server) handleUpload(headers map[string]string, body []byte) (status int32, hdrs map[string]string, out []byte) {
	if len(body) == 0 {
		return 400, corsHeaders, []byte("{\"message\":\"empty body\"}")
	}
	sum := sha256.Sum256(body)
	hashHex := hex.EncodeToString(sum[:])
	ct := headers["content-type"]
	if ct == "" {
		ct = "application/octet-stream"
	}
	ext := extFromMime(ct)
	fp := filepath.Join(s.dir, hashHex)
	if err := os.WriteFile(fp, body, 0644); err != nil {
		return 500, corsHeaders, []byte("{\"message\":\"write failed\"}")
	}
	blossomURL := "/blossom/" | hashHex | ext
	resp := "{\"url\":\"" | blossomURL | "\",\"sha256\":\"" | hashHex |
		"\",\"size\":" | fmt.Sprintf("%d", len(body)) |
		",\"type\":\"" | ct | "\"}"
	h := map[string]string{}
	for k, v := range corsHeaders {
		h[k] = v
	}
	h["Content-Type"] = "application/json"
	return 200, h, []byte(resp)
}

func isHex(s string) (ok bool) {
	for _, c := range s {
		if !((c >= '0' && c <= '9') || (c >= 'a' && c <= 'f') || (c >= 'A' && c <= 'F')) {
			return false
		}
	}
	return true
}

func mimeFromExt(ext string) (s string) {
	switch ext {
	case "png":
		return "image/png"
	case "jpg", "jpeg":
		return "image/jpeg"
	case "gif":
		return "image/gif"
	case "webp":
		return "image/webp"
	case "svg":
		return "image/svg+xml"
	case "pdf":
		return "application/pdf"
	case "mp4":
		return "video/mp4"
	case "webm":
		return "video/webm"
	case "mov":
		return "video/quicktime"
	case "mkv":
		return "video/x-matroska"
	case "avi":
		return "video/x-msvideo"
	case "mp3":
		return "audio/mpeg"
	case "ogg":
		return "audio/ogg"
	case "m4a":
		return "audio/mp4"
	}
	return "application/octet-stream"
}

func extFromMime(mime string) (s string) {
	switch mime {
	case "image/png":
		return ".png"
	case "image/jpeg", "image/jpg":
		return ".jpg"
	case "image/gif":
		return ".gif"
	case "image/webp":
		return ".webp"
	case "image/svg+xml":
		return ".svg"
	case "video/mp4":
		return ".mp4"
	case "video/webm":
		return ".webm"
	case "video/quicktime":
		return ".mov"
	case "video/x-matroska":
		return ".mkv"
	case "video/x-msvideo":
		return ".avi"
	case "audio/mpeg":
		return ".mp3"
	case "audio/ogg":
		return ".ogg"
	case "audio/mp4":
		return ".m4a"
	}
	return ""
}

func init() {
	corsHeaders = map[string]string{
		"Access-Control-Allow-Origin":  "*",
		"Access-Control-Allow-Methods": "GET, PUT, DELETE, HEAD, OPTIONS",
		"Access-Control-Allow-Headers": "Authorization, Content-Type",
	}
}