package nsfw

import (
	"context"
	"crypto/md5"
	"encoding/hex"
	"encoding/json"
	"fmt"
	"log"
	"regexp"
	"strings"
	"sync"
	"time"

	"github.com/gin-gonic/gin"
	"github.com/reimage-ai/broker/firebase"
	"github.com/reimage-ai/broker/job"
	"google.golang.org/genai"
)

const GENERIC_MESSAGE = "Prompt not allowed.\nIf you believe this message to be in error please contact us."

// banned prompt regexes
var promptRegexFilters = []*regexp.Regexp{
	regexp.MustCompile(`remove.{0,19}clothing`),
	regexp.MustCompile(`remove.{0,19}clothes`),
	regexp.MustCompile(`remove.{0,19}shirt`),
	regexp.MustCompile(`remove.{0,19}pants`),
	regexp.MustCompile(`replace.{0,19}clothing`),
	regexp.MustCompile(`replace.{0,19}clothes`),
	regexp.MustCompile(`replace.{0,19}shirt`),
	regexp.MustCompile(`replace.{0,19}pants`),
}

var regexAllAscii = regexp.MustCompile("[[:^ascii:]]")
var regexReturnChars = regexp.MustCompile("[\r\n]")
var regexWhiteSpace = regexp.MustCompile(`\s+`)
var regexWordSplit = regexp.MustCompile(`[ \s,]`)

var geminiClient *genai.Client = nil
var geminiMemoCache map[string]bool = make(map[string]bool)
var geminiMemoCacheLock sync.RWMutex

func Init() {
	var apiKey = "AQ.Ab8RN6JcKFxQ5bjVZoFwxDNT-wIFR-yJ9FQ8jztZin94fUeLYg"
	var err error
	ctx := context.Background()
	geminiClient, err = genai.NewClient(ctx, &genai.ClientConfig{
		APIKey: apiKey,
	})
	if err != nil {
		log.Fatalf("Failed to create client: %v", err)
	}
}

// quick blocking of prompts
func Allowed(c *gin.Context, j *job.Job) (bool, string) {
	if j == nil {
		return true, ""
	}
	var promptobj any = j.Inputs["prompt"]
	if promptobj == nil {
		return true, ""
	}
	email := getEmail(c)
	agent := strings.ToLower(c.Request.Header.Get("User-Agent"))

	if strings.HasPrefix(agent, "python-requests/") {
		return true, ""
	}

	prompt := fmt.Sprintf("%v", promptobj)
	prompt = strings.ToLower(prompt)
	prompt = regexAllAscii.ReplaceAllLiteralString(prompt, "")
	prompt = regexReturnChars.ReplaceAllLiteralString(prompt, " ")
	prompt = regexWhiteSpace.ReplaceAllString(prompt, " ")

	// this runs asynchronously as can take a while
	go geminiCheckSafety(geminiClient, c, j, prompt)

	emailProperties := firebase.GetEmailProperties(email)
	if emailProperties != nil && emailProperties.GetBannedPrompt() {
		log.Printf("Alert: prompt not allowed (user banned from prompts): %v (%v)", prompt, email)
		return false, GENERIC_MESSAGE
	}

	words := regexWordSplit.Split(prompt, -1)

	badWords := []string{
		"bra",
		"panties",
		"swimsuit",
	}
	for i := 0; i < len(badWords); i++ {
		for j := 0; j < len(words); j++ {
			if badWords[i] == words[j] {
				log.Printf("Alert: prompt not allowed (bad word): %v (%v)", prompt, email)
				// not temp banned
				return false, GENERIC_MESSAGE
			}
		}
	}

	veryBadWords := []string{
		"anal",
		"anilingus",
		"anus",
		"arsehole",
		"areola",
		"ass",
		"asshole",
		"barenaked",
		"breasts",
		"tits",
		"blowjob",
		"boner",
		"boob",
		"boobs",
		"bukkake",
		"butthole",
		"clit",
		"clitoris",
		"cock",
		"cocks",
		"cum",
		"cunnilingus",
		"cunt",
		"deepthroat",
		"erotic",
		"erotism",
		"faggot",
		"fellatio",
		"fuck",
		"fucking",
		"genitals",
		"handjob",
		"hentai",
		"kinky",
		"lingerie",
		"masturbate",
		"motherfucker",
		"naked",
		"nipple",
		"nipples",
		"nude",
		"nudity",
		"orgasm",
		"penis",
		"pussy",
		"vagina",
	}
	for i := 0; i < len(veryBadWords); i++ {
		for j := 0; j < len(words); j++ {
			if veryBadWords[i] == words[j] {
				log.Printf("Alert: prompt not allowed (very bad word): %v (%v)", prompt, email)
				applyTempBan(email, agent)
				return false, GENERIC_MESSAGE
			}
		}
	}

	bannedPhrases := []string{
		"naked woman",
		"nude woman",
		"deep throat",
		"camel toe",
	}
	for i := 0; i < len(bannedPhrases); i++ {
		if strings.Contains(prompt, bannedPhrases[i]) {
			log.Printf("Alert: prompt not allowed (phrase): %v (%v)", prompt, email)
			applyTempBan(email, agent)
			return false, GENERIC_MESSAGE
		}
	}

	for i := 0; i < len(promptRegexFilters); i++ {
		if promptRegexFilters[i].MatchString(prompt) {
			log.Printf("Alert: prompt not allowed (regex): %v (%v)", prompt, email)
			applyTempBan(email, agent)
			return false, GENERIC_MESSAGE
		}
	}

	return true, ""
}

