package main

import (
	"encoding/binary"
	"fmt"
	"io"
	"math"
	"os"
	"strconv"
	"strings"
)

const (
	GGUFMagic        uint32 = 0x46554747
	DefaultAlignment uint64 = 32
)

type ValueType uint32

const (
	TypeUint8   ValueType = 0
	TypeInt8    ValueType = 1
	TypeUint16  ValueType = 2
	TypeInt16   ValueType = 3
	TypeUint32  ValueType = 4
	TypeInt32   ValueType = 5
	TypeFloat32 ValueType = 6
	TypeBool    ValueType = 7
	TypeString  ValueType = 8
	TypeArray   ValueType = 9
	TypeUint64  ValueType = 10
	TypeInt64   ValueType = 11
	TypeFloat64 ValueType = 12
)

var typeNames = map[ValueType]string{
	TypeUint8: "uint8", TypeInt8: "int8",
	TypeUint16: "uint16", TypeInt16: "int16",
	TypeUint32: "uint32", TypeInt32: "int32",
	TypeFloat32: "float32", TypeBool: "bool",
	TypeString: "string", TypeArray: "array",
	TypeUint64: "uint64", TypeInt64: "int64",
	TypeFloat64: "float64",
}

// ScalarTypes returns all non-array value types for the type picker.
var ScalarTypes = []ValueType{
	TypeString, TypeBool,
	TypeUint8, TypeInt8, TypeUint16, TypeInt16,
	TypeUint32, TypeInt32, TypeUint64, TypeInt64,
	TypeFloat32, TypeFloat64,
}

func (t ValueType) String() string {
	if s, ok := typeNames[t]; ok {
		return s
	}
	return fmt.Sprintf("unknown(%d)", t)
}

// MetaValue holds a single GGUF metadata value.
type MetaValue struct {
	Type     ValueType
	Data     interface{}
	ElemType ValueType // only for arrays
}

// DisplayString returns a human-readable representation of the value.
func (v MetaValue) DisplayString() string {
	switch v.Type {
	case TypeBool:
		if v.Data.(bool) {
			return "true"
		}
		return "false"
	case TypeString:
		return v.Data.(string)
	case TypeFloat32:
		return strconv.FormatFloat(float64(v.Data.(float32)), 'g', -1, 32)
	case TypeFloat64:
		return strconv.FormatFloat(v.Data.(float64), 'g', -1, 64)
	case TypeArray:
		arr := v.Data.([]MetaValue)
		return fmt.Sprintf("[%s x %d]", v.ElemType, len(arr))
	default:
		return fmt.Sprintf("%v", v.Data)
	}
}

// TypeName returns the display name for this value's type.
func (v MetaValue) TypeName() string {
	if v.Type == TypeArray {
		return fmt.Sprintf("array[%s]", v.ElemType)
	}
	return v.Type.String()
}

// KVPair is a metadata key-value pair.
type KVPair struct {
	Key   string
	Value MetaValue
}

// TensorInfo describes a single tensor in the file.
type TensorInfo struct {
	Name   string
	NDims  uint32
	Dims   []uint64
	DType  uint32
	Offset uint64
}

// GGUFFile represents a parsed GGUF file's header and metadata.
// Tensor data is never loaded into memory.
type GGUFFile struct {
	Version     uint32
	TensorCount uint64
	Metadata    []KVPair
	Tensors     []TensorInfo
	Alignment   uint64

	sourcePath      string
	tensorDataStart int64
	fileSize        int64
}

