Goondex/internal/ml/analysis.go

package ml

import (
	"context"
	"database/sql"
	"encoding/json"
	"fmt"
	"log"
	"time"

	"git.leaktechnologies.dev/stu/Goondex/internal/db"
	"git.leaktechnologies.dev/stu/Goondex/internal/model"
)

// ScenePrediction represents ML prediction data for a scene
type ScenePrediction struct {
	ID             int64              `json:"id"`
	PredictionType string             `json:"prediction_type"`
	Predictions    map[string]float64 `json:"predictions"` // tag -> confidence
	OverallScore   float64            `json:"overall_score"`
	Model          string             `json:"model"`
	Confidence     float64            `json:"confidence"`
	CreatedAt      interface{}        `json:"created_at"`
	UpdatedAt      interface{}        `json:"updated_at"`
}

// MLAnalysisService handles ML-powered scene analysis
type MLAnalysisService struct {
	db *db.DB
}

// NewMLAnalysisService creates a new ML service
func NewMLAnalysisService(database *db.DB) *MLAnalysisService {
	return &MLAnalysisService{
		db: database,
	}
}

// AnalyzeScene runs ML analysis on a scene and stores results
func (ml *MLAnalysisService) AnalyzeScene(ctx context.Context, sceneID int64, imageData []byte, modelVersion string) (*ScenePrediction, error) {
	// For now, simulate ML analysis based on basic image processing
	// In a real implementation, this would call your ML model

	// Simulate detecting various attributes
	predictions := make(map[string]float64)

	// Detect hair-related attributes (based on your requirements)
	predictions["shaved"] = ml.analyzeHairStyle(imageData)
	predictions["natural_hair"] = ml.analyzeHairStyle(imageData)
	predictions["bushy"] = ml.analyzeHairStyle(imageData)

	// Detect gender attributes
	predictions["male"] = ml.analyzeGender(imageData)
	predictions["circumcised"] = ml.analyzeCircumcision(imageData)

	// Detect body attributes
	predictions["athletic"] = ml.analyzeBodyType(imageData, "athletic")
	predictions["slim"] = ml.analyzeBodyType(imageData, "slim")
	predictions["curvy"] = ml.analyzeBodyType(imageData, "curvy")
	predictions["bbw"] = ml.analyzeBodyType(imageData, "bbw")

	// Detect age categories
	predictions["teen"] = ml.analyzeAgeCategory(imageData, "teen")
	predictions["milf"] = ml.analyzeAgeCategory(imageData, "milf")
	predictions["mature"] = ml.analyzeAgeCategory(imageData, "mature")

	// Detect clothing
	predictions["pink_clothing"] = ml.analyzeClothingColor(imageData, "pink")
	predictions["black_clothing"] = ml.analyzeClothingColor(imageData, "black")
	predictions["red_clothing"] = ml.analyzeClothingColor(imageData, "red")
	predictions["blue_clothing"] = ml.analyzeClothingColor(imageData, "blue")
	predictions["white_clothing"] = ml.analyzeClothingColor(imageData, "white")
	predictions["thong"] = ml.analyzeClothingType(imageData, "thong")
	predictions["panties"] = ml.analyzeClothingType(imageData, "panties")
	predictions["lingerie"] = ml.analyzeClothingType(imageData, "lingerie")
	predictions["dress"] = ml.analyzeClothingType(imageData, "dress")
	predictions["skirt"] = ml.analyzeClothingType(imageData, "skirt")
	predictions["heels"] = ml.analyzeClothingType(imageData, "heels")
	predictions["boots"] = ml.analyzeClothingType(imageData, "boots")
	predictions["stockings"] = ml.analyzeClothingType(imageData, "stockings")

	// Detect actions/positions
	predictions["creampie"] = ml.analyzeSexualAct(imageData, "creampie")
	predictions["blowjob"] = ml.analyzeSexualAct(imageData, "blowjob")
	predictions["cowgirl"] = ml.analyzePosition(imageData, "cowgirl")
	predictions["doggy"] = ml.analyzePosition(imageData, "doggy")

	// Detect settings
	predictions["bedroom"] = ml.analyzeSetting(imageData, "bedroom")
	predictions["couch"] = ml.analyzeSetting(imageData, "couch")
	predictions["office"] = ml.analyzeSetting(imageData, "office")
	predictions["kitchen"] = ml.analyzeSetting(imageData, "kitchen")
	predictions["bathroom"] = ml.analyzeSetting(imageData, "bathroom")
	predictions["car"] = ml.analyzeSetting(imageData, "car")
	predictions["outdoor"] = ml.analyzeSetting(imageData, "outdoor")

	// Detect objects/furniture
	predictions["sofa"] = ml.analyzeObject(imageData, "sofa")
	predictions["bed"] = ml.analyzeObject(imageData, "bed")
	predictions["table"] = ml.analyzeObject(imageData, "table")

	// Calculate overall confidence score
	overallScore := ml.calculateOverallScore(predictions)

	prediction := &ScenePrediction{
		PredictionType: "comprehensive",
		Predictions:    predictions,
		OverallScore:   overallScore,
		Model:          modelVersion,
		Confidence:     overallScore,
	}

	// Store analysis results
	if err := ml.storeSceneAnalysis(ctx, sceneID, prediction); err != nil {
		return nil, fmt.Errorf("failed to store scene analysis: %w", err)
	}

	log.Printf("ML analysis complete for scene %d: overall score %.2f, %d predictions",
		sceneID, overallScore, len(predictions))

	return prediction, nil
}

