package ollama

import (
	"bytes"
	"encoding/json"
	"fmt"
	"net/http"
	"net/http/httptest"
	"os"
	"strings"
	"testing"
)

type MockOllama struct {
    SendRequestFunc func(userMessage string, fileID string) (string, error)
}

func (m *MockOllama) SendRequest(userMessage string, fileID string) (string, error) {
    return m.SendRequestFunc(userMessage, fileID)
}

func TestNewOllama(t *testing.T) {
	// Save original environment and restore it after test
	origEndpoint := os.Getenv("OLLAMA_API_ENDPOINT")
	origKey := os.Getenv("OLLAMA_API_KEY")
	defer func() {
		os.Setenv("OLLAMA_API_ENDPOINT", origEndpoint)
		os.Setenv("OLLAMA_API_KEY", origKey)
	}()

	// Test with empty environment
	os.Unsetenv("OLLAMA_API_ENDPOINT")
	os.Setenv("OLLAMA_API_KEY", "test-key")

	ollama := NewOllama()
	if ollama.OllamaEndpoint != "http://localhost:3000/api/chat/completions" {
		t.Errorf("Expected default endpoint, got %s", ollama.OllamaEndpoint)
	}
	if ollama.OllamaKey != "test-key" {
		t.Errorf("Expected API key 'test-key', got %s", ollama.OllamaKey)
	}

	// Test with custom environment
	os.Setenv("OLLAMA_API_ENDPOINT", "https://custom-endpoint.com")
	os.Setenv("OLLAMA_API_KEY", "custom-key")

	ollama = NewOllama()
	if ollama.OllamaEndpoint != "https://custom-endpoint.com" {
		t.Errorf("Expected custom endpoint, got %s", ollama.OllamaEndpoint)
	}
	if ollama.OllamaKey != "custom-key" {
		t.Errorf("Expected API key 'custom-key', got %s", ollama.OllamaKey)
	}
}

func TestSendRequest(t *testing.T) {
	// Create mock server
	server := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
		// Verify headers
		if r.Header.Get("Content-Type") != "application/json" {
			t.Errorf("Expected Content-Type header to be application/json")
		}
		if r.Header.Get("Authorization") != "Bearer test-key" {
			t.Errorf("Expected Authorization header to be 'Bearer test-key'")
		}

		// Verify request body
		var reqBody ChatCompletionRequest
		if err := json.NewDecoder(r.Body).Decode(&reqBody); err != nil {
			t.Errorf("Error decoding request body: %v", err)
		}

		if reqBody.Model != "gemma3:12b" {
			t.Errorf("Expected model 'gemma3:12b', got %s", reqBody.Model)
		}
		if len(reqBody.Messages) != 1 {
			t.Errorf("Expected 1 message, got %d", len(reqBody.Messages))
		}
		if reqBody.Messages[0].Role != "user" {
			t.Errorf("Expected role 'user', got %s", reqBody.Messages[0].Role)
		}

		// Check if fileID was properly included
		if r.URL.Query().Get("fileId") == "test-file" {
			if len(reqBody.Files) != 1 {
				t.Errorf("Expected 1 file, got %d", len(reqBody.Files))
			}
		}

		// Return mock response
		w.Header().Set("Content-Type", "application/json")
		w.WriteHeader(http.StatusOK)
		w.Write([]byte(`{"response": "This is a test response"}`))
	}))
	defer server.Close()

	// Create Ollama client with mock server URL
	ollama := &Ollama{
		OllamaKey:      "test-key",
		OllamaEndpoint: server.URL,
	}

	// Test without fileID
	resp, err := ollama.SendRequest("Hello", "")
	if err != nil {
		t.Errorf("SendRequest failed: %v", err)
	}
	if resp != `{"response": "This is a test response"}` {
		t.Errorf("Unexpected response: %s", resp)
	}

	// Test with fileID
	resp, err = ollama.SendRequest("Hello with file", "test-file")
	if err != nil {
		t.Errorf("SendRequest with fileID failed: %v", err)
	}
}

func TestSendRequestError(t *testing.T) {
	// Create mock server that returns an error
	server := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
		w.WriteHeader(http.StatusInternalServerError)
		w.Write([]byte("Internal Server Error"))
	}))
	defer server.Close()

	// Create Ollama client with mock server URL
	ollama := &Ollama{
		OllamaKey:      "test-key",
		OllamaEndpoint: server.URL,
	}

	// Test error handling
	_, err := ollama.SendRequest("Hello", "")
	if err == nil {
		t.Error("Expected error, got nil")
	}
}

