package codex

import (
	"bufio"
	"bytes"
	"crypto/rand"
	"encoding/hex"
	"encoding/json"
	"fmt"
	"io"
	"net/http"
	"strings"
	"time"
)

const toolNameLimit = 64

// --- Tool name shortening ---

// shortenName shortens a tool name to fit within limit characters.
// For MCP names (mcp__server__tool), strips the server part.
func shortenName(name string, limit int) string {
	if len(name) <= limit {
		return name
	}
	if strings.HasPrefix(name, "mcp__") {
		if idx := strings.LastIndex(name, "__"); idx >= 0 {
			candidate := "mcp__" + name[idx+2:]
			if len(candidate) <= limit {
				return candidate
			}
			return candidate[:limit]
		}
	}
	if len(name) > limit {
		return name[:limit]
	}
	return name
}

// buildShortNameMap builds a map from original names to shortened names,
// resolving collisions by appending _1, _2, etc.
func buildShortNameMap(names []string) map[string]string {
	m := make(map[string]string, len(names))
	used := make(map[string]bool, len(names))
	for _, name := range names {
		candidate := shortenName(name, toolNameLimit)
		if !used[candidate] {
			m[name] = candidate
			used[candidate] = true
			continue
		}
		for i := 1; i < 10000; i++ {
			suffix := fmt.Sprintf("_%d", i)
			maxBase := toolNameLimit - len(suffix)
			base := candidate
			if len(base) > maxBase {
				base = base[:maxBase]
			}
			attempt := base + suffix
			if !used[attempt] {
				m[name] = attempt
				used[attempt] = true
				break
			}
		}
	}
	return m
}

// buildReverseNameMap inverts a short→original mapping.
func buildReverseNameMap(shortMap map[string]string) map[string]string {
	rev := make(map[string]string, len(shortMap))
	for orig, short := range shortMap {
		rev[short] = orig
	}
	return rev
}

// --- Request translation ---

// translateRequest converts an OpenAI Chat Completions JSON body to a Codex Responses API JSON body.
// Returns the translated JSON and the tool name map (original→shortened).
func translateRequest(jsonData []byte) ([]byte, map[string]string, error) {
	var req map[string]interface{}
	if err := json.Unmarshal(jsonData, &req); err != nil {
		return nil, nil, fmt.Errorf("failed to parse chat completion request: %w", err)
	}

	// Collect tool names for shortening
	var toolNames []string
	if tools, ok := req["tools"].([]interface{}); ok {
		for _, t := range tools {
			if tm, ok := t.(map[string]interface{}); ok {
				if tm["type"] == "function" {
					if fn, ok := tm["function"].(map[string]interface{}); ok {
						if name, ok := fn["name"].(string); ok {
							toolNames = append(toolNames, name)
						}
					}
				}
			}
		}
	}
	nameMap := buildShortNameMap(toolNames)

	// Build Responses API request
	out := map[string]interface{}{
		"instructions":        "",
		"stream":              true,
		"model":               "",
		"reasoning":           map[string]interface{}{"effort": "medium", "summary": "auto"},
		"store":               false,
		"parallel_tool_calls": true,
		"include":             []string{"reasoning.encrypted_content"},
	}

	if model, ok := req["model"].(string); ok {
		out["model"] = model
	}

	// reasoning_effort → reasoning
	if re, ok := req["reasoning_effort"].(string); ok && re != "" {
		out["reasoning"] = map[string]interface{}{
			"effort":  re,
			"summary": "auto",
		}
	}

	// messages → input
	input, err := translateMessages(req, nameMap)
	if err != nil {
		return nil, nil, err
	}
	out["input"] = input

	// tools
	if tools, ok := req["tools"].([]interface{}); ok && len(tools) > 0 {
		translated, err := translateTools(tools, nameMap)
		if err != nil {
			return nil, nil, err
		}
		out["tools"] = translated
	}

	// tool_choice
	if tc, ok := req["tool_choice"]; ok {
		out["tool_choice"] = translateToolChoice(tc, nameMap)
	}

	// response_format → text
	if rf, ok := req["response_format"].(map[string]interface{}); ok {
		text := translateResponseFormat(rf)
		if text != nil {
			out["text"] = text
		}
	}

	result, err := json.Marshal(out)
	if err != nil {
		return nil, nil, fmt.Errorf("failed to marshal codex request: %w", err)
	}
	return result, nameMap, nil
}

