package protocol

import (
	"crypto/rand"
	"encoding/json"
	"fmt"
	"strings"
	"time"
)

func newID() string { return fmt.Sprintf("resp_%x", randomID()) }
func randomID() []byte {
	b := make([]byte, 16)
	if _, e := rand.Read(b); e != nil {
		panic(e)
	}
	return b
}

func contentBlock(k Kind, p part, role string) (object, error) {
	switch p.typ {
	case "text":
		typ := "text"
		if k == Responses {
			typ = "input_text"
			if role == "assistant" {
				typ = "output_text"
			}
		}
		return object{"type": typ, "text": p.text}, nil
	case "image":
		if role != "user" {
			return nil, fmt.Errorf("images only supported in user messages")
		}
		switch k {
		case Chat:
			return object{"type": "image_url", "image_url": object{"url": p.url}}, nil
		case Responses:
			return object{"type": "input_image", "image_url": p.url}, nil
		default:
			src := object{"type": "url", "url": p.url}
			if strings.HasPrefix(p.url, "data:") {
				a, b, ok := strings.Cut(strings.TrimPrefix(p.url, "data:"), ";base64,")
				if !ok || a == "" || b == "" {
					return nil, fmt.Errorf("unsupported image data URL")
				}
				src = object{"type": "base64", "media_type": a, "data": b}
			}
			return object{"type": "image", "source": src}, nil
		}
	case "call":
		if role != "assistant" {
			return nil, fmt.Errorf("tool call outside assistant message")
		}
		switch k {
		case Chat:
			return object{"id": p.id, "type": "function", "function": object{"name": p.name, "arguments": p.args}}, nil
		case Responses:
			return object{"type": "function_call", "call_id": p.id, "name": p.name, "arguments": p.args}, nil
		default:
			args, e := decode([]byte(p.args))
			if e != nil {
				return nil, e
			}
			return object{"type": "tool_use", "id": p.id, "name": p.name, "input": args}, nil
		}
	case "result":
		if role != "tool" && role != "user" {
			return nil, fmt.Errorf("tool result outside user/tool message")
		}
		switch k {
		case Chat:
			return object{"role": "tool", "tool_call_id": p.id, "content": p.text}, nil
		case Responses:
			return object{"type": "function_call_output", "call_id": p.id, "output": p.text}, nil
		default:
			return object{"type": "tool_result", "tool_use_id": p.id, "content": p.text}, nil
		}
	}
	return nil, fmt.Errorf("unsupported part")
}
func writeRequest(k Kind, r request) (object, error) {
	out := object{}
	for key, val := range r.options {
		switch key {
		case "max":
			switch k {
			case Chat:
				key = "max_completion_tokens"
			case Responses:
				key = "max_output_tokens"
			case Anthropic:
				key = "max_tokens"
			}
		case "stop":
			if k == Responses {
				return nil, fmt.Errorf("Responses does not support stop sequences")
			}
			if k == Anthropic {
				key = "stop_sequences"
			}
		case "parallel_tool_calls":
			if k == Anthropic {
				continue
			}
		}
		out[key] = val
	}
	if k == Anthropic {
		if _, ok := out["max_tokens"]; !ok {
			return nil, fmt.Errorf("Anthropic requires explicit max_tokens")
		}
	}
	ms := []any{}
	sys := []any{}
	seenNonSystem := false
	for _, m := range r.messages {
		if k == Anthropic && (m.role == "system" || m.role == "developer") {
			if m.role == "developer" {
				return nil, fmt.Errorf("developer role cannot be represented by Anthropic")
			}
			if seenNonSystem {
				return nil, fmt.Errorf("mid-conversation system message cannot be represented by Anthropic")
			}
			for _, p := range m.parts {
				if p.typ != "text" {
					return nil, fmt.Errorf("non-text system content")
				}
				sys = append(sys, object{"type": "text", "text": p.text})
			}
			continue
		}
		seenNonSystem = true
		role := m.role
		if k == Anthropic && role == "tool" {
			role = "user"
		}
		cs := []any{}
		calls := []any{}
		flush := func() {
			if len(cs) > 0 {
				ms = append(ms, object{"role": role, "content": cs})
				cs = []any{}
			}
		}
		for _, p := range m.parts {
			b, e := contentBlock(k, p, m.role)
			if e != nil {
				return nil, e
			}
			switch {
			case k == Chat && p.typ == "call":
				calls = append(calls, b)
			case (k == Chat && p.typ == "result") || (k == Responses && (p.typ == "call" || p.typ == "result")):
				flush()
				ms = append(ms, b)
			default:
				cs = append(cs, b)
			}
		}
		if k == Chat && len(calls) > 0 {
			x := object{"role": role, "tool_calls": calls, "content": nil}
			if len(cs) > 0 {
				x["content"] = cs
			}
			ms = append(ms, x)
		} else {
			flush()
		}
	}
	if k == Anthropic {
		merged := []any{}
		for _, v := range ms {
			m := obj(v)
			if len(merged) > 0 && obj(merged[len(merged)-1])["role"] == m["role"] {
				last := obj(merged[len(merged)-1])
				last["content"] = append(arr(last["content"]), arr(m["content"])...)
			} else {
				merged = append(merged, m)
			}
		}
		ms = merged
		if len(sys) > 0 {
			out["system"] = sys
		}
	}
	if k == Responses {
		out["input"] = ms
		out["store"] = false
	} else {
		out["messages"] = ms
	}
	if len(r.tools) > 0 {
		ts := []any{}
		for _, f := range r.tools {
			switch k {
			case Chat:
				ts = append(ts, object{"type": "function", "function": f})
			case Responses:
				t := object{"type": "function"}
				for key, v := range f {
					t[key] = v
				}
				ts = append(ts, t)
			case Anthropic:
				t := object{}
				for key, v := range f {
					if key == "strict" {
						if v != false {
							return nil, fmt.Errorf("strict tool schema cannot be enforced by Anthropic")
						}
						continue
					}
					if key == "parameters" {
						key = "input_schema"
					}
					t[key] = v
				}
				ts = append(ts, t)
			}
		}
		out["tools"] = ts
	}
	if r.choice != nil {
		if s, ok := r.choice.(string); ok {
			if k == Anthropic {
				if s == "required" {
					s = "any"
				}
				out["tool_choice"] = object{"type": s}
			} else {
				out["tool_choice"] = s
			}
		} else {
			name := obj(r.choice)["name"]
			switch k {
			case Chat:
				out["tool_choice"] = object{"type": "function", "function": object{"name": name}}
			case Responses:
				out["tool_choice"] = object{"type": "function", "name": name}
			case Anthropic:
				out["tool_choice"] = object{"type": "tool", "name": name}
			}
		}
	}
	if k == Anthropic {
		if v, ok := r.options["parallel_tool_calls"]; ok {
			b, ok := v.(bool)
			if !ok {
				return nil, fmt.Errorf("invalid parallel_tool_calls")
			}
			tc := obj(out["tool_choice"])
			if tc == nil {
				tc = object{"type": "auto"}
				out["tool_choice"] = tc
			}
			tc["disable_parallel_tool_use"] = !b
		}
	}
	if k == Chat && out["stream"] == true {
		out["stream_options"] = object{"include_usage": true}
	}
	return out, nil
}