// ParseGGUF reads only the header, metadata, and tensor info from a GGUF file.
// No tensor data is loaded into memory.
func ParseGGUF(path string) (*GGUFFile, error) {
	f, err := os.Open(path)
	if err != nil {
		return nil, err
	}
	defer f.Close()

	info, err := f.Stat()
	if err != nil {
		return nil, err
	}

	g := &GGUFFile{
		sourcePath: path,
		fileSize:   info.Size(),
		Alignment:  DefaultAlignment,
	}

	var magic uint32
	if err := binary.Read(f, binary.LittleEndian, &magic); err != nil {
		return nil, fmt.Errorf("reading magic: %w", err)
	}
	if magic != GGUFMagic {
		return nil, fmt.Errorf("not a GGUF file (magic: 0x%08X, expected 0x%08X)", magic, GGUFMagic)
	}

	if err := binary.Read(f, binary.LittleEndian, &g.Version); err != nil {
		return nil, fmt.Errorf("reading version: %w", err)
	}
	if g.Version < 2 || g.Version > 3 {
		return nil, fmt.Errorf("unsupported GGUF version %d (supported: 2, 3)", g.Version)
	}

	if err := binary.Read(f, binary.LittleEndian, &g.TensorCount); err != nil {
		return nil, fmt.Errorf("reading tensor count: %w", err)
	}
	var metaCount uint64
	if err := binary.Read(f, binary.LittleEndian, &metaCount); err != nil {
		return nil, fmt.Errorf("reading metadata count: %w", err)
	}

	// Read metadata KV pairs
	g.Metadata = make([]KVPair, 0, metaCount)
	for i := uint64(0); i < metaCount; i++ {
		key, err := readString(f)
		if err != nil {
			return nil, fmt.Errorf("metadata[%d] key: %w", i, err)
		}
		var vtype uint32
		if err := binary.Read(f, binary.LittleEndian, &vtype); err != nil {
			return nil, fmt.Errorf("metadata[%d] type: %w", i, err)
		}
		val, err := readValue(f, ValueType(vtype))
		if err != nil {
			return nil, fmt.Errorf("metadata[%d] (%s) value: %w", i, key, err)
		}
		g.Metadata = append(g.Metadata, KVPair{Key: key, Value: val})

		if key == "general.alignment" {
			switch val.Type {
			case TypeUint32:
				g.Alignment = uint64(val.Data.(uint32))
			case TypeUint64:
				g.Alignment = val.Data.(uint64)
			}
		}
	}

	// Read tensor infos
	g.Tensors = make([]TensorInfo, 0, g.TensorCount)
	for i := uint64(0); i < g.TensorCount; i++ {
		name, err := readString(f)
		if err != nil {
			return nil, fmt.Errorf("tensor[%d] name: %w", i, err)
		}
		var ndims uint32
		if err := binary.Read(f, binary.LittleEndian, &ndims); err != nil {
			return nil, fmt.Errorf("tensor[%d] ndims: %w", i, err)
		}
		dims := make([]uint64, ndims)
		for d := uint32(0); d < ndims; d++ {
			if err := binary.Read(f, binary.LittleEndian, &dims[d]); err != nil {
				return nil, fmt.Errorf("tensor[%d] dim[%d]: %w", i, d, err)
			}
		}
		var dtype uint32
		if err := binary.Read(f, binary.LittleEndian, &dtype); err != nil {
			return nil, fmt.Errorf("tensor[%d] dtype: %w", i, err)
		}
		var offset uint64
		if err := binary.Read(f, binary.LittleEndian, &offset); err != nil {
			return nil, fmt.Errorf("tensor[%d] offset: %w", i, err)
		}
		g.Tensors = append(g.Tensors, TensorInfo{
			Name: name, NDims: ndims, Dims: dims, DType: dtype, Offset: offset,
		})
	}

	// Record where tensor data begins (aligned)
	pos, err := f.Seek(0, io.SeekCurrent)
	if err != nil {
		return nil, fmt.Errorf("getting position: %w", err)
	}
	g.tensorDataStart = alignUp(pos, int64(g.Alignment))

	return g, nil
}

