package proxy

import (
	"bytes"
	"encoding/json"
	"fmt"
	"io"
	"net/http"
	"net/http/httptest"
	"strings"
	"testing"
	"time"

	"github.com/Riotpiaole/homelab-frontend/internal/config"
)

// OpenAI-style tool definition
type Tool struct {
	Type     string       `json:"type"`
	Function ToolFunction `json:"function"`
}

type ToolFunction struct {
	Name        string                 `json:"name"`
	Description string                 `json:"description"`
	Parameters  map[string]interface{} `json:"parameters"`
}

// OpenAI chat completion with tools request
type ChatCompletionRequest struct {
	Model    string    `json:"model"`
	Messages []Message `json:"messages"`
	Tools    []Tool    `json:"tools,omitempty"`
	Stream   bool      `json:"stream,omitempty"`
}

type Message struct {
	Role      string      `json:"role"`
	Content   interface{} `json:"content"`
	ToolCalls []ToolCall  `json:"tool_calls,omitempty"`
}

type ToolCall struct {
	ID       string       `json:"id"`
	Type     string       `json:"type"`
	Function FunctionCall `json:"function"`
}

type FunctionCall struct {
	Name      string `json:"name"`
	Arguments string `json:"arguments"`
}

// Test OpenAI-style tool calling (GET /v1/chat/completions with tools)
func TestToolCallOpenAIStyle(t *testing.T) {
	upstreamReceived := false
	var receivedBody []byte

	upstreamServer := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
		upstreamReceived = true
		receivedBody, _ = io.ReadAll(r.Body)

		w.Header().Set("Content-Type", "application/json")
		w.WriteHeader(http.StatusOK)

		response := map[string]interface{}{
			"id":     "chatcmpl-123",
			"object": "chat.completion",
			"model":  "reasoning",
			"choices": []map[string]interface{}{
				{
					"index": 0,
					"message": map[string]interface{}{
						"role":    "assistant",
						"content": nil,
						"tool_calls": []map[string]interface{}{
							{
								"id":   "call_abc123",
								"type": "function",
								"function": map[string]interface{}{
									"name":      "get_weather",
									"arguments": `{"location":"San Francisco","unit":"celsius"}`,
								},
							},
						},
					},
					"finish_reason": "tool_calls",
				},
			},
		}

		json.NewEncoder(w).Encode(response)
	}))
	defer upstreamServer.Close()

	upstreamAddr := strings.TrimPrefix(upstreamServer.URL, "http://")

	cfg := &config.Config{
		Models: map[string]*config.ModelUpstream{
			"reasoning": {
				Name:    "reasoning",
				Address: upstreamAddr,
				Path:    "/v1/chat/completions",
			},
		},
		Routes: make(map[string]*config.Route),
	}

	handler := New(cfg)
	defer handler.Close()

	server := httptest.NewServer(handler)
	defer server.Close()

	// Create request with tools
	toolRequest := ChatCompletionRequest{
		Model: "reasoning",
		Messages: []Message{
			{
				Role:    "user",
				Content: "What's the weather in San Francisco?",
			},
		},
		Tools: []Tool{
			{
				Type: "function",
				Function: ToolFunction{
					Name:        "get_weather",
					Description: "Get weather for a location",
					Parameters: map[string]interface{}{
						"type": "object",
						"properties": map[string]interface{}{
							"location": map[string]interface{}{
								"type":        "string",
								"description": "City name",
							},
							"unit": map[string]interface{}{
								"type":        "string",
								"description": "Temperature unit",
								"enum":        []string{"celsius", "fahrenheit"},
							},
						},
						"required": []string{"location"},
					},
				},
			},
		},
	}

	body, _ := json.Marshal(toolRequest)
	resp, err := http.Post(server.URL+"/v1/chat/completions", "application/json", bytes.NewReader(body))
	if err != nil {
		t.Fatalf("request failed: %v", err)
	}
	defer resp.Body.Close()

	// Verify upstream received the request
	if !upstreamReceived {
		t.Errorf("expected upstream to receive tool call request")
	}

	// Verify the tool definition was forwarded
	var receivedRequest ChatCompletionRequest
	json.Unmarshal(receivedBody, &receivedRequest)
	if len(receivedRequest.Tools) != 1 {
		t.Errorf("expected 1 tool in upstream request, got %d", len(receivedRequest.Tools))
	}

	// Verify response contains tool_calls
	var respBody map[string]interface{}
	json.NewDecoder(resp.Body).Decode(&respBody)
	choices := respBody["choices"].([]interface{})
	choice := choices[0].(map[string]interface{})
	message := choice["message"].(map[string]interface{})
	toolCalls := message["tool_calls"].([]interface{})

	if len(toolCalls) == 0 {
		t.Errorf("expected tool_calls in response")
	}

	toolCall := toolCalls[0].(map[string]interface{})
	if toolCall["type"] != "function" {
		t.Errorf("expected tool_call type 'function', got %s", toolCall["type"])
	}
}