type response struct {
	id, model, finish string
	parts             []part
	usage             Usage
}

func finishReason(k Kind, s string) (string, error) {
	switch k {
	case Chat:
		switch s {
		case "stop":
			return "stop", nil
		case "tool_calls":
			return "tool", nil
		case "length":
			return "length", nil
		}
	case Anthropic:
		switch s {
		case "end_turn":
			return "stop", nil
		case "tool_use":
			return "tool", nil
		case "max_tokens":
			return "length", nil
		}
	case Responses:
		if s == "completed" {
			return "stop", nil
		}
	}
	return "", fmt.Errorf("unsupported finish reason %q", s)
}
func wireFinish(k Kind, s string) string {
	switch k {
	case Chat:
		switch s {
		case "tool":
			return "tool_calls"
		case "length":
			return "length"
		}
		return "stop"
	case Anthropic:
		switch s {
		case "tool":
			return "tool_use"
		case "length":
			return "max_tokens"
		}
		return "end_turn"
	default:
		if s == "length" {
			return "incomplete"
		}
		return "completed"
	}
}
func ConvertResponse(from, to Kind, body []byte) ([]byte, error) {
	if !valid(from) || !valid(to) {
		return nil, fmt.Errorf("unknown protocol")
	}
	m, e := decode(body)
	if e != nil {
		return nil, e
	}
	if from == to {
		return body, nil
	}
	r, e := readResponse(from, m)
	if e != nil {
		return nil, e
	}
	out, e := writeResponse(to, r)
	if e != nil {
		return nil, e
	}
	return json.Marshal(out)
}
func readResponse(k Kind, m object) (response, error) {
	r := response{id: str(m["id"]), model: str(m["model"]), usage: usage(k, obj(m["usage"]))}
	if m["error"] != nil {
		return r, fmt.Errorf("upstream error: %v", m["error"])
	}
	var e error
	switch k {
	case Chat:
		if e = fields(m, "id", "object", "created", "model", "choices", "usage", "system_fingerprint", "service_tier"); e != nil {
			return r, e
		}
		cs := arr(m["choices"])
		if len(cs) != 1 {
			return r, fmt.Errorf("expected one choice")
		}
		c := obj(cs[0])
		if e = fields(c, "index", "message", "finish_reason", "logprobs"); e != nil {
			return r, e
		}
		if c["logprobs"] != nil {
			return r, fmt.Errorf("logprobs unsupported")
		}
		message, e := chatResponseMessage(obj(c["message"]))
		if e != nil {
			return r, e
		}
		mm, e := readMessage(k, message)
		if e != nil {
			return r, e
		}
		r.parts = mm.parts
		r.finish, e = finishReason(k, str(c["finish_reason"]))
		if e != nil {
			return r, e
		}
	case Anthropic:
		if e = fields(m, "id", "type", "role", "model", "content", "stop_reason", "stop_sequence", "usage"); e != nil {
			return r, e
		}
		if m["stop_sequence"] != nil {
			return r, fmt.Errorf("stop sequence cannot be preserved")
		}
		content, e := anthropicOutput(m["content"])
		if e != nil {
			return r, e
		}
		r.parts, e = readContent(k, content)
		if e != nil {
			return r, e
		}
		r.finish, e = finishReason(k, str(m["stop_reason"]))
		if e != nil {
			return r, e
		}
	case Responses:
		// Responses echoes request configuration; only output items carry generated content.
		for _, v := range arr(m["output"]) {
			x := obj(v)
			switch str(x["type"]) {
			case "reasoning":
				if e = validateReasoning(x); e != nil {
					return r, e
				}
			case "message":
				x = responseMessage(x)
				mm, e := readMessage(k, x)
				if e != nil {
					return r, e
				}
				r.parts = append(r.parts, mm.parts...)
			case "function_call":
				if e = fields(x, "type", "id", "call_id", "name", "arguments", "status"); e != nil {
					return r, e
				}
				p, e := callPart(x, "call_id", "arguments")
				if e != nil {
					return r, e
				}
				r.parts = append(r.parts, p)
			default:
				return r, fmt.Errorf("unsupported output item %q", str(x["type"]))
			}
		}
		if str(m["status"]) == "incomplete" {
			if str(obj(m["incomplete_details"])["reason"]) != "max_output_tokens" {
				return r, fmt.Errorf("unsupported incomplete reason")
			}
			r.finish = "length"
		} else {
			r.finish, e = finishReason(k, str(m["status"]))
			if e != nil {
				return r, e
			}
			for _, p := range r.parts {
				if p.typ == "call" {
					r.finish = "tool"
				}
			}
		}
	}
	return r, nil
}
func writeResponse(k Kind, r response) (object, error) {
	if r.id == "" {
		r.id = newID()
	}
	out := object{"id": r.id, "model": r.model}
	cs := []any{}
	calls := []any{}
	text := ""
	for i, p := range r.parts {
		if p.typ != "text" && p.typ != "call" {
			return nil, fmt.Errorf("unsupported response content")
		}
		b, e := contentBlock(k, p, "assistant")
		if e != nil {
			return nil, e
		}
		switch k {
		case Chat:
			if p.typ == "call" {
				calls = append(calls, b)
			} else {
				text += p.text
			}
		case Anthropic:
			cs = append(cs, b)
		case Responses:
			if p.typ == "call" {
				b["id"] = fmt.Sprintf("fc_%s_%d", r.id, i)
				b["status"] = "completed"
				cs = append(cs, b)
			} else {
				b["annotations"] = []any{}
				cs = append(cs, object{"type": "message", "id": fmt.Sprintf("msg_%s_%d", r.id, i), "role": "assistant", "status": "completed", "content": []any{b}})
			}
		}
	}
	switch k {
	case Chat:
		m := object{"role": "assistant", "content": text}
		if len(calls) > 0 {
			m["tool_calls"] = calls
			if text == "" {
				m["content"] = nil
			}
		}
		out["object"] = "chat.completion"
		out["created"] = 0
		out["choices"] = []any{object{"index": 0, "message": m, "finish_reason": wireFinish(k, r.finish)}}
	case Anthropic:
		out["type"] = "message"
		out["role"] = "assistant"
		out["content"] = cs
		out["stop_reason"] = wireFinish(k, r.finish)
		out["stop_sequence"] = nil
	case Responses:
		out["object"] = "response"
		out["created_at"] = time.Now().Unix()
		out["status"] = wireFinish(k, r.finish)
		out["output"] = cs
		out["error"] = nil
		out["incomplete_details"] = nil
		if r.finish == "length" {
			out["incomplete_details"] = object{"reason": "max_output_tokens"}
		}
	}
	if r.usage.Known || k == Anthropic {
		out["usage"] = usageJSON(k, r.usage)
	}
	return out, nil
}