// Save writes the GGUF file with (potentially modified) metadata to outPath.
// Tensor data is streamed from the original file — never fully loaded.
func (g *GGUFFile) Save(outPath string) error {
	out, err := os.Create(outPath)
	if err != nil {
		return err
	}
	defer out.Close()

	cw := &countingWriter{w: out}

	// Header
	if err := binary.Write(cw, binary.LittleEndian, GGUFMagic); err != nil {
		return fmt.Errorf("writing magic: %w", err)
	}
	if err := binary.Write(cw, binary.LittleEndian, g.Version); err != nil {
		return fmt.Errorf("writing version: %w", err)
	}
	if err := binary.Write(cw, binary.LittleEndian, g.TensorCount); err != nil {
		return fmt.Errorf("writing tensor count: %w", err)
	}
	if err := binary.Write(cw, binary.LittleEndian, uint64(len(g.Metadata))); err != nil {
		return fmt.Errorf("writing metadata count: %w", err)
	}

	// Metadata KV pairs
	for _, kv := range g.Metadata {
		if err := writeString(cw, kv.Key); err != nil {
			return fmt.Errorf("writing key %q: %w", kv.Key, err)
		}
		if err := binary.Write(cw, binary.LittleEndian, uint32(kv.Value.Type)); err != nil {
			return fmt.Errorf("writing type for %q: %w", kv.Key, err)
		}
		if err := writeValue(cw, kv.Value); err != nil {
			return fmt.Errorf("writing value for %q: %w", kv.Key, err)
		}
	}

	// Tensor infos
	for _, t := range g.Tensors {
		if err := writeString(cw, t.Name); err != nil {
			return fmt.Errorf("writing tensor name: %w", err)
		}
		if err := binary.Write(cw, binary.LittleEndian, t.NDims); err != nil {
			return err
		}
		for _, d := range t.Dims {
			if err := binary.Write(cw, binary.LittleEndian, d); err != nil {
				return err
			}
		}
		if err := binary.Write(cw, binary.LittleEndian, t.DType); err != nil {
			return err
		}
		if err := binary.Write(cw, binary.LittleEndian, t.Offset); err != nil {
			return err
		}
	}

	// Alignment padding
	aligned := alignUp(cw.count, int64(g.Alignment))
	padLen := aligned - cw.count
	if padLen > 0 {
		if _, err := out.Write(make([]byte, padLen)); err != nil {
			return fmt.Errorf("writing padding: %w", err)
		}
	}

	// Stream-copy tensor data from original file
	src, err := os.Open(g.sourcePath)
	if err != nil {
		return fmt.Errorf("opening source for tensor data: %w", err)
	}
	defer src.Close()

	if _, err := src.Seek(g.tensorDataStart, io.SeekStart); err != nil {
		return fmt.Errorf("seeking to tensor data: %w", err)
	}

	tensorLen := g.fileSize - g.tensorDataStart
	if tensorLen > 0 {
		if _, err := io.CopyN(out, src, tensorLen); err != nil {
			return fmt.Errorf("copying tensor data: %w", err)
		}
	}

	return nil
}

// ParseValueFromString parses a user-entered string into a MetaValue of the given type.
func ParseValueFromString(s string, t ValueType) (MetaValue, error) {
	s = strings.TrimSpace(s)
	switch t {
	case TypeUint8:
		v, err := strconv.ParseUint(s, 10, 8)
		return MetaValue{Type: t, Data: uint8(v)}, err
	case TypeInt8:
		v, err := strconv.ParseInt(s, 10, 8)
		return MetaValue{Type: t, Data: int8(v)}, err
	case TypeUint16:
		v, err := strconv.ParseUint(s, 10, 16)
		return MetaValue{Type: t, Data: uint16(v)}, err
	case TypeInt16:
		v, err := strconv.ParseInt(s, 10, 16)
		return MetaValue{Type: t, Data: int16(v)}, err
	case TypeUint32:
		v, err := strconv.ParseUint(s, 10, 32)
		return MetaValue{Type: t, Data: uint32(v)}, err
	case TypeInt32:
		v, err := strconv.ParseInt(s, 10, 32)
		return MetaValue{Type: t, Data: int32(v)}, err
	case TypeUint64:
		v, err := strconv.ParseUint(s, 10, 64)
		return MetaValue{Type: t, Data: v}, err
	case TypeInt64:
		v, err := strconv.ParseInt(s, 10, 64)
		return MetaValue{Type: t, Data: v}, err
	case TypeFloat32:
		v, err := strconv.ParseFloat(s, 32)
		return MetaValue{Type: t, Data: float32(v)}, err
	case TypeFloat64:
		v, err := strconv.ParseFloat(s, 64)
		return MetaValue{Type: t, Data: v}, err
	case TypeBool:
		v, err := strconv.ParseBool(s)
		return MetaValue{Type: t, Data: v}, err
	case TypeString:
		return MetaValue{Type: t, Data: s}, nil
	default:
		return MetaValue{}, fmt.Errorf("cannot parse type %s from string", t)
	}
}