func translateMessages(req map[string]interface{}, nameMap map[string]string) ([]interface{}, error) {
	messages, ok := req["messages"].([]interface{})
	if !ok || len(messages) == 0 {
		return []interface{}{}, nil
	}

	var input []interface{}
	for _, msg := range messages {
		m, ok := msg.(map[string]interface{})
		if !ok {
			continue
		}
		role, _ := m["role"].(string)

		switch role {
		case "tool":
			// Tool result → function_call_output
			callID, _ := m["tool_call_id"].(string)
			input = append(input, map[string]interface{}{
				"type":    "function_call_output",
				"call_id": callID,
				"output":  contentToOutput(m["content"]),
			})

		default:
			// system → developer, others unchanged
			mappedRole := role
			if role == "system" {
				mappedRole = "developer"
			}

			// Build content parts
			contentParts := translateContentParts(m["content"], role)
			input = append(input, map[string]interface{}{
				"type":    "message",
				"role":    mappedRole,
				"content": contentParts,
			})

			// Assistant tool_calls → separate function_call items
			if role == "assistant" {
				if toolCalls, ok := m["tool_calls"].([]interface{}); ok {
					for _, tc := range toolCalls {
						tcMap, ok := tc.(map[string]interface{})
						if !ok {
							continue
						}
						if tcMap["type"] == "function" {
							fn, _ := tcMap["function"].(map[string]interface{})
							if fn == nil {
								continue
							}
							origName, _ := fn["name"].(string)
							shortName := shortenName(origName, toolNameLimit)
							if sn, ok := nameMap[origName]; ok {
								shortName = sn
							}
							args, _ := fn["arguments"].(string)
							callID, _ := tcMap["id"].(string)
							input = append(input, map[string]interface{}{
								"type":      "function_call",
								"call_id":   callID,
								"name":      shortName,
								"arguments": args,
							})
						}
					}
				}
			}
		}
	}
	return input, nil
}

func translateContentParts(content interface{}, role string) []interface{} {
	if content == nil {
		return []interface{}{}
	}

	// String content
	if s, ok := content.(string); ok {
		if s == "" {
			return []interface{}{}
		}
		textType := "input_text"
		if role == "assistant" {
			textType = "output_text"
		}
		return []interface{}{
			map[string]interface{}{"type": textType, "text": s},
		}
	}

	// Array content
	parts, ok := content.([]interface{})
	if !ok {
		return []interface{}{}
	}

	var result []interface{}
	for _, part := range parts {
		pm, ok := part.(map[string]interface{})
		if !ok {
			continue
		}
		t, _ := pm["type"].(string)

		switch t {
		case "text":
			textType := "input_text"
			if role == "assistant" {
				textType = "output_text"
			}
			text, _ := pm["text"].(string)
			result = append(result, map[string]interface{}{
				"type": textType,
				"text": text,
			})
		case "image_url":
			if role == "user" {
				if iu, ok := pm["image_url"].(map[string]interface{}); ok {
					if url, ok := iu["url"].(string); ok && url != "" {
						result = append(result, map[string]interface{}{
							"type":      "input_image",
							"image_url": url,
						})
					}
				}
			}
		}
	}
	return result
}

func translateTools(tools []interface{}, nameMap map[string]string) ([]interface{}, error) {
	var result []interface{}
	for _, tool := range tools {
		t, ok := tool.(map[string]interface{})
		if !ok {
			continue
		}
		if t["type"] != "function" {
			// Non-function tools pass through
			result = append(result, t)
			continue
		}
		fn, _ := t["function"].(map[string]interface{})
		if fn == nil {
			continue
		}
		origName, _ := fn["name"].(string)
		shortName := shortenName(origName, toolNameLimit)
		if sn, ok := nameMap[origName]; ok {
			shortName = sn
		}
		ft := map[string]interface{}{
			"type": "function",
			"name": shortName,
		}
		if desc, ok := fn["description"].(string); ok && desc != "" {
			ft["description"] = desc
		}
		if params, ok := fn["parameters"]; ok {
			ft["parameters"] = params
		}
		if strict, ok := fn["strict"]; ok {
			ft["strict"] = strict
		}
		result = append(result, ft)
	}
	return result, nil
}

func translateToolChoice(tc interface{}, nameMap map[string]string) interface{} {
	// String: "auto", "required", "none"
	if s, ok := tc.(string); ok {
		return s
	}
	// Object
	m, ok := tc.(map[string]interface{})
	if !ok {
		return tc
	}
	if m["type"] == "function" {
		choice := map[string]interface{}{"type": "function"}
		if fn, ok := m["function"].(map[string]interface{}); ok {
			origName, _ := fn["name"].(string)
			shortName := shortenName(origName, toolNameLimit)
			if sn, ok := nameMap[origName]; ok {
				shortName = sn
			}
			if shortName != "" {
				choice["name"] = shortName
			}
		}
		return choice
	}
	return tc
}

