package ollama

import (
	"bytes"
	"encoding/json"
	"fmt"
	"io"
	"net/http"
	"os"
	"encoding/base64"

	"gitea.miguelmuniz.com/rogueking/junk2jive-server/internal/response"
	"github.com/joho/godotenv"
)

type Ollama struct {
	OllamaKey      string
	OllamaEndpoint string
	Image          base64.Encoding
}

// ChatMessage represents a message in the chat completion request
type ChatMessage struct {
	Role    string `json:"role"`
	Content string `json:"content"`
}

// FileInfo represents a file in the chat completion request
type FileInfo struct {
	Type string `json:"type"`
	ID   string `json:"id"`
}

// ChatCompletionRequest represents the request body for the chat completion API
type ChatCompletionRequest struct {
	Model    string        `json:"model"`
	Messages []ChatMessage `json:"messages"`
	Files    []FileInfo    `json:"files,omitempty"`
}

func init() {
	// Load .env file if it exists
	_ = godotenv.Load()
}

func NewOllama() *Ollama {
	// Get API endpoint from environment with default fallback
	endpoint := os.Getenv("OLLAMA_API_ENDPOINT")
	if endpoint == "" {
		endpoint = "http://localhost:3000/api/chat/completions"
	}

	return &Ollama{
		OllamaKey:      os.Getenv("OLLAMA_API_KEY"),
		OllamaEndpoint: endpoint,
	}
}

func (o *Ollama) SendRequest(userMessage string, fileID string) (string, error) {
    // Prepare request body
    reqBody := ChatCompletionRequest{
        Model: "gemma3:12b",
        Messages: []ChatMessage{
            {Role: "user", Content: userMessage},
        },
    }

    // Add file if fileID is provided
    if fileID != "" {
        reqBody.Files = []FileInfo{
            {Type: "file", ID: fileID},
        }
    }

    // Marshal the request body to JSON
    jsonBody, err := json.Marshal(reqBody)
    if err != nil {
        // Remove the call to RespondWithError and just return the error
        return "", fmt.Errorf("error marshalling request: %w", err)
    }

    // Create the HTTP request
    req, err := http.NewRequest("POST", o.OllamaEndpoint, bytes.NewBuffer(jsonBody))
    if err != nil {
        // Remove the call to RespondWithError and just return the error
        return "", fmt.Errorf("error creating request: %w", err)
    }

    // Add headers
    req.Header.Set("Content-Type", "application/json")
    req.Header.Set("Authorization", "Bearer "+o.OllamaKey)

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

    // Read the response
    body, err := io.ReadAll(resp.Body)
    if err != nil {
        // Remove the call to RespondWithError and just return the error
        return "", fmt.Errorf("error reading response: %w", err)
    }

    // Check for non-200 status code
    if resp.StatusCode != http.StatusOK {
        // Remove the call to RespondWithError and just return the error
        return "", fmt.Errorf("API request failed with status %d: %s", resp.StatusCode, string(body))
    }

    return string(body), nil
}

func OllamaRequest(w http.ResponseWriter, r *http.Request) {
    // Parse request body
    var requestData struct {
        Message string `json:"message"`
        FileID  string `json:"fileId,omitempty"`
    }

    err := json.NewDecoder(r.Body).Decode(&requestData)
    if err != nil {
        response.RespondWithError(w, r, http.StatusBadRequest, "Invalid request body", err)
        return
    }

    // Create ollama client using API key from environment
    ollama := NewOllama()
    
    // Send request to Ollama API
    apiResponse, err := ollama.SendRequest(requestData.Message, requestData.FileID)
    if err != nil {
        response.RespondWithError(w, r, http.StatusInternalServerError, "Error sending request to Ollama API", err)
        return
    }

    // Return response
    response.RespondWithJSON(w, http.StatusOK, map[string]string{"response": apiResponse})
}