// GetSceneAnalysis retrieves stored ML analysis for a scene
func (ml *MLAnalysisService) GetSceneAnalysis(ctx context.Context, sceneID int64) ([]ScenePrediction, error) {
	rows, err := ml.db.Conn().Query(`
		SELECT id, model_version, prediction_type, predictions, confidence_score, created_at, updated_at
		FROM scene_ml_analysis
		WHERE scene_id = ?
		ORDER BY created_at DESC
	`, sceneID)

	if err != nil {
		return nil, fmt.Errorf("failed to retrieve scene analysis: %w", err)
	}
	defer rows.Close()

	var predictions []ScenePrediction
	for rows.Next() {
		var prediction ScenePrediction
		var predictionsJSON string
		var createdAt, updatedAt string

		err := rows.Scan(
			&prediction.ID, &prediction.Model, &prediction.PredictionType,
			&predictionsJSON, &prediction.OverallScore, &prediction.Confidence,
			&createdAt, &updatedAt,
		)

		if err != nil {
			continue
		}

		// Parse predictions JSON
		if err := json.Unmarshal([]byte(predictionsJSON), &prediction.Predictions); err != nil {
			continue
		}

		// Parse timestamps (for now, store as strings)
		prediction.CreatedAt = parseTime(createdAt)
		prediction.UpdatedAt = parseTime(updatedAt)

		predictions = append(predictions, prediction)
	}

	return predictions, nil
}

// UpdateSceneTags applies ML predictions to scene_tags table
func (ml *MLAnalysisService) UpdateSceneTags(ctx context.Context, sceneID int64, minConfidence float64) error {
	predictions, err := ml.GetSceneAnalysis(ctx, sceneID)
	if err != nil {
		return fmt.Errorf("failed to get scene analysis: %w", err)
	}

	if len(predictions) == 0 {
		return nil
	}

	// Get the latest high-confidence predictions
	latest := predictions[0]
	for _, prediction := range predictions {
		if prediction.Confidence > latest.Confidence {
			latest = prediction
		}
	}

	// Apply predictions to scene_tags table
	tagStore := db.NewTagStore(ml.db)

	for tagName, confidence := range latest.Predictions {
		if confidence < minConfidence {
			continue // Skip low-confidence predictions
		}

		// Find or create the tag
		tag, err := tagStore.FindOrCreate(tagName, "ml")
		if err != nil {
			log.Printf("Failed to find/create tag %s: %v", tagName, err)
			continue
		}

		// Link tag to scene with ML source and confidence
		if err := ml.linkSceneToTag(ctx, sceneID, tag.ID, confidence, "ml"); err != nil {
			log.Printf("Failed to link scene %d to tag %d: %v", sceneID, tag.ID, err)
		}
	}

	log.Printf("Applied %d ML predictions to scene %d", len(latest.Predictions), sceneID)
	return nil
}

// Mock ML analysis functions (replace with real ML model calls)
func (ml *MLAnalysisService) analyzeHairStyle(imageData []byte) float64 {
	// Simulate hair style analysis
	return 0.7 // Mock confidence
}

