package main

import (
	"bytes"
	"encoding/base64"
	"encoding/binary"
	"encoding/json"
	"errors"
	"io/ioutil"
	"log"
	"math/bits"
	"os"
	"strings"

	"lukechampine.com/muse"
	"lukechampine.com/us/merkle"
	"lukechampine.com/us/renter/proto"
	"lukechampine.com/us/renter/renterutil"
	"lukechampine.com/us/renterhost"
)

func computeLinkBits(length uint32) uint16 {
	minlength := uint32(1 << 21)
	mode := uint32(1 << 19)
	for length <= minlength && mode > (1<<12) {
		mode >>= 1
		minlength >>= 1
	}
	lengthAlign := uint32(1 << 12)
	if mode > 1<<13 {
		lengthAlign = mode >> 1
	}
	if mode > 1<<12 {
		length -= lengthAlign * 8
	}
	if length != 0 && length%lengthAlign == 0 {
		length--
	}
	length &^= lengthAlign - 1
	bitwiselength := uint16(length / lengthAlign)
	bitfield := bitwiselength << 1
	baseAlign := uint32(1 << 12)
	for baseAlign < mode {
		baseAlign <<= 1
		bitfield = (bitfield << 1) | 1
	}
	bitfield <<= 2
	return bitfield
}

func decodeLinkBits(bitfield uint16) (uint32, error) {
	version := bitfield&0b11 + 1
	bitfield >>= 2
	modeBits := bits.TrailingZeros16(^bitfield)
	bitfield >>= modeBits + 1
	length := uint32(bitfield&0b111) + 1
	bitfield >>= 3
	offset := uint32(bitfield)

	if version != 1 {
		return 0, errors.New("unrecognized version")
	} else if modeBits > 7 {
		return 0, errors.New("invalid mode bits")
	}

	// compute offset and length
	offset *= uint32(4096) << modeBits
	if modeBits > 0 {
		fetchSizeAlign := uint32(4096) << (modeBits - 1)
		length = length*fetchSizeAlign + fetchSizeAlign<<3
	} else {
		length *= 4096
	}
	if offset+length > renterhost.SectorSize {
		return 0, errors.New("invalid offset+length")
	} else if offset > 0 {
		return 0, errors.New("can't handle non-zero offset")
	}
	return length, nil
}

type linkfileMetadata struct {
	Filename string      `json:"filename,omitempty"`
	Mode     os.FileMode `json:"mode,omitempty"`
}

type linkfileLayout struct {
	Version      uint8
	Filesize     uint64
	MetadataSize uint64
	FanoutSize   uint64
	Unused       [74]byte
}

func encodeLinkFile(name string, data []byte) *[renterhost.SectorSize]byte {
	metadata, _ := json.Marshal(linkfileMetadata{
		Filename: name,
		Mode:     0666,
	})
	layout := linkfileLayout{
		Version:      1,
		Filesize:     uint64(len(data)),
		MetadataSize: uint64(len(metadata)),
	}
	var sector [renterhost.SectorSize]byte
	buf := bytes.NewBuffer(sector[:0])
	binary.Write(buf, binary.LittleEndian, layout)
	buf.Write(metadata)
	buf.Write(data)
	return &sector
}

func uploadSector(mc *muse.Client, sector *[renterhost.SectorSize]byte) error {
	currentHeight, err := mc.SHARD().ChainHeight()
	if err != nil {
		return err
	}
	contracts, err := mc.Contracts("default")
	if err != nil {
		return err
	}
	results := make(chan *renterutil.HostError)
	for _, c := range contracts {
		go func(c muse.Contract) {
			s, err := proto.NewSession(c.HostAddress, c.HostKey, c.ID, c.RenterKey, currentHeight)
			if err == nil {
				// Append is a convenient wrapper around the Write RPC
				_, err = s.Append(sector)
				s.Close()
			}
			results <- &renterutil.HostError{HostKey: c.HostKey, Err: err}
		}(c)
	}
	var errs renterutil.HostErrorSet
	for range contracts {
		if err := <-results; err.Err != nil {
			errs = append(errs, err)
		}
	}
	if len(errs) == len(contracts) {
		return errs // none of the uploads succeeded
	}
	return nil
}

func encodeLink(data []byte, sector *[renterhost.SectorSize]byte) string {
	var linkdata [34]byte
	binary.LittleEndian.PutUint16(linkdata[:2], computeLinkBits(uint32(len(data))))
	root := merkle.SectorRoot(sector)
	copy(linkdata[2:], root[:])
	return "sia://" + base64.RawURLEncoding.EncodeToString(linkdata[:])
}

func decodeLink(link string) (req renterhost.RPCReadRequestSection, err error) {
	data, err := base64.RawURLEncoding.DecodeString(strings.TrimPrefix(link, "sia://"))
	if err != nil {
		return
	} else if len(data) != 34 {
		err = errors.New("invalid Link")
		return
	}
	req.Length, err = decodeLinkBits(binary.LittleEndian.Uint16(data[:2]))
	copy(req.MerkleRoot[:], data[2:])
	return
}

func downloadSector(mc *muse.Client, req renterhost.RPCReadRequestSection) ([]byte, error) {
	contracts, err := mc.Contracts("default")
	if err != nil {
		return nil, err
	}
	var errs renterutil.HostErrorSet
	for _, c := range contracts {
		s, err := proto.NewSession(c.HostAddress, c.HostKey, c.ID, c.RenterKey, 0)
		if err == nil {
			var buf bytes.Buffer
			err = s.Read(&buf, []renterhost.RPCReadRequestSection{req})
			s.Close()
			if err == nil {
				return buf.Bytes(), nil
			}
		}
		errs = append(errs, &renterutil.HostError{HostKey: c.HostKey, Err: err})
	}
	return nil, errs
}

func decodeLinkFile(sector []byte) (meta linkfileMetadata, data []byte, err error) {
	buf := bytes.NewBuffer(sector)
	var lfl linkfileLayout
	if err = binary.Read(buf, binary.LittleEndian, &lfl); err != nil {
		return
	}
	if lfl.FanoutSize > 0 {
		err = errors.New("can only handle single-sector files")
		return
	}
	if err = json.Unmarshal(buf.Next(int(lfl.MetadataSize)), &meta); err != nil {
		return
	}
	data = buf.Next(int(lfl.Filesize))
	return
}

func main() {
	mc := muse.NewClient("<my muse server>")
	switch os.Args[1] {
	case "upload":
		data, err := ioutil.ReadFile(os.Args[2])
		if err != nil {
			log.Fatal(err)
		}
		sector := encodeLinkFile(os.Args[2], data)
		err = uploadSector(mc, sector)
		log.Println(encodeLink(data, sector), err)

	case "download":
		req, err := decodeLink(os.Args[2])
		if err != nil {
			log.Fatal(err)
		}
		sector, err := downloadSector(mc, req)
		if err != nil {
			log.Fatal(err)
		}
		meta, data, err := decodeLinkFile(sector)
		if err != nil {
			log.Fatal(err)
		}
		ioutil.WriteFile(meta.Filename, data, meta.Mode)
	}
}