// Test tool calling with streaming (SSE)
func TestToolCallStreaming(t *testing.T) {
	upstreamServer := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
		// Parse request to check stream flag
		var req ChatCompletionRequest
		json.NewDecoder(r.Body).Decode(&req)

		if !req.Stream {
			t.Errorf("expected stream=true in request")
		}

		w.Header().Set("Content-Type", "text/event-stream")
		w.Header().Set("Cache-Control", "no-cache")
		w.WriteHeader(http.StatusOK)

		rc := http.NewResponseController(w)

		// Stream multiple chunks with tool_calls
		chunks := []map[string]interface{}{
			{
				"id":     "chatcmpl-123",
				"object": "chat.completion.chunk",
				"model":  "reasoning",
				"choices": []map[string]interface{}{
					{
						"index": 0,
						"delta": map[string]interface{}{
							"role":    "assistant",
							"content": nil,
						},
						"finish_reason": nil,
					},
				},
			},
			{
				"id":     "chatcmpl-123",
				"object": "chat.completion.chunk",
				"choices": []map[string]interface{}{
					{
						"index": 0,
						"delta": map[string]interface{}{
							"tool_calls": []map[string]interface{}{
								{
									"index": 0,
									"id":    "call_abc123",
									"type":  "function",
									"function": map[string]interface{}{
										"name":      "get_weather",
										"arguments": `{"location":"SF"}`,
									},
								},
							},
						},
						"finish_reason": nil,
					},
				},
			},
			{
				"id":     "chatcmpl-123",
				"object": "chat.completion.chunk",
				"choices": []map[string]interface{}{
					{
						"index":         0,
						"delta":         map[string]interface{}{},
						"finish_reason": "tool_calls",
					},
				},
			},
		}

		for _, chunk := range chunks {
			data, _ := json.Marshal(chunk)
			fmt.Fprintf(w, "data: %s\n\n", string(data))
			rc.Flush()
			time.Sleep(10 * time.Millisecond)
		}

		fmt.Fprint(w, "data: [DONE]\n\n")
		rc.Flush()
	}))
	defer upstreamServer.Close()

	upstreamAddr := strings.TrimPrefix(upstreamServer.URL, "http://")

	cfg := &config.Config{
		Models: map[string]*config.ModelUpstream{
			"reasoning": {
				Name:    "reasoning",
				Address: upstreamAddr,
			},
		},
		Routes: make(map[string]*config.Route),
	}

	handler := New(cfg)
	defer handler.Close()

	server := httptest.NewServer(handler)
	defer server.Close()

	toolRequest := ChatCompletionRequest{
		Model: "reasoning",
		Messages: []Message{
			{
				Role:    "user",
				Content: "Call get_weather",
			},
		},
		Tools: []Tool{
			{
				Type: "function",
				Function: ToolFunction{
					Name:        "get_weather",
					Description: "Get weather",
					Parameters:  map[string]interface{}{},
				},
			},
		},
		Stream: true,
	}

	body, _ := json.Marshal(toolRequest)
	resp, err := http.Post(server.URL+"/v1/chat/completions", "application/json", bytes.NewReader(body))
	if err != nil {
		t.Fatalf("request failed: %v", err)
	}
	defer resp.Body.Close()

	// Read streaming response
	respBody, _ := io.ReadAll(resp.Body)
	respStr := string(respBody)

	// Verify we got tool_calls in the stream
	if !strings.Contains(respStr, "get_weather") {
		t.Errorf("expected 'get_weather' in streaming response")
	}

	if !strings.Contains(respStr, "[DONE]") {
		t.Errorf("expected [DONE] sentinel in streaming response")
	}
}