func translateResponseFormat(rf map[string]interface{}) interface{} {
	t, _ := rf["type"].(string)
	switch t {
	case "text":
		return map[string]interface{}{
			"format": map[string]interface{}{"type": "text"},
		}
	case "json_schema":
		js, _ := rf["json_schema"].(map[string]interface{})
		if js == nil {
			return map[string]interface{}{
				"format": map[string]interface{}{"type": "json_schema"},
			}
		}
		format := map[string]interface{}{"type": "json_schema"}
		if name, ok := js["name"]; ok {
			format["name"] = name
		}
		if strict, ok := js["strict"]; ok {
			format["strict"] = strict
		}
		if schema, ok := js["schema"]; ok {
			format["schema"] = schema
		}
		return map[string]interface{}{"format": format}
	}
	return nil
}

func contentToOutput(v interface{}) interface{} {
	if v == nil {
		return ""
	}
	if s, ok := v.(string); ok {
		return s
	}
	return v
}

// extractStreamField extracts the "stream" field from a Chat Completions JSON body.
func extractStreamField(jsonData []byte) bool {
	var partial struct {
		Stream *bool `json:"stream"`
	}
	json.Unmarshal(jsonData, &partial)
	return partial.Stream != nil && *partial.Stream
}

// --- Streaming SSE translation ---

// streamState tracks state across SSE events.
type streamState struct {
	responseID      string
	createdAt       int64
	model           string
	functionCallIdx int // -1 = none seen
	reverseNameMap  map[string]string
}

// SSETranslator wraps an upstream SSE stream and translates Responses API
// events into Chat Completions SSE chunks.
type SSETranslator struct {
	scanner  *bufio.Scanner
	body     io.Closer
	state    *streamState
	buf      bytes.Buffer
	done     bool
	doneSent bool
}

func newSSETranslator(body io.ReadCloser, nameMap map[string]string) *SSETranslator {
	rev := buildReverseNameMap(nameMap)
	scanner := bufio.NewScanner(body)
	scanner.Buffer(make([]byte, 0, 64*1024), 1*1024*1024)
	return &SSETranslator{
		scanner: scanner,
		body:    body,
		state: &streamState{
			functionCallIdx: -1,
			reverseNameMap:  rev,
		},
	}
}

func (t *SSETranslator) Close() error {
	if t.body != nil {
		return t.body.Close()
	}
	return nil
}

func (t *SSETranslator) Read(p []byte) (int, error) {
	if t.buf.Len() > 0 {
		return t.buf.Read(p)
	}
	if t.done {
		return 0, io.EOF
	}

	for t.scanner.Scan() {
		line := t.scanner.Text()

		// Skip empty lines
		if line == "" {
			continue
		}

		// Only process data: lines. Accept both "data: " and "data:".
		if !strings.HasPrefix(line, "data:") {
			continue
		}
		data := strings.TrimSpace(line[5:])

		// [DONE] sentinel
		if data == "[DONE]" {
			t.done = true
			t.writeDone()
			return t.buf.Read(p)
		}

		var event map[string]interface{}
		if err := json.Unmarshal([]byte(data), &event); err != nil {
			continue
		}

		eventType, _ := event["type"].(string)
		output := t.handleEvent(eventType, event)
		if output == nil {
			continue
		}

		chunk, err := json.Marshal(output)
		if err != nil {
			continue
		}
		t.buf.WriteString("data: ")
		t.buf.Write(chunk)
		t.buf.WriteString("\n\n")
		return t.buf.Read(p)
	}

	t.done = true
	if err := t.scanner.Err(); err != nil {
		return 0, err
	}
	t.writeDone()
	if t.buf.Len() > 0 {
		return t.buf.Read(p)
	}
	return 0, io.EOF
}

func (t *SSETranslator) writeDone() {
	if t.doneSent {
		return
	}
	t.buf.WriteString("data: [DONE]\n\n")
	t.doneSent = true
}

func (t *SSETranslator) handleEvent(eventType string, event map[string]interface{}) interface{} {
	switch eventType {
	case "response.created":
		return t.handleResponseCreated(event)
	case "response.reasoning_summary_text.delta":
		return t.handleReasoningDelta(event)
	case "response.reasoning_summary_text.done":
		return t.handleReasoningDone()
	case "response.output_text.delta":
		return t.handleOutputTextDelta(event)
	case "response.output_item.done":
		return t.handleOutputItemDone(event)
	case "response.completed":
		return t.handleCompleted(event)
	}
	return nil
}