// --- internal helpers ---

func readString(r io.Reader) (string, error) {
	var length uint64
	if err := binary.Read(r, binary.LittleEndian, &length); err != nil {
		return "", err
	}
	if length > 1<<30 { // sanity limit: 1 GB
		return "", fmt.Errorf("string length %d exceeds sanity limit", length)
	}
	buf := make([]byte, length)
	if _, err := io.ReadFull(r, buf); err != nil {
		return "", err
	}
	return string(buf), nil
}

func readValue(r io.Reader, t ValueType) (MetaValue, error) {
	switch t {
	case TypeUint8:
		var v uint8
		err := binary.Read(r, binary.LittleEndian, &v)
		return MetaValue{Type: t, Data: v}, err
	case TypeInt8:
		var v int8
		err := binary.Read(r, binary.LittleEndian, &v)
		return MetaValue{Type: t, Data: v}, err
	case TypeUint16:
		var v uint16
		err := binary.Read(r, binary.LittleEndian, &v)
		return MetaValue{Type: t, Data: v}, err
	case TypeInt16:
		var v int16
		err := binary.Read(r, binary.LittleEndian, &v)
		return MetaValue{Type: t, Data: v}, err
	case TypeUint32:
		var v uint32
		err := binary.Read(r, binary.LittleEndian, &v)
		return MetaValue{Type: t, Data: v}, err
	case TypeInt32:
		var v int32
		err := binary.Read(r, binary.LittleEndian, &v)
		return MetaValue{Type: t, Data: v}, err
	case TypeFloat32:
		var v float32
		err := binary.Read(r, binary.LittleEndian, &v)
		return MetaValue{Type: t, Data: v}, err
	case TypeBool:
		var v uint8
		err := binary.Read(r, binary.LittleEndian, &v)
		return MetaValue{Type: t, Data: v != 0}, err
	case TypeString:
		s, err := readString(r)
		return MetaValue{Type: t, Data: s}, err
	case TypeUint64:
		var v uint64
		err := binary.Read(r, binary.LittleEndian, &v)
		return MetaValue{Type: t, Data: v}, err
	case TypeInt64:
		var v int64
		err := binary.Read(r, binary.LittleEndian, &v)
		return MetaValue{Type: t, Data: v}, err
	case TypeFloat64:
		var v float64
		err := binary.Read(r, binary.LittleEndian, &v)
		return MetaValue{Type: t, Data: v}, err
	case TypeArray:
		var elemType uint32
		if err := binary.Read(r, binary.LittleEndian, &elemType); err != nil {
			return MetaValue{}, fmt.Errorf("array elem type: %w", err)
		}
		var count uint64
		if err := binary.Read(r, binary.LittleEndian, &count); err != nil {
			return MetaValue{}, fmt.Errorf("array count: %w", err)
		}
		if count > 1<<30 {
			return MetaValue{}, fmt.Errorf("array count %d exceeds sanity limit", count)
		}
		arr := make([]MetaValue, count)
		for i := uint64(0); i < count; i++ {
			v, err := readValue(r, ValueType(elemType))
			if err != nil {
				return MetaValue{}, fmt.Errorf("array[%d]: %w", i, err)
			}
			arr[i] = v
		}
		return MetaValue{Type: TypeArray, Data: arr, ElemType: ValueType(elemType)}, nil
	default:
		return MetaValue{}, fmt.Errorf("unknown value type: %d", t)
	}
}