// Test multi-turn conversation with tool use and results
func TestToolCallMultiTurn(t *testing.T) {
	upstreamServer := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
		var req ChatCompletionRequest
		json.NewDecoder(r.Body).Decode(&req)

		// Check if request includes tool result from previous assistant call
		hasToolResult := false
		for _, msg := range req.Messages {
			if msg.Role == "tool" {
				hasToolResult = true
				break
			}
		}

		w.Header().Set("Content-Type", "application/json")
		w.WriteHeader(http.StatusOK)

		if hasToolResult {
			// Second turn: assistant responds with final answer
			response := map[string]interface{}{
				"id": "chatcmpl-456",
				"choices": []map[string]interface{}{
					{
						"index": 0,
						"message": map[string]interface{}{
							"role":    "assistant",
							"content": "The weather in San Francisco is 22°C and cloudy.",
						},
						"finish_reason": "stop",
					},
				},
			}
			json.NewEncoder(w).Encode(response)
		} else {
			// First turn: assistant calls tool
			response := map[string]interface{}{
				"id": "chatcmpl-123",
				"choices": []map[string]interface{}{
					{
						"index": 0,
						"message": map[string]interface{}{
							"role":    "assistant",
							"content": nil,
							"tool_calls": []map[string]interface{}{
								{
									"id":   "call_abc123",
									"type": "function",
									"function": map[string]interface{}{
										"name":      "get_weather",
										"arguments": `{"location":"San Francisco"}`,
									},
								},
							},
						},
						"finish_reason": "tool_calls",
					},
				},
			}
			json.NewEncoder(w).Encode(response)
		}
	}))
	defer upstreamServer.Close()

	upstreamAddr := strings.TrimPrefix(upstreamServer.URL, "http://")

	cfg := &config.Config{
		Models: map[string]*config.ModelUpstream{
			"reasoning": {
				Name:    "reasoning",
				Address: upstreamAddr,
			},
		},
		Routes: make(map[string]*config.Route),
	}

	handler := New(cfg)
	defer handler.Close()

	server := httptest.NewServer(handler)
	defer server.Close()

	// Turn 1: User asks question
	turn1 := ChatCompletionRequest{
		Model: "reasoning",
		Messages: []Message{
			{
				Role:    "user",
				Content: "What's the weather in San Francisco?",
			},
		},
		Tools: []Tool{
			{
				Type: "function",
				Function: ToolFunction{
					Name:        "get_weather",
					Description: "Get weather for a location",
					Parameters:  map[string]interface{}{},
				},
			},
		},
	}

	body1, _ := json.Marshal(turn1)
	resp1, err := http.Post(server.URL+"/v1/chat/completions", "application/json", bytes.NewReader(body1))
	if err != nil {
		t.Fatalf("request failed: %v", err)
	}
	var turn1Resp map[string]interface{}
	json.NewDecoder(resp1.Body).Decode(&turn1Resp)
	resp1.Body.Close()

	// Verify first turn has tool_calls
	choices1 := turn1Resp["choices"].([]interface{})
	msg1 := choices1[0].(map[string]interface{})["message"].(map[string]interface{})
	if msg1["tool_calls"] == nil {
		t.Errorf("expected tool_calls in first turn")
	}

	// Turn 2: Send tool result
	turn2 := ChatCompletionRequest{
		Model: "reasoning",
		Messages: []Message{
			{
				Role:    "user",
				Content: "What's the weather in San Francisco?",
			},
			{
				Role:      "assistant",
				ToolCalls: []ToolCall{{ID: "call_abc123", Type: "function", Function: FunctionCall{Name: "get_weather", Arguments: `{"location":"San Francisco"}`}}},
			},
			{
				Role:    "tool",
				Content: `{"temperature":22,"condition":"cloudy"}`,
			},
		},
		Tools: []Tool{
			{
				Type: "function",
				Function: ToolFunction{
					Name:        "get_weather",
					Description: "Get weather for a location",
					Parameters:  map[string]interface{}{},
				},
			},
		},
	}

	body2, _ := json.Marshal(turn2)
	resp2, err := http.Post(server.URL+"/v1/chat/completions", "application/json", bytes.NewReader(body2))
	if err != nil {
		t.Fatalf("request failed: %v", err)
	}
	var turn2Resp map[string]interface{}
	json.NewDecoder(resp2.Body).Decode(&turn2Resp)
	resp2.Body.Close()

	// Verify second turn has final answer (no tool_calls)
	choices2 := turn2Resp["choices"].([]interface{})
	msg2 := choices2[0].(map[string]interface{})["message"].(map[string]interface{})
	content := msg2["content"].(string)

	if !strings.Contains(content, "cloudy") || !strings.Contains(content, "22") {
		t.Errorf("expected final answer with weather info, got: %s", content)
	}
}

