package main

import (
	"bytes"
	"encoding/base64"
	"encoding/json"
	"flag"
	"fmt"
	"io"
	"io/ioutil"
	"log"
	"mime/multipart"
	"net/http"
	"os"
	"path/filepath"
	"regexp"
	"strings"
	"unicode"

	"gopkg.in/yaml.v3"
)

// ImageExtensionRegex is a regular expression to validate image file extensions.
var ImageExtensionRegex = regexp.MustCompile(`(?i)\.(jpg|jpeg|png|gif)$`)

// OpenAIRequest represents the structure of the request to the OpenAI API.
type OpenAIRequest struct {
	Model       string    `json:"model"`
	Messages    []Message `json:"messages"`
	Temperature float64   `json:"temperature,omitempty"`
}

// Message represents a message in the OpenAI API request.
type Message struct {
	Role    string    `json:"role"`
	Content []Content `json:"content"`
}

type Content struct {
	Type     string    `json:"type"`
	Text     string    `json:"text,omitempty"`
	ImageURL *ImageURL `json:"image_url,omitempty"`
}

type ImageURL struct {
	URL string `json:"url"`
}

// OpenAIResponse represents the structure of the response from the OpenAI API.
type OpenAIResponse struct {
	Choices []struct {
		Message struct {
			Content string `json:"content"`
		} `json:"message"`
	} `json:"choices"`
}

// Config represents the structure of the configuration file.
type Config struct {
	APIKey      string  `yaml:"apikey"`
	Model       string  `yaml:"model"`
	Endpoint    string  `yaml:"endpoint"`
	Temperature float64 `yaml:"temperature"`
}

func main() {
	var configFile string
	flag.StringVar(&configFile, "c", "", "Path to config file")
	flag.Parse()

	// Get the config file path
	configFilePath, err := getConfigFilePath(configFile)
	if err != nil {
		log.Fatalf("Error getting config file path: %v", err)
	}

	// Load the configuration from the file
	config, err := loadConfig(configFilePath)
	if err != nil {
		log.Fatalf("Error loading config: %v", err)
	}

	// Get the image filename from the command-line arguments
	if len(flag.Args()) < 1 {
		fmt.Println("Usage: gonamer [-c config.yaml] <image_filename>")
		flag.PrintDefaults()
		return
	}
	imageFilename := flag.Arg(0)

	// Validate the image filename extension
	if !ImageExtensionRegex.MatchString(imageFilename) {
		fmt.Println("Error: Invalid image file extension. Only .jpg, .jpeg, .png, and .gif are supported.")
		return
	}

	// Get the suggested filename from the LLM
	newFilename, err := getSuggestedFilename(imageFilename, config)
	if err != nil {
		fmt.Println("Error getting suggested filename:", err)
		return
	}

	// Rename the file
	err = renameFile(imageFilename, newFilename)
	if err != nil {
		fmt.Println("Error renaming file:", err)
		return
	}

	fmt.Printf("Successfully renamed '%s' to '%s'\n", imageFilename, newFilename)
}

// getConfigFilePath returns the path to the config file.
func getConfigFilePath(configFile string) (string, error) {
	if configFile != "" {
		return configFile, nil
	}
	homeDir, err := os.UserHomeDir()
	if err != nil {
		return "", err
	}
	return filepath.Join(homeDir, ".config", "gonamer.yaml"), nil
}

// loadConfig loads the configuration from the specified file.
func loadConfig(configFilePath string) (*Config, error) {
	data, err := ioutil.ReadFile(configFilePath)
	if err != nil {
		return nil, err
	}

	var config Config
	err = yaml.Unmarshal(data, &config)
	if err != nil {
		return nil, err
	}

	return &config, nil
}