func (t *SSETranslator) handleResponseCreated(event map[string]interface{}) interface{} {
	resp, _ := event["response"].(map[string]interface{})
	if resp != nil {
		if id, ok := resp["id"].(string); ok {
			t.state.responseID = id
		}
		if created, ok := resp["created_at"].(float64); ok {
			t.state.createdAt = int64(created)
		}
		if model, ok := resp["model"].(string); ok {
			t.state.model = model
		}
	}
	return nil // no chunk emitted
}

func (t *SSETranslator) handleReasoningDelta(event map[string]interface{}) interface{} {
	delta, _ := event["delta"].(string)
	return t.makeChunk(map[string]interface{}{
		"role":              "assistant",
		"reasoning_content": delta,
	}, nil)
}

func (t *SSETranslator) handleReasoningDone() interface{} {
	return t.makeChunk(map[string]interface{}{
		"role":              "assistant",
		"reasoning_content": "\n\n",
	}, nil)
}

func (t *SSETranslator) handleOutputTextDelta(event map[string]interface{}) interface{} {
	delta, _ := event["delta"].(string)
	return t.makeChunk(map[string]interface{}{
		"role":    "assistant",
		"content": delta,
	}, nil)
}

func (t *SSETranslator) handleOutputItemDone(event map[string]interface{}) interface{} {
	item, _ := event["item"].(map[string]interface{})
	if item == nil || item["type"] != "function_call" {
		return nil
	}
	t.state.functionCallIdx++
	name, _ := item["name"].(string)
	if rev, ok := t.state.reverseNameMap[name]; ok {
		name = rev
	}
	callID, _ := item["call_id"].(string)
	args, _ := item["arguments"].(string)

	return t.makeChunk(map[string]interface{}{
		"role": "assistant",
		"tool_calls": []interface{}{
			map[string]interface{}{
				"index": t.state.functionCallIdx,
				"id":    callID,
				"type":  "function",
				"function": map[string]interface{}{
					"name":      name,
					"arguments": args,
				},
			},
		},
	}, nil)
}

func (t *SSETranslator) handleCompleted(event map[string]interface{}) interface{} {
	resp, _ := event["response"].(map[string]interface{})
	finishReason := "stop"
	if t.state.functionCallIdx >= 0 {
		finishReason = "tool_calls"
	}

	chunk := t.makeChunk(map[string]interface{}{}, &finishReason)

	// Extract usage
	if resp != nil {
		if usage, ok := resp["usage"].(map[string]interface{}); ok {
			chunk["usage"] = translateUsage(usage)
		}
	}
	return chunk
}

func (t *SSETranslator) makeChunk(delta map[string]interface{}, finishReason *string) map[string]interface{} {
	chunk := map[string]interface{}{
		"id":      t.state.responseID,
		"object":  "chat.completion.chunk",
		"created": t.state.createdAt,
		"model":   t.state.model,
		"choices": []interface{}{
			map[string]interface{}{
				"index":         0,
				"delta":         delta,
				"finish_reason": finishReason,
			},
		},
	}
	return chunk
}

// --- Non-streaming response translation ---

// translateNonStreamingResponse reads the full SSE stream, extracts response.completed,
// and returns an *http.Response with Chat Completions JSON body.
func translateNonStreamingResponse(resp *http.Response, nameMap map[string]string) (*http.Response, error) {
	defer resp.Body.Close()

	body, err := io.ReadAll(resp.Body)
	if err != nil {
		return nil, fmt.Errorf("failed to read codex response: %w", err)
	}

	// Find response.completed event in SSE stream
	var completed map[string]interface{}
	for _, line := range strings.Split(string(body), "\n") {
		line = strings.TrimSpace(line)
		if !strings.HasPrefix(line, "data:") {
			continue
		}
		data := strings.TrimSpace(line[5:])
		if data == "[DONE]" {
			continue
		}
		var event map[string]interface{}
		if err := json.Unmarshal([]byte(data), &event); err != nil {
			continue
		}
		if event["type"] == "response.completed" {
			completed = event
			break
		}
	}

	if completed == nil {
		return nil, fmt.Errorf("no response.completed event found in codex response")
	}

	chatResp := buildChatCompletionResponse(completed, nameMap)
	respBody, err := json.Marshal(chatResp)
	if err != nil {
		return nil, fmt.Errorf("failed to marshal chat completion response: %w", err)
	}

	return &http.Response{
		StatusCode:    http.StatusOK,
		Header:        http.Header{"Content-Type": []string{"application/json"}},
		Body:          io.NopCloser(bytes.NewReader(respBody)),
		ContentLength: int64(len(respBody)),
	}, nil
}