// Test parallel tool calls (multiple tools at once)
func TestParallelToolCalls(t *testing.T) {
	upstreamServer := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
		w.Header().Set("Content-Type", "application/json")
		w.WriteHeader(http.StatusOK)

		response := map[string]interface{}{
			"id": "chatcmpl-789",
			"choices": []map[string]interface{}{
				{
					"index": 0,
					"message": map[string]interface{}{
						"role":    "assistant",
						"content": nil,
						"tool_calls": []map[string]interface{}{
							{
								"id":   "call_1",
								"type": "function",
								"function": map[string]interface{}{
									"name":      "get_weather",
									"arguments": `{"location":"New York"}`,
								},
							},
							{
								"id":   "call_2",
								"type": "function",
								"function": map[string]interface{}{
									"name":      "get_weather",
									"arguments": `{"location":"London"}`,
								},
							},
							{
								"id":   "call_3",
								"type": "function",
								"function": map[string]interface{}{
									"name":      "get_weather",
									"arguments": `{"location":"Tokyo"}`,
								},
							},
						},
					},
					"finish_reason": "tool_calls",
				},
			},
		}

		json.NewEncoder(w).Encode(response)
	}))
	defer upstreamServer.Close()

	upstreamAddr := strings.TrimPrefix(upstreamServer.URL, "http://")

	cfg := &config.Config{
		Models: map[string]*config.ModelUpstream{
			"reasoning": {
				Name:    "reasoning",
				Address: upstreamAddr,
			},
		},
		Routes: make(map[string]*config.Route),
	}

	handler := New(cfg)
	defer handler.Close()

	server := httptest.NewServer(handler)
	defer server.Close()

	request := ChatCompletionRequest{
		Model: "reasoning",
		Messages: []Message{
			{
				Role:    "user",
				Content: "Get weather for New York, London, and Tokyo",
			},
		},
		Tools: []Tool{
			{
				Type: "function",
				Function: ToolFunction{
					Name:        "get_weather",
					Description: "Get weather for a location",
					Parameters:  map[string]interface{}{},
				},
			},
		},
	}

	body, _ := json.Marshal(request)
	resp, err := http.Post(server.URL+"/v1/chat/completions", "application/json", bytes.NewReader(body))
	if err != nil {
		t.Fatalf("request failed: %v", err)
	}
	var respData map[string]interface{}
	json.NewDecoder(resp.Body).Decode(&respData)
	resp.Body.Close()

	// Verify multiple tool calls
	choices := respData["choices"].([]interface{})
	msg := choices[0].(map[string]interface{})["message"].(map[string]interface{})
	toolCalls := msg["tool_calls"].([]interface{})

	if len(toolCalls) != 3 {
		t.Errorf("expected 3 parallel tool calls, got %d", len(toolCalls))
	}

	// Verify each call is present
	callNames := []string{}
	for _, call := range toolCalls {
		tc := call.(map[string]interface{})
		fn := tc["function"].(map[string]interface{})
		args := fn["arguments"].(string)
		callNames = append(callNames, args)
	}

	if !strings.Contains(strings.Join(callNames, ""), "New York") {
		t.Errorf("expected New York in parallel calls")
	}
	if !strings.Contains(strings.Join(callNames, ""), "London") {
		t.Errorf("expected London in parallel calls")
	}
	if !strings.Contains(strings.Join(callNames, ""), "Tokyo") {
		t.Errorf("expected Tokyo in parallel calls")
	}
}