// getSuggestedFilename sends the image to the LLM and returns the suggested filename.
func getSuggestedFilename(imageFilename string, config *Config) (string, error) {
	// 1. Encode the image to base64
	imageData, err := os.ReadFile(imageFilename)
	if err != nil {
		return "", fmt.Errorf("error reading image file: %w", err)
	}
	base64Image := base64.StdEncoding.EncodeToString(imageData)

	// 2. Create the request payload
	requestData := OpenAIRequest{
		Model: config.Model,
		Messages: []Message{
			{
				Role: "user",
				Content: []Content{
					{
						Type: "text",
						Text: "Suggest a useful, descriptive, and unique filename for this image, considering its content and context. Only provide the filename, no other text or explanation. Do not include the file extension.",
					},
					{
						Type: "image_url",
						ImageURL: &ImageURL{
							URL: fmt.Sprintf("data:image/%s;base64,%s", strings.TrimPrefix(filepath.Ext(imageFilename), "."), base64Image),
						},
					},
				},
			},
		},
		Temperature: config.Temperature,
	}

	requestBody, err := json.Marshal(requestData)
	if err != nil {
		return "", fmt.Errorf("error marshalling request data: %w", err)
	}

	// 3. Send the request to the LLM (using JSON request as default)
	originalExtension := filepath.Ext(imageFilename)
	newFilename, err := sendJSONRequest(requestBody, originalExtension, config)
	if err != nil {
		return "", err
	}

	return newFilename, nil
}

// sendJSONRequest sends the request as JSON to the LLM endpoint
func sendJSONRequest(requestBody []byte, originalExtension string, config *Config) (string, error) {
	req, err := http.NewRequest("POST", config.Endpoint, bytes.NewBuffer(requestBody))
	if err != nil {
		return "", fmt.Errorf("error creating HTTP request: %w", err)
	}

	req.Header.Set("Content-Type", "application/json")
	req.Header.Set("Authorization", "Bearer "+config.APIKey)

	client := &http.Client{}
	resp, err := client.Do(req)
	if err != nil {
		return "", fmt.Errorf("error sending HTTP request: %w", err)
	}
	defer resp.Body.Close()

	// 4. Process the response
	responseBody, err := io.ReadAll(resp.Body)
	if err != nil {
		return "", fmt.Errorf("error reading response body: %w", err)
	}

	// Print the raw response body
	//fmt.Println("Raw Response Body:")
	//fmt.Println(string(responseBody))

	if resp.StatusCode != http.StatusOK {
		return "", fmt.Errorf("API error (status %d): %s", resp.StatusCode, responseBody)
	}

	var apiResponse OpenAIResponse
	err = json.Unmarshal(responseBody, &apiResponse)
	if err != nil {
		return "", fmt.Errorf("error unmarshalling response: %w", err)
	}

	if len(apiResponse.Choices) == 0 || apiResponse.Choices[0].Message.Content == "" {
		return "", fmt.Errorf("empty response from API")
	}

	suggestedFilename := strings.TrimSpace(apiResponse.Choices[0].Message.Content)

	// Remove any existing extension from the suggested filename
	suggestedFilename = strings.TrimSuffix(suggestedFilename, filepath.Ext(suggestedFilename))

	// Ensure the suggested filename has the correct extension
	if !strings.HasSuffix(strings.ToLower(suggestedFilename), strings.ToLower(originalExtension)) {
		suggestedFilename += originalExtension
	}

	return suggestedFilename, nil
}