func writeString(w io.Writer, s string) error {
	if err := binary.Write(w, binary.LittleEndian, uint64(len(s))); err != nil {
		return err
	}
	_, err := w.Write([]byte(s))
	return err
}

func writeValue(w io.Writer, v MetaValue) error {
	switch v.Type {
	case TypeUint8:
		return binary.Write(w, binary.LittleEndian, v.Data.(uint8))
	case TypeInt8:
		return binary.Write(w, binary.LittleEndian, v.Data.(int8))
	case TypeUint16:
		return binary.Write(w, binary.LittleEndian, v.Data.(uint16))
	case TypeInt16:
		return binary.Write(w, binary.LittleEndian, v.Data.(int16))
	case TypeUint32:
		return binary.Write(w, binary.LittleEndian, v.Data.(uint32))
	case TypeInt32:
		return binary.Write(w, binary.LittleEndian, v.Data.(int32))
	case TypeFloat32:
		return binary.Write(w, binary.LittleEndian, v.Data.(float32))
	case TypeBool:
		b := uint8(0)
		if v.Data.(bool) {
			b = 1
		}
		return binary.Write(w, binary.LittleEndian, b)
	case TypeString:
		return writeString(w, v.Data.(string))
	case TypeUint64:
		return binary.Write(w, binary.LittleEndian, v.Data.(uint64))
	case TypeInt64:
		return binary.Write(w, binary.LittleEndian, v.Data.(int64))
	case TypeFloat64:
		return binary.Write(w, binary.LittleEndian, v.Data.(float64))
	case TypeArray:
		arr := v.Data.([]MetaValue)
		if err := binary.Write(w, binary.LittleEndian, uint32(v.ElemType)); err != nil {
			return err
		}
		if err := binary.Write(w, binary.LittleEndian, uint64(len(arr))); err != nil {
			return err
		}
		for i, elem := range arr {
			if err := writeValue(w, elem); err != nil {
				return fmt.Errorf("array[%d]: %w", i, err)
			}
		}
		return nil
	default:
		return fmt.Errorf("unknown type: %d", v.Type)
	}
}

func alignUp(pos, align int64) int64 {
	return (pos + align - 1) & ^(align - 1)
}

type countingWriter struct {
	w     io.Writer
	count int64
}

func (cw *countingWriter) Write(p []byte) (int, error) {
	n, err := cw.w.Write(p)
	cw.count += int64(n)
	return n, err
}

// humanSize formats a byte count as a human-readable string.
func humanSize(b int64) string {
	const unit = 1024
	if b < unit {
		return fmt.Sprintf("%d B", b)
	}
	div, exp := int64(unit), 0
	for n := b / unit; n >= unit; n /= unit {
		div *= unit
		exp++
	}
	return fmt.Sprintf("%.1f %ciB", float64(b)/float64(div), "KMGTPE"[exp])
}

// ggufTypeName returns the GGUF tensor type name.
func ggufTypeName(t uint32) string {
	names := map[uint32]string{
		0: "F32", 1: "F16", 2: "Q4_0", 3: "Q4_1",
		6: "Q5_0", 7: "Q5_1", 8: "Q8_0", 9: "Q8_1",
		10: "Q2_K", 11: "Q3_K", 12: "Q4_K", 13: "Q5_K", 14: "Q6_K", 15: "Q8_K",
		16: "IQ2_XXS", 17: "IQ2_XS", 18: "IQ3_XXS", 19: "IQ1_S", 20: "IQ4_NL",
		21: "IQ3_S", 22: "IQ2_S", 23: "IQ4_XS",
		24: "I8", 25: "I16", 26: "I32", 27: "I64",
		28: "F64", 29: "IQ1_M",
	}
	if s, ok := names[t]; ok {
		return s
	}
	return fmt.Sprintf("type_%d", t)
}

// NaN/Inf-safe float display
func init() {
	_ = math.NaN()
}