// Test Anthropic-style tool use (different format)
func TestAnthropicToolUse(t *testing.T) {
	upstreamServer := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
		// Anthropic uses different request format
		var body map[string]interface{}
		json.NewDecoder(r.Body).Decode(&body)

		// Check for tools in request
		if _, hasTools := body["tools"]; !hasTools {
			t.Errorf("expected tools in Anthropic request")
		}

		w.Header().Set("Content-Type", "application/json")
		w.WriteHeader(http.StatusOK)

		// Anthropic response format with tool_use block
		response := map[string]interface{}{
			"id":   "msg_123",
			"type": "message",
			"role": "assistant",
			"content": []map[string]interface{}{
				{
					"type": "text",
					"text": "I'll check the weather for you.",
				},
				{
					"type": "tool_use",
					"id":   "toolu_123",
					"name": "get_weather",
					"input": map[string]interface{}{
						"location": "San Francisco",
						"unit":     "celsius",
					},
				},
			},
			"model":       "claude-3-5-sonnet",
			"stop_reason": "tool_use",
			"usage": map[string]interface{}{
				"input_tokens":  100,
				"output_tokens": 50,
			},
		}

		json.NewEncoder(w).Encode(response)
	}))
	defer upstreamServer.Close()

	upstreamAddr := strings.TrimPrefix(upstreamServer.URL, "http://")

	// For Anthropic, we need to configure a regular route (not body-based dispatch)
	// since Anthropic uses different path structure (/llm/v1/messages)
	cfg := &config.Config{
		Routes: map[string]*config.Route{
			"anthropic": {
				Name: "anthropic",
				Upstream: config.Upstream{
					Address:        upstreamAddr,
					PathRewrite:    "/v1/messages",
					ConnectTimeout: 5 * time.Second,
					ReadTimeout:    30 * time.Second,
					WriteTimeout:   30 * time.Second,
					MaxBodySize:    10 * 1024 * 1024,
					AuthRequired:   false,
				},
			},
		},
		Models: make(map[string]*config.ModelUpstream),
	}

	handler := New(cfg)
	defer handler.Close()

	server := httptest.NewServer(handler)
	defer server.Close()

	// Anthropic request format
	anthropicRequest := map[string]interface{}{
		"model": "claude",
		"messages": []map[string]interface{}{
			{
				"role":    "user",
				"content": "What's the weather?",
			},
		},
		"tools": []map[string]interface{}{
			{
				"name":        "get_weather",
				"description": "Get weather information",
				"input_schema": map[string]interface{}{
					"type": "object",
					"properties": map[string]interface{}{
						"location": map[string]interface{}{
							"type": "string",
						},
					},
				},
			},
		},
		"max_tokens": 1024,
	}

	body, _ := json.Marshal(anthropicRequest)
	// Note: For now we route through a generic path
	// In Phase 2.9+, this would be integrated with the Anthropic dialect handler
	resp, err := http.Post(server.URL+"/v1/messages", "application/json", bytes.NewReader(body))
	if err != nil {
		t.Fatalf("request failed: %v", err)
	}
	var respData map[string]interface{}
	json.NewDecoder(resp.Body).Decode(&respData)
	resp.Body.Close()

	// Verify Anthropic tool_use format
	content := respData["content"].([]interface{})
	if len(content) < 2 {
		t.Errorf("expected at least 2 content blocks, got %d", len(content))
	}

	// Find tool_use block
	foundToolUse := false
	for _, block := range content {
		b := block.(map[string]interface{})
		if b["type"] == "tool_use" {
			foundToolUse = true
			if b["name"] != "get_weather" {
				t.Errorf("expected tool name 'get_weather', got %s", b["name"])
			}
		}
	}

	if !foundToolUse {
		t.Errorf("expected tool_use block in Anthropic response")
	}
}

