package protocol

import (
	"bytes"
	"encoding/json"
	"fmt"
	"io"
)

type Kind string

const (
	Chat      Kind = "openai-chat"
	Responses Kind = "openai-responses"
	Anthropic Kind = "anthropic"
)

type Usage struct {
	Input      int64           `json:"input_tokens"`
	Output     int64           `json:"output_tokens"`
	CacheRead  int64           `json:"cache_read_tokens"`
	CacheWrite int64           `json:"cache_write_tokens"`
	Known      bool            `json:"known"`
	Raw        json.RawMessage `json:"raw,omitempty"`
}
type object = map[string]any

// maxJSONNodes bounds how many JSON tokens one body may contain before it is
// materialized into maps. Without it a 16 MiB body of micro-values expands
// into hundreds of MB of Go objects per request.
const maxJSONNodes = 1 << 18

func valid(k Kind) bool { return k == Chat || k == Responses || k == Anthropic }

// scanNodeCount walks b with the streaming decoder and fails once more than
// limit tokens have been seen. Malformed JSON is not an error here; the
// materializing pass reports its usual message.
func scanNodeCount(b []byte, limit int) error {
	d := json.NewDecoder(bytes.NewReader(b))
	n := 0
	for {
		_, err := d.Token()
		if err != nil {
			return nil
		}
		if n++; n > limit {
			return fmt.Errorf("JSON exceeds %d nodes", limit)
		}
	}
}
func decode(b []byte) (object, error) {
	if err := scanNodeCount(b, maxJSONNodes); err != nil {
		return nil, err
	}
	d := json.NewDecoder(bytes.NewReader(b))
	d.UseNumber()
	var v object
	if err := d.Decode(&v); err != nil {
		return nil, err
	}
	if v == nil {
		return nil, fmt.Errorf("expected JSON object")
	}
	var extra any
	if err := d.Decode(&extra); err != io.EOF {
		return nil, fmt.Errorf("trailing JSON data")
	}
	return v, nil
}
func obj(v any) object { m, _ := v.(map[string]any); return m }
func str(v any) string { s, _ := v.(string); return s }
func arr(v any) []any  { a, _ := v.([]any); return a }
func num(v any) int64 {
	switch n := v.(type) {
	case json.Number:
		i, _ := n.Int64()
		return i
	case int64:
		return n
	case int:
		return int64(n)
	case float64:
		return int64(n)
	}
	return 0
}
func fields(m object, names ...string) error {
	allowed := map[string]bool{}
	for _, n := range names {
		allowed[n] = true
	}
	for k := range m {
		if !allowed[k] {
			return fmt.Errorf("unsupported field %q", k)
		}
	}
	return nil
}
func required(m object, k string) (string, error) {
	s, ok := m[k].(string)
	if !ok || s == "" {
		return "", fmt.Errorf("missing or invalid %s", k)
	}
	return s, nil
}
func ExtractUsage(kind Kind, body []byte) Usage {
	m, e := decode(body)
	if e != nil {
		return Usage{}
	}
	if r := obj(m["response"]); r != nil {
		m = r
	}
	if r := obj(m["message"]); r != nil {
		m = r
	}
	return usage(kind, obj(m["usage"]))
}
func usage(k Kind, m object) Usage {
	if m == nil {
		return Usage{}
	}
	raw, _ := json.Marshal(m)
	u := Usage{Raw: raw}
	switch k {
	case Chat:
		_, a := m["prompt_tokens"]
		_, b := m["completion_tokens"]
		u.Known = a || b
		u.Input = num(m["prompt_tokens"])
		u.Output = num(m["completion_tokens"])
		u.CacheRead = num(obj(m["prompt_tokens_details"])["cached_tokens"])
	case Responses:
		_, a := m["input_tokens"]
		_, b := m["output_tokens"]
		u.Known = a || b
		u.Input = num(m["input_tokens"])
		u.Output = num(m["output_tokens"])
		u.CacheRead = num(obj(m["input_tokens_details"])["cached_tokens"])
	case Anthropic:
		_, a := m["input_tokens"]
		_, b := m["output_tokens"]
		u.Known = a || b
		u.CacheRead = num(m["cache_read_input_tokens"])
		u.CacheWrite = num(m["cache_creation_input_tokens"])
		u.Input = num(m["input_tokens"]) + u.CacheRead + u.CacheWrite
		u.Output = num(m["output_tokens"])
	}
	return u
}
func usageJSON(k Kind, u Usage) object {
	switch k {
	case Chat:
		return object{"prompt_tokens": u.Input, "completion_tokens": u.Output, "total_tokens": u.Input + u.Output, "prompt_tokens_details": object{"cached_tokens": u.CacheRead}}
	case Responses:
		return object{"input_tokens": u.Input, "output_tokens": u.Output, "total_tokens": u.Input + u.Output, "input_tokens_details": object{"cached_tokens": u.CacheRead}}
	default:
		return object{"input_tokens": u.Input - u.CacheRead - u.CacheWrite, "output_tokens": u.Output, "cache_read_input_tokens": u.CacheRead, "cache_creation_input_tokens": u.CacheWrite}
	}
}

type part struct{ typ, text, id, name, args, url string }
type message struct {
	role  string
	parts []part
}
type request struct {
	messages []message
	tools    []object
	options  object
	choice   any
}

func textPart(s string) part { return part{typ: "text", text: s} }

func ConvertRequest(from, to Kind, body []byte, model string) ([]byte, error) {
	if !valid(from) || !valid(to) {
		return nil, fmt.Errorf("unknown protocol")
	}
	m, err := decode(body)
	if err != nil {
		return nil, err
	}
	if from == to {
		if model != "" {
			m["model"] = model
		}
		return json.Marshal(m)
	}
	r, err := readRequest(from, m)
	if err != nil {
		return nil, err
	}
	out, err := writeRequest(to, r)
	if err != nil {
		return nil, err
	}
	if model == "" {
		model = str(m["model"])
	}
	out["model"] = model
	return json.Marshal(out)
}
func readRequest(k Kind, m object) (request, error) {
	r := request{options: object{}}
	var err error
	common := []string{"model", "stream", "temperature", "top_p", "tools", "tool_choice"}
	switch k {
	case Chat:
		err = fields(m, append(common, "messages", "max_tokens", "max_completion_tokens", "stop", "stream_options", "parallel_tool_calls", "n")...)
	case Responses:
		err = fields(m, append(common, "input", "instructions", "max_output_tokens", "parallel_tool_calls", "store")...)
	case Anthropic:
		err = fields(m, append(common, "messages", "system", "max_tokens", "stop_sequences")...)
	}
	if err != nil {
		return r, err
	}
	if n, ok := m["n"]; ok && num(n) != 1 {
		return r, fmt.Errorf("multiple choices unsupported")
	}
	if v, ok := m["store"]; ok && v != false {
		return r, fmt.Errorf("stored Responses unsupported")
	}
	if v, ok := m["stream_options"]; ok {
		o := obj(v)
		if o == nil {
			return r, fmt.Errorf("invalid stream_options")
		}
		if e := fields(o, "include_usage"); e != nil {
			return r, e
		}
	}
	for _, key := range []string{"stream", "temperature", "top_p", "parallel_tool_calls"} {
		if v, ok := m[key]; ok {
			r.options[key] = v
		}
	}
	for _, key := range []string{"max_tokens", "max_completion_tokens", "max_output_tokens"} {
		if v, ok := m[key]; ok {
			if _, exists := r.options["max"]; exists {
				return r, fmt.Errorf("conflicting token limits")
			}
			r.options["max"] = v
		}
	}
	for _, key := range []string{"stop", "stop_sequences"} {
		if v, ok := m[key]; ok {
			if s, ok := v.(string); ok {
				v = []any{s}
			}
			if _, ok := v.([]any); !ok {
				return r, fmt.Errorf("invalid stop sequences")
			}
			r.options["stop"] = v
		}
	}
	if k == Responses {
		if s, ok := m["instructions"]; ok {
			v, ok := s.(string)
			if !ok {
				return r, fmt.Errorf("invalid instructions")
			}
			r.messages = append(r.messages, message{"system", []part{textPart(v)}})
		}
		if s, ok := m["input"].(string); ok {
			r.messages = append(r.messages, message{"user", []part{textPart(s)}})
		} else {
			items, ok := m["input"].([]any)
			if !ok {
				return r, fmt.Errorf("invalid input")
			}
			for _, v := range items {
				x := obj(v)
				typ := str(x["type"])
				switch typ {
				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.messages = append(r.messages, message{"assistant", []part{p}})
				case "function_call_output":
					if e := fields(x, "type", "id", "call_id", "output", "status"); e != nil {
						return r, e
					}
					id, e := required(x, "call_id")
					if e != nil {
						return r, e
					}
					s, ok := x["output"].(string)
					if !ok {
						return r, fmt.Errorf("non-text function output unsupported")
					}
					r.messages = append(r.messages, message{"tool", []part{{typ: "result", id: id, text: s}}})
				case "", "message":
					mm, e := readMessage(k, x)
					if e != nil {
						return r, e
					}
					r.messages = append(r.messages, mm)
				default:
					return r, fmt.Errorf("unsupported input item %q", typ)
				}
			}
		}
	} else {
		if k == Anthropic {
			if s, ok := m["system"]; ok {
				ps, e := readContent(k, s)
				if e != nil {
					return r, e
				}
				for _, p := range ps {
					if p.typ != "text" {
						return r, fmt.Errorf("non-text system")
					}
				}
				r.messages = append(r.messages, message{"system", ps})
			}
		}
		ms, ok := m["messages"].([]any)
		if !ok {
			return r, fmt.Errorf("invalid messages")
		}
		for _, v := range ms {
			mm, e := readMessage(k, obj(v))
			if e != nil {
				return r, e
			}
			r.messages = append(r.messages, mm)
		}
	}
	if v, ok := m["tools"]; ok {
		ts, ok := v.([]any)
		if !ok {
			return r, fmt.Errorf("invalid tools")
		}
		for _, v := range ts {
			t := obj(v)
			var f object
			switch k {
			case Chat:
				if e := fields(t, "type", "function"); e != nil {
					return r, e
				}
				if str(t["type"]) != "function" {
					return r, fmt.Errorf("unsupported tool")
				}
				f = obj(t["function"])
			case Responses:
				if str(t["type"]) != "function" {
					return r, fmt.Errorf("unsupported tool")
				}
				f = object{}
				for key, val := range t {
					if key != "type" {
						f[key] = val
					}
				}
			case Anthropic:
				f = object{}
				for key, val := range t {
					if key == "input_schema" {
						key = "parameters"
					}
					f[key] = val
				}
			}
			if e := fields(f, "name", "description", "parameters", "strict"); e != nil {
				return r, e
			}
			if _, e := required(f, "name"); e != nil {
				return r, e
			}
			if obj(f["parameters"]) == nil {
				return r, fmt.Errorf("missing tool schema")
			}
			r.tools = append(r.tools, f)
		}
	}
	if v, ok := m["tool_choice"]; ok {
		if s, ok := v.(string); ok {
			if s != "auto" && s != "none" && s != "required" {
				return r, fmt.Errorf("unsupported tool_choice")
			}
			r.choice = s
		} else {
			x := obj(v)
			if k == Chat {
				if e := fields(x, "type", "function"); e != nil {
					return r, e
				}
				if str(x["type"]) != "function" {
					return r, fmt.Errorf("unsupported tool_choice")
				}
				f := obj(x["function"])
				if e := fields(f, "name"); e != nil {
					return r, e
				}
				s, e := required(f, "name")
				if e != nil {
					return r, e
				}
				r.choice = object{"name": s}
			} else if k == Responses {
				if e := fields(x, "type", "name"); e != nil {
					return r, e
				}
				if str(x["type"]) != "function" {
					return r, fmt.Errorf("unsupported tool_choice")
				}
				s, e := required(x, "name")
				if e != nil {
					return r, e
				}
				r.choice = object{"name": s}
			} else {
				if e := fields(x, "type", "name", "disable_parallel_tool_use"); e != nil {
					return r, e
				}
				switch str(x["type"]) {
				case "auto", "none":
					r.choice = x["type"]
				case "any":
					r.choice = "required"
				case "tool":
					s, e := required(x, "name")
					if e != nil {
						return r, e
					}
					r.choice = object{"name": s}
				default:
					return r, fmt.Errorf("unsupported tool_choice")
				}
				if b, ok := x["disable_parallel_tool_use"].(bool); ok {
					r.options["parallel_tool_calls"] = !b
				}
			}
		}
	}
	return r, nil
}
func callPart(x object, idKey, argsKey string) (part, error) {
	id, e := required(x, idKey)
	if e != nil {
		return part{}, e
	}
	name, e := required(x, "name")
	if e != nil {
		return part{}, e
	}
	s, ok := x[argsKey].(string)
	if !ok {
		return part{}, fmt.Errorf("invalid tool arguments")
	}
	if _, e := decode([]byte(s)); e != nil {
		return part{}, fmt.Errorf("tool arguments must be JSON object: %w", e)
	}
	return part{typ: "call", id: id, name: name, args: s}, nil
}
func readMessage(k Kind, x object) (message, error) {
	m := message{role: str(x["role"])}
	if x == nil {
		return m, fmt.Errorf("invalid message")
	}
	if e := fields(x, "role", "content", "type", "id", "status", "tool_calls", "tool_call_id", "refusal"); e != nil {
		return m, e
	}
	if x["refusal"] != nil {
		return m, fmt.Errorf("refusal content cannot be preserved")
	}
	switch m.role {
	case "user", "assistant", "system", "developer", "tool":
	default:
		return m, fmt.Errorf("unsupported role %q", m.role)
	}
	if m.role == "tool" {
		if k != Chat {
			return m, fmt.Errorf("invalid tool role")
		}
		id, e := required(x, "tool_call_id")
		if e != nil {
			return m, e
		}
		s, ok := x["content"].(string)
		if !ok {
			return m, fmt.Errorf("non-text tool result unsupported")
		}
		m.parts = []part{{typ: "result", id: id, text: s}}
		return m, nil
	}
	if c := x["content"]; c != nil {
		ps, e := readContent(k, c)
		if e != nil {
			return m, e
		}
		m.parts = ps
	}
	if tc, ok := x["tool_calls"]; ok {
		if k != Chat || m.role != "assistant" {
			return m, fmt.Errorf("invalid tool_calls")
		}
		ts, ok := tc.([]any)
		if !ok {
			return m, fmt.Errorf("invalid tool_calls")
		}
		for _, v := range ts {
			t := obj(v)
			if e := fields(t, "id", "type", "function"); e != nil {
				return m, e
			}
			if str(t["type"]) != "function" {
				return m, fmt.Errorf("unsupported tool call")
			}
			f := obj(t["function"])
			if e := fields(f, "name", "arguments"); e != nil {
				return m, e
			}
			p, e := callPart(object{"id": t["id"], "name": f["name"], "arguments": f["arguments"]}, "id", "arguments")
			if e != nil {
				return m, e
			}
			m.parts = append(m.parts, p)
		}
	}
	return m, nil
}
func readContent(k Kind, v any) ([]part, error) {
	if s, ok := v.(string); ok {
		return []part{textPart(s)}, nil
	}
	a, ok := v.([]any)
	if !ok {
		return nil, fmt.Errorf("invalid content")
	}
	var ps []part
	for _, v := range a {
		x := obj(v)
		switch str(x["type"]) {
		case "text", "input_text", "output_text":
			if e := fields(x, "type", "text", "annotations"); e != nil {
				return nil, e
			}
			if an, ok := x["annotations"]; ok && len(arr(an)) != 0 {
				return nil, fmt.Errorf("annotations unsupported")
			}
			s, ok := x["text"].(string)
			if !ok {
				return nil, fmt.Errorf("invalid text")
			}
			ps = append(ps, textPart(s))
		case "image_url":
			if k != Chat {
				return nil, fmt.Errorf("invalid image block")
			}
			if e := fields(x, "type", "image_url"); e != nil {
				return nil, e
			}
			im := obj(x["image_url"])
			if e := fields(im, "url", "detail"); e != nil {
				return nil, e
			}
			if d := str(im["detail"]); d != "" && d != "auto" {
				return nil, fmt.Errorf("image detail cannot be preserved")
			}
			s, e := required(im, "url")
			if e != nil {
				return nil, e
			}
			ps = append(ps, part{typ: "image", url: s})
		case "input_image":
			if k != Responses {
				return nil, fmt.Errorf("invalid image block")
			}
			if e := fields(x, "type", "image_url", "detail"); e != nil {
				return nil, e
			}
			if d := str(x["detail"]); d != "" && d != "auto" {
				return nil, fmt.Errorf("image detail cannot be preserved")
			}
			s, e := required(x, "image_url")
			if e != nil {
				return nil, e
			}
			ps = append(ps, part{typ: "image", url: s})
		case "image":
			if k != Anthropic {
				return nil, fmt.Errorf("invalid image block")
			}
			if e := fields(x, "type", "source"); e != nil {
				return nil, e
			}
			src := obj(x["source"])
			var url string
			switch str(src["type"]) {
			case "url":
				if e := fields(src, "type", "url"); e != nil {
					return nil, e
				}
				url = str(src["url"])
			case "base64":
				if e := fields(src, "type", "media_type", "data"); e != nil {
					return nil, e
				}
				if str(src["media_type"]) == "" || str(src["data"]) == "" {
					return nil, fmt.Errorf("invalid image source")
				}
				url = "data:" + str(src["media_type"]) + ";base64," + str(src["data"])
			default:
				return nil, fmt.Errorf("unsupported image source")
			}
			if url == "" {
				return nil, fmt.Errorf("empty image URL")
			}
			ps = append(ps, part{typ: "image", url: url})
		case "tool_use":
			if k != Anthropic {
				return nil, fmt.Errorf("invalid tool block")
			}
			if e := fields(x, "type", "id", "name", "input"); e != nil {
				return nil, e
			}
			b, e := json.Marshal(x["input"])
			if e != nil {
				return nil, e
			}
			p, e := callPart(object{"id": x["id"], "name": x["name"], "arguments": string(b)}, "id", "arguments")
			if e != nil {
				return nil, e
			}
			ps = append(ps, p)
		case "tool_result":
			if k != Anthropic {
				return nil, fmt.Errorf("invalid result block")
			}
			if e := fields(x, "type", "tool_use_id", "content", "is_error"); e != nil {
				return nil, e
			}
			if x["is_error"] == true {
				return nil, fmt.Errorf("tool error flag cannot be preserved")
			}
			id, e := required(x, "tool_use_id")
			if e != nil {
				return nil, e
			}
			s, ok := x["content"].(string)
			if !ok {
				pp, e := readContent(k, x["content"])
				if e != nil {
					return nil, e
				}
				for _, p := range pp {
					if p.typ != "text" {
						return nil, fmt.Errorf("non-text tool result unsupported")
					}
					s += p.text
				}
			}
			ps = append(ps, part{typ: "result", id: id, text: s})
		default:
			return nil, fmt.Errorf("unsupported content type %q", str(x["type"]))
		}
	}
	return ps, nil
}