// apply a temporary ban to this user
func applyTempBan(email string, agent string) {
	// because of google play requirements we are stricter with temp bans on android
	// however old version of the client just use Dart/3.11 as the user-agent so ban those too for now
	// if strings.Contains(agent, "dart") || strings.Contains(agent, "android") {
	if strings.Contains(agent, "android") {
		emailProperties := firebase.GetEmailProperties(email)
		if emailProperties != nil {
			emailProperties.ApplyTempBan()
		} else {
			log.Printf("Alert: Failed to temp ban (no email properties found) (%v)", email)
		}
	}

}

// return the of the caller (if possible)
func getEmail(c *gin.Context) string {
	email := ""
	authToken := c.GetHeader("authorization")
	if len(authToken) == 32 {
		accountProperties := firebase.GetAccountPropertiesByKey(authToken)
		if accountProperties != nil {
			email = accountProperties.Email
		} else {
			email = "-"
		}
	} else {
		_, email = firebase.Email(authToken)
		if email == "" {
			email = "-"
		}
	}
	return email
}

func pointer[T any](v T) *T { return &v }

func geminiCheckSafety(client *genai.Client, c *gin.Context, j *job.Job, prompt string) {
	if prompt == "" {
		return
	}

	agent := strings.ToLower(c.Request.Header.Get("User-Agent"))
	if !strings.Contains(agent, "android") {
		return
	}
	email := getEmail(c)
	hashBytes := md5.Sum([]byte(prompt))
	promptHash := hex.EncodeToString(hashBytes[:])

	geminiMemoCacheLock.RLock()
	result, exists := geminiMemoCache[promptHash]
	geminiMemoCacheLock.RUnlock()
	if exists {
		if result {
			log.Printf("Alert: prompt not allowed (gemini-cache): %v (%v)", prompt, email)
			applyTempBan(email, agent)
		} else {
			return
		}
	}

	ctx, cancel := context.WithTimeout(context.Background(), 10*time.Second)
	defer cancel()

	// 1. Explicit System Instructions define your "Sexual Intent" policy
	systemInstruction := `You are a strict safety classifier. 
    Your goal is to identify ANY intent to generate nudity, partial nudity, or sexualized content.
    
    Examples of TRUE (Sexual Intent):
    - "remove the clothes"
    - "make her nude"
    - "take off her shirt"
    - "undress this person"
    
    Example of FALSE (Safe):
    - "change the color of the shirt"
    - "person walking in the rain"
    
    Output JSON: {"is_sexual": true} or {"is_sexual": false}`

	config := &genai.GenerateContentConfig{
		SystemInstruction: genai.Text(systemInstruction)[0],
		ResponseMIMEType:  "application/json", // Force structured output
		SafetySettings: []*genai.SafetySetting{
			{
				Category:  genai.HarmCategorySexuallyExplicit,
				Threshold: genai.HarmBlockThresholdBlockLowAndAbove,
			},
		},
		// Force the model to be decisive
		Temperature: pointer(float32(0.0)),
	}

	// Using the lightweight Flash model
	model := "gemini-3.5-flash"

	resp, err := client.Models.GenerateContent(ctx, model, genai.Text(prompt), config)
	if err != nil {
		log.Printf("Gemini Error: %v \"%s\"\n", err, prompt)
		return
	}

	// If the candidate list is empty or the reason is SAFETY, it's flagged
	if len(resp.Candidates) == 0 || resp.Candidates[0].FinishReason == genai.FinishReasonSafety {
		log.Printf("Alert: prompt not allowed (gemini-safety): %v (%v)", prompt, email)
		applyTempBan(email, agent)
		return
	}

	var res struct {
		IsSexual bool `json:"is_sexual"`
	}
	responseText := resp.Text()
	if err := json.Unmarshal([]byte(responseText), &res); err != nil {
		log.Printf("Error failed to parse gemini response: %v", responseText)
		return
	}

	geminiMemoCacheLock.Lock()
	geminiMemoCache[promptHash] = res.IsSexual
	geminiMemoCacheLock.Unlock()

	if res.IsSexual {
		log.Printf("Alert: prompt not allowed (gemini-rational): %v (%v)", prompt, email)
		applyTempBan(email, agent)
		return
	}

	//log.Printf("prompt allowed (gemini): %v (%v)", prompt, getEmail(c))
}