// Test tool call with complex nested arguments
func TestComplexToolArguments(t *testing.T) {
	upstreamServer := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
		var req ChatCompletionRequest
		json.NewDecoder(r.Body).Decode(&req)

		w.Header().Set("Content-Type", "application/json")
		w.WriteHeader(http.StatusOK)

		// Response with complex nested tool arguments
		response := map[string]interface{}{
			"id": "chatcmpl-999",
			"choices": []map[string]interface{}{
				{
					"index": 0,
					"message": map[string]interface{}{
						"role":    "assistant",
						"content": nil,
						"tool_calls": []map[string]interface{}{
							{
								"id":   "call_complex",
								"type": "function",
								"function": map[string]interface{}{
									"name": "create_event",
									"arguments": `{
										"title": "Team Meeting",
										"time": "2025-08-20T14:00:00Z",
										"attendees": [
											{"name": "Alice", "email": "alice@example.com"},
											{"name": "Bob", "email": "bob@example.com"}
										],
										"location": {
											"address": "123 Main St",
											"city": "San Francisco",
											"country": "USA"
										},
										"tags": ["important", "recurring"]
									}`,
								},
							},
						},
					},
					"finish_reason": "tool_calls",
				},
			},
		}

		json.NewEncoder(w).Encode(response)
	}))
	defer upstreamServer.Close()

	upstreamAddr := strings.TrimPrefix(upstreamServer.URL, "http://")

	cfg := &config.Config{
		Models: map[string]*config.ModelUpstream{
			"reasoning": {
				Name:    "reasoning",
				Address: upstreamAddr,
			},
		},
		Routes: make(map[string]*config.Route),
	}

	handler := New(cfg)
	defer handler.Close()

	server := httptest.NewServer(handler)
	defer server.Close()

	request := ChatCompletionRequest{
		Model: "reasoning",
		Messages: []Message{
			{
				Role:    "user",
				Content: "Create a team meeting",
			},
		},
		Tools: []Tool{
			{
				Type: "function",
				Function: ToolFunction{
					Name:        "create_event",
					Description: "Create a calendar event",
					Parameters:  map[string]interface{}{},
				},
			},
		},
	}

	body, _ := json.Marshal(request)
	resp, err := http.Post(server.URL+"/v1/chat/completions", "application/json", bytes.NewReader(body))
	if err != nil {
		t.Fatalf("request failed: %v", err)
	}
	var respData map[string]interface{}
	json.NewDecoder(resp.Body).Decode(&respData)
	resp.Body.Close()

	// Verify complex nested arguments are preserved
	choices := respData["choices"].([]interface{})
	msg := choices[0].(map[string]interface{})["message"].(map[string]interface{})
	toolCalls := msg["tool_calls"].([]interface{})
	toolCall := toolCalls[0].(map[string]interface{})
	fn := toolCall["function"].(map[string]interface{})
	args := fn["arguments"].(string)

	// Verify the structure is intact
	if !strings.Contains(args, "alice@example.com") {
		t.Errorf("expected email in tool arguments")
	}
	if !strings.Contains(args, "San Francisco") {
		t.Errorf("expected location in tool arguments")
	}
	if !strings.Contains(args, "recurring") {
		t.Errorf("expected tags in tool arguments")
	}
}