// sendMultipartRequest sends the request as multipart/form-data to the LLM endpoint (alternative)
func sendMultipartRequest(requestBody []byte, originalExtension string, config *Config) (string, error) {
	// Prepare a form
	body := &bytes.Buffer{}
	writer := multipart.NewWriter(body)

	// Add the JSON part
	jsonPart, err := writer.CreateFormField("json")
	if err != nil {
		return "", fmt.Errorf("error creating form field: %w", err)
	}
	_, err = jsonPart.Write(requestBody)
	if err != nil {
		return "", fmt.Errorf("error writing to form field: %w", err)
	}

	// Close the multipart writer
	err = writer.Close()
	if err != nil {
		return "", fmt.Errorf("error closing multipart writer: %w", err)
	}

	// Create a new request
	req, err := http.NewRequest("POST", config.Endpoint, body)
	if err != nil {
		return "", fmt.Errorf("error creating HTTP request: %w", err)
	}

	// Set headers
	req.Header.Set("Content-Type", writer.FormDataContentType())
	req.Header.Set("Authorization", "Bearer "+config.APIKey)

	// Send the request
	client := &http.Client{}
	resp, err := client.Do(req)
	if err != nil {
		return "", fmt.Errorf("error sending HTTP request: %w", err)
	}
	defer resp.Body.Close()

	// Process the response (same as in sendJSONRequest)
	responseBody, err := io.ReadAll(resp.Body)
	if err != nil {
		return "", fmt.Errorf("error reading response body: %w", err)
	}

	// Print the raw response body
	//fmt.Println("Raw Response Body:")
	//fmt.Println(string(responseBody))

	if resp.StatusCode != http.StatusOK {
		return "", fmt.Errorf("API error (status %d): %s", resp.StatusCode, responseBody)
	}

	var apiResponse OpenAIResponse
	err = json.Unmarshal(responseBody, &apiResponse)
	if err != nil {
		return "", fmt.Errorf("error unmarshalling response: %w", err)
	}

	if len(apiResponse.Choices) == 0 || apiResponse.Choices[0].Message.Content == "" {
		return "", fmt.Errorf("empty response from API")
	}

	suggestedFilename := strings.TrimSpace(apiResponse.Choices[0].Message.Content)

	// Remove any existing extension from the suggested filename
	suggestedFilename = strings.TrimSuffix(suggestedFilename, filepath.Ext(suggestedFilename))

	// Ensure the suggested filename has the correct extension
	if !strings.HasSuffix(strings.ToLower(suggestedFilename), strings.ToLower(originalExtension)) {
		suggestedFilename += originalExtension
	}

	return suggestedFilename, nil
}

// renameFile renames the file, handling potential issues like existing files.
func renameFile(oldFilename, newFilename string) error {
	// Sanitize the new filename (remove invalid characters, etc.)
	newFilename = sanitizeFilename(newFilename)

	// Check if a file with the new name already exists
	if _, err := os.Stat(newFilename); err == nil {
		// If it exists, generate a unique name (e.g., by adding a number)
		newFilename = makeUniqueFilename(newFilename)
	}

	// Perform the rename operation
	return os.Rename(oldFilename, newFilename)
}

// sanitizeFilename removes invalid characters from a filename and ensures the first character is a letter.
func sanitizeFilename(filename string) string {
	// Extract the extension
	ext := filepath.Ext(filename)
	filenameWithoutExt := filename[:len(filename)-len(ext)]

	// Define a regular expression to match invalid filename characters
	invalidChars := regexp.MustCompile(`[\\/:*?"<>|]`)

	// Replace invalid characters with underscores
	sanitizedFilename := invalidChars.ReplaceAllString(filenameWithoutExt, "_")

	// Remove non-letter characters from the beginning of the filename
	for len(sanitizedFilename) > 0 && !unicode.IsLetter(rune(sanitizedFilename[0])) {
		sanitizedFilename = sanitizedFilename[1:]
	}

	// Ensure the first character is a letter
	if len(sanitizedFilename) == 0 || !unicode.IsLetter(rune(sanitizedFilename[0])) {
		sanitizedFilename = "A_" + sanitizedFilename
	}

	// Remove leading dots and underscores
	for len(sanitizedFilename) > 0 && (sanitizedFilename[0] == '.' || sanitizedFilename[0] == '_') {
		sanitizedFilename = sanitizedFilename[1:]
	}

	// If the filename is now empty (e.g., it was only dots or underscores), set it to "A"
	if len(sanitizedFilename) == 0 {
		sanitizedFilename = "A"
	}

	return sanitizedFilename + ext
}

// makeUniqueFilename adds a number to the filename to make it unique.
func makeUniqueFilename(filename string) string {
	ext := filepath.Ext(filename)
	name := filename[:len(filename)-len(ext)]

	counter := 1
	newFilename := filename
	for {
		if _, err := os.Stat(newFilename); os.IsNotExist(err) {
			return newFilename
		}
		newFilename = fmt.Sprintf("%s_%d%s", name, counter, ext)
		counter++
	}
}