// TestOllamaRequestHandlerIntegration tests the complete flow of the OllamaRequest handler
func TestOllamaRequestHandlerIntegration(t *testing.T) {
	// Save original NewOllama function and restore it after the test
	originalNewOllama := NewOllama
	defer func() {
		NewOllama = originalNewOllama
	}()

	// Set up test cases
	testCases := []struct {
		name           string
		requestBody    string
		mockResponse   string
		mockError      error
		expectedStatus int
		expectedBody   string
	}{
		{
			name:           "Successful request",
			requestBody:    `{"message": "Hello, AI!", "fileId": "test-file"}`,
			mockResponse:   `{"choices":[{"message":{"content":"Hello, human!"}}]}`,
			mockError:      nil,
			expectedStatus: http.StatusOK,
			expectedBody:   `{"response":"{\"choices\":[{\"message\":{\"content\":\"Hello, human!\"}}]}"}`,
		},
		{
			name:           "Invalid request body",
			requestBody:    `{"message": }`, // Invalid JSON
			mockResponse:   "",
			mockError:      nil,
			expectedStatus: http.StatusBadRequest,
			expectedBody:   `{"error":"Invalid request body"}`,
		},
		{
			name:           "API Error",
			requestBody:    `{"message": "Hello"}`,
			mockResponse:   "",
			mockError:      fmt.Errorf("API error"),
			expectedStatus: http.StatusInternalServerError,
			expectedBody:   `{"error":"Error sending request to Ollama API"}`,
		},
	}

	for _, tc := range testCases {
		t.Run(tc.name, func(t *testing.T) {
			// Override the NewOllama function to return our mock
			NewOllama = func() *Ollama {
				return &Ollama{
					OllamaKey:      "mock-key",
					OllamaEndpoint: "mock-endpoint",
				}
			}

			// Create a mock SendRequest function
			originalSendRequest := (*Ollama).SendRequest
			defer func() {
				(*Ollama).SendRequest = originalSendRequest
			}()

			(*Ollama).SendRequest = func(o *Ollama, userMessage string, fileID string) (string, error) {
				return tc.mockResponse, tc.mockError
			}

			// Create a request with the test case body
			req := httptest.NewRequest(http.MethodPost, "/api/ollama", strings.NewReader(tc.requestBody))
			req.Header.Set("Content-Type", "application/json")

			// Create a recorder to capture the response
			w := httptest.NewRecorder()

			// Call the handler
			OllamaRequest(w, req)

			// Get the result
			resp := w.Result()
			defer resp.Body.Close()

			// Check the status code
			if resp.StatusCode != tc.expectedStatus {
				t.Errorf("Expected status %d, got %d", tc.expectedStatus, resp.StatusCode)
			}

			// Check the response body
			var respBody map[string]string
			if err := json.NewDecoder(resp.Body).Decode(&respBody); err != nil {
				t.Fatalf("Failed to decode response body: %v", err)
			}

			// For successful case, verify the response contains the expected content
			if tc.expectedStatus == http.StatusOK {
				if _, ok := respBody["response"]; !ok {
					t.Errorf("Response doesn't contain 'response' field")
				}
			} else {
				// For error cases, check if error message is included
				if _, ok := respBody["error"]; !ok {
					t.Errorf("Error response doesn't contain 'error' field")
				}
			}
		})
	}
}

// TestEndToEndIntegration tests the actual API connection if environment variables are set
func TestEndToEndIntegration(t *testing.T) {
	// Skip this test unless explicitly enabled with an environment variable
	if os.Getenv("ENABLE_E2E_TESTS") != "true" {
		t.Skip("Skipping end-to-end test; set ENABLE_E2E_TESTS=true to run")
	}

	// Create a real request
	reqBody := map[string]string{
		"message": "Hello, this is a test message",
	}
	jsonBody, err := json.Marshal(reqBody)
	if err != nil {
		t.Fatalf("Failed to marshal request body: %v", err)
	}

	// Create a request to our handler
	req := httptest.NewRequest(http.MethodPost, "/api/ollama", bytes.NewBuffer(jsonBody))
	req.Header.Set("Content-Type", "application/json")

	// Create a recorder
	w := httptest.NewRecorder()

	// Call the actual handler (which will make a real API call)
	OllamaRequest(w, req)

	// Check response
	resp := w.Result()
	defer resp.Body.Close()

	// Basic assertions
	if resp.StatusCode != http.StatusOK {
		t.Errorf("Expected status OK, got %v", resp.Status)

		// Print error body for debugging
		var errorResp map[string]interface{}
		if err := json.NewDecoder(resp.Body).Decode(&errorResp); err == nil {
			t.Logf("Error response: %v", errorResp)
		}
	} else {
		var respBody map[string]interface{}
		if err := json.NewDecoder(resp.Body).Decode(&respBody); err != nil {
			t.Errorf("Failed to decode response body: %v", err)
		} else {
			t.Logf("Successful API response received: %v", respBody)
		}
	}
}