func buildChatCompletionResponse(completed map[string]interface{}, nameMap map[string]string) map[string]interface{} {
	rev := buildReverseNameMap(nameMap)
	resp, _ := completed["response"].(map[string]interface{})

	var responseID string
	var createdAt int64
	var model string
	if resp != nil {
		responseID, _ = resp["id"].(string)
		if ca, ok := resp["created_at"].(float64); ok {
			createdAt = int64(ca)
		}
		model, _ = resp["model"].(string)
	}
	if responseID == "" {
		responseID = fmt.Sprintf("chatcmpl-%s", randomHex(24))
	}
	if createdAt == 0 {
		createdAt = time.Now().Unix()
	}

	msg := map[string]interface{}{
		"role":    "assistant",
		"content": nil,
	}
	var toolCalls []interface{}
	finishReason := "stop"
	var reasoningContent string

	// Process output items
	var outputs []interface{}
	if resp != nil {
		outputs, _ = resp["output"].([]interface{})
	}
	for _, item := range outputs {
		im, ok := item.(map[string]interface{})
		if !ok {
			continue
		}
		switch im["type"] {
		case "message":
			// Extract last output_text content
			content, _ := im["content"].([]interface{})
			for _, c := range content {
				cm, ok := c.(map[string]interface{})
				if !ok {
					continue
				}
				if cm["type"] == "output_text" {
					if text, ok := cm["text"].(string); ok {
						msg["content"] = text
					}
				}
			}
		case "reasoning":
			// Extract summary text
			summary, _ := im["summary"].([]interface{})
			for _, s := range summary {
				sm, ok := s.(map[string]interface{})
				if !ok {
					continue
				}
				if sm["type"] == "summary_text" {
					if text, ok := sm["text"].(string); ok {
						reasoningContent = text
					}
				}
			}
		case "function_call":
			name, _ := im["name"].(string)
			if orig, ok := rev[name]; ok {
				name = orig
			}
			callID, _ := im["call_id"].(string)
			args, _ := im["arguments"].(string)
			toolCalls = append(toolCalls, map[string]interface{}{
				"id":   callID,
				"type": "function",
				"function": map[string]interface{}{
					"name":      name,
					"arguments": args,
				},
			})
			finishReason = "tool_calls"
		}
	}

	if reasoningContent != "" {
		msg["reasoning_content"] = reasoningContent
	}
	if len(toolCalls) > 0 {
		msg["tool_calls"] = toolCalls
	}

	result := map[string]interface{}{
		"id":      responseID,
		"object":  "chat.completion",
		"created": createdAt,
		"model":   model,
		"choices": []interface{}{
			map[string]interface{}{
				"index":         0,
				"message":       msg,
				"finish_reason": finishReason,
			},
		},
	}

	// Usage
	if resp != nil {
		if usage, ok := resp["usage"].(map[string]interface{}); ok {
			result["usage"] = translateUsage(usage)
		}
	}

	return result
}

func translateUsage(usage map[string]interface{}) map[string]interface{} {
	u := map[string]interface{}{
		"prompt_tokens":     intFromFloat(usage["input_tokens"]),
		"completion_tokens": intFromFloat(usage["output_tokens"]),
		"total_tokens":      intFromFloat(usage["total_tokens"]),
	}

	if details, ok := usage["input_tokens_details"].(map[string]interface{}); ok {
		if cached, ok := details["cached_tokens"]; ok && cached != nil {
			u["prompt_tokens_details"] = map[string]interface{}{
				"cached_tokens": intFromFloat(cached),
			}
		}
	}
	if details, ok := usage["output_tokens_details"].(map[string]interface{}); ok {
		if reasoning, ok := details["reasoning_tokens"]; ok && reasoning != nil {
			u["completion_tokens_details"] = map[string]interface{}{
				"reasoning_tokens": intFromFloat(reasoning),
			}
		}
	}

	return u
}

func intFromFloat(v interface{}) int {
	if v == nil {
		return 0
	}
	switch n := v.(type) {
	case float64:
		return int(n)
	case int:
		return n
	case json.Number:
		i, _ := n.Int64()
		return int(i)
	}
	return 0
}

func randomHex(n int) string {
	b := make([]byte, n/2+1)
	rand.Read(b)
	return hex.EncodeToString(b)[:n]
}