func (ml *MLAnalysisService) analyzeGender(imageData []byte) float64 {
	// Simulate gender analysis
	return 0.8 // Mock confidence
}

func (ml *MLAnalysisService) analyzeCircumcision(imageData []byte) float64 {
	// Simulate circumcision detection
	return 0.6 // Mock confidence
}

func (ml *MLAnalysisService) analyzeBodyType(imageData []byte, bodyType string) float64 {
	// Simulate body type analysis
	switch bodyType {
	case "athletic", "slim":
		return 0.8
	case "curvy":
		return 0.7
	case "bbw":
		return 0.9
	default:
		return 0.5
	}
}

func (ml *MLAnalysisService) analyzeAgeCategory(imageData []byte, ageCat string) float64 {
	// Simulate age category analysis
	switch ageCat {
	case "teen", "milf", "mature":
		return 0.9
	default:
		return 0.5
	}
}

func (ml *MLAnalysisService) analyzeClothingColor(imageData []byte, color string) float64 {
	// Simulate clothing color detection
	switch color {
	case "pink", "black", "red", "blue":
		return 0.9
	default:
		return 0.5
	}
}

func (ml *MLAnalysisService) analyzeClothingType(imageData []byte, clothingType string) float64 {
	// Simulate clothing type detection
	switch clothingType {
	case "thong", "heels":
		return 0.85
	case "stockings", "lingerie":
		return 0.75
	default:
		return 0.5
	}
}

func (ml *MLAnalysisService) analyzeSexualAct(imageData []byte, act string) float64 {
	// Simulate sexual act detection
	switch act {
	case "creampie", "blowjob", "cowgirl", "doggy":
		return 0.9
	default:
		return 0.5
	}
}

func (ml *MLAnalysisService) analyzePosition(imageData []byte, position string) float64 {
	// Simulate position detection
	switch position {
	case "cowgirl", "doggy":
		return 0.85
	default:
		return 0.5
	}
}

func (ml *MLAnalysisService) analyzeSetting(imageData []byte, setting string) float64 {
	// Simulate setting detection
	switch setting {
	case "bedroom", "couch":
		return 0.8
	case "office":
		return 0.6
	case "kitchen":
		return 0.6
	case "bathroom":
		return 0.6
	case "car":
		return 0.7
	case "outdoor":
		return 0.7
	default:
		return 0.5
	}
}

func (ml *MLAnalysisService) analyzeObject(imageData []byte, objectType string) float64 {
	// Simulate object detection
	switch objectType {
	case "sofa":
		return 0.8
	case "bed", "table":
		return 0.9
	default:
		return 0.5
	}
}

func (ml *MLAnalysisService) calculateOverallScore(predictions map[string]float64) float64 {
	if len(predictions) == 0 {
		return 0.0
	}

	total := 0.0
	count := 0

	for _, confidence := range predictions {
		total += confidence
		count++
	}

	// Weighted average with bonus for having multiple predictions
	average := total / float64(count)
	multiplier := 1.0 + (float64(count)-1.0)*0.1 // Bonus for comprehensive coverage

	return average * multiplier
}

func (ml *MLAnalysisService) storeSceneAnalysis(ctx context.Context, sceneID int64, prediction *ScenePrediction) error {
	predictionsJSON, err := json.Marshal(prediction.Predictions)
	if err != nil {
		return fmt.Errorf("failed to marshal predictions: %w", err)
	}

	_, err = ml.db.Conn().Exec(`
		INSERT INTO scene_ml_analysis (scene_id, model_version, prediction_type, predictions, confidence_score, created_at, updated_at)
		VALUES (?, ?, ?, ?, ?, datetime('now'), datetime('now'))
	`, sceneID, prediction.Model, prediction.PredictionType, predictionsJSON, prediction.OverallScore)

	return err
}

func (ml *MLAnalysisService) linkSceneToTag(ctx context.Context, sceneID, tagID int64, confidence float64, source string) error {
	_, err := ml.db.Conn().Exec(`
		INSERT OR REPLACE INTO scene_tags (scene_id, tag_id, confidence, source, verified, created_at)
		VALUES (?, ?, ?, ?, ?, 0, datetime('now'))
	`, sceneID, tagID, confidence, source)

	return err
}

func parseTime(timeStr string) interface{} {
	// For now, return as string. In real implementation, parse to time.Time
	return timeStr
}