package llm

import (
	"encoding/json"
	"fmt"
	"reflect"
	"slices"
	"strings"
	"time"
)

// SchemaFor derives a JSON Schema from Go type T for structured output.
//
// The emitted schema is strict-mode-compatible across providers: every
// object lists all its properties as required and sets
// additionalProperties:false; optionality is expressed by pointer fields,
// which become nullable (anyOf with null) and may simply be returned as
// null by the model.
//
// Field tags: `json:"name"` / `json:"-"` control naming and omission;
// `description:"..."` documents a field for the model; `enum:"a,b,c"`
// constrains a string field to fixed values.
//
// Supported kinds: strings, bools, integer and float kinds, structs
// (nested), slices/arrays, map[string]V, pointers (nullable), time.Time
// (string, date-time), json.RawMessage and interface{} (any). Recursive
// types are rejected — no provider's structured-output mode accepts them.
func SchemaFor[T any]() (json.RawMessage, error) {
	t := reflect.TypeFor[T]()
	schema, err := schemaOf(t, "", "", nil)
	if err != nil {
		return nil, fmt.Errorf("llm: SchemaFor[%s]: %w", t, err)
	}
	out, err := json.Marshal(schema)
	if err != nil {
		return nil, fmt.Errorf("llm: SchemaFor[%s]: encode: %w", t, err)
	}
	return out, nil
}

var (
	timeType = reflect.TypeFor[time.Time]()
	rawType  = reflect.TypeFor[json.RawMessage]()
)

// schemaOf builds the schema map for one type. visiting tracks the struct
// types on the current path for cycle detection.
func schemaOf(t reflect.Type, description, enum string, visiting []reflect.Type) (map[string]any, error) {
	s := make(map[string]any)
	if description != "" {
		s["description"] = description
	}

	switch {
	case t == timeType:
		s["type"] = "string"
		s["format"] = "date-time"
		return s, nil
	case t == rawType:
		// Any JSON value.
		return s, nil
	}

	switch t.Kind() {
	case reflect.Pointer:
		inner, err := schemaOf(t.Elem(), "", enum, visiting)
		if err != nil {
			return nil, err
		}
		s["anyOf"] = []any{inner, map[string]any{"type": "null"}}
		return s, nil

	case reflect.String:
		s["type"] = "string"
		if enum != "" {
			var vals []any
			for v := range strings.SplitSeq(enum, ",") {
				vals = append(vals, strings.TrimSpace(v))
			}
			s["enum"] = vals
		}
		return s, nil

	case reflect.Bool:
		s["type"] = "boolean"
		return s, nil

	case reflect.Int, reflect.Int8, reflect.Int16, reflect.Int32, reflect.Int64,
		reflect.Uint, reflect.Uint8, reflect.Uint16, reflect.Uint32, reflect.Uint64:
		s["type"] = "integer"
		return s, nil

	case reflect.Float32, reflect.Float64:
		s["type"] = "number"
		return s, nil

	case reflect.Slice, reflect.Array:
		if t.Elem().Kind() == reflect.Uint8 {
			// []byte: base64 text on the wire.
			s["type"] = "string"
			return s, nil
		}
		items, err := schemaOf(t.Elem(), "", "", visiting)
		if err != nil {
			return nil, err
		}
		s["type"] = "array"
		s["items"] = items
		return s, nil

	case reflect.Map:
		if t.Key().Kind() != reflect.String {
			return nil, fmt.Errorf("map key type %s unsupported (only string keys)", t.Key())
		}
		values, err := schemaOf(t.Elem(), "", "", visiting)
		if err != nil {
			return nil, err
		}
		s["type"] = "object"
		s["additionalProperties"] = values
		return s, nil

	case reflect.Interface:
		// Any JSON value.
		return s, nil

	case reflect.Struct:
		if slices.Contains(visiting, t) {
			return nil, fmt.Errorf("recursive type %s (structured-output schemas cannot recurse)", t)
		}
		visiting = append(visiting, t)

		properties := make(map[string]any)
		required := []string{}
		for i := range t.NumField() {
			f := t.Field(i)
			if !f.IsExported() {
				continue
			}
			name := f.Name
			if tag, ok := f.Tag.Lookup("json"); ok {
				base, _, _ := strings.Cut(tag, ",")
				if base == "-" {
					continue
				}
				if base != "" {
					name = base
				}
			}
			fs, err := schemaOf(f.Type, f.Tag.Get("description"), f.Tag.Get("enum"), visiting)
			if err != nil {
				return nil, fmt.Errorf("field %s: %w", f.Name, err)
			}
			properties[name] = fs
			required = append(required, name)
		}
		s["type"] = "object"
		s["properties"] = properties
		s["required"] = required
		s["additionalProperties"] = false
		return s, nil

	default:
		return nil, fmt.Errorf("kind %s unsupported in structured-output schemas", t.Kind())
	}
}