package bmap

import (
	"archive/tar"
	"bytes"
	"compress/bzip2"
	"compress/gzip"
	"crypto/sha1"
	"encoding/hex"
	"encoding/xml"
	"errors"
	"fmt"
	"io"
	"io/ioutil"
	"net/http"
	"os"
	"regexp"
	"strings"
	"syscall"

	"github.com/frostschutz/go-fibmap"
	//"github.com/uli-go/xz/lzma"
	"code.google.com/p/lzma"
)

const defaultBMAPHash = "0000000000000000000000000000000000000000"

// BlockRange represents a range of blocks of written data
type BlockRange struct {
	XMLName xml.Name `xml:"Range"`
	Start   int64
	End     int64
	Hash    string `xml:"sha1,attr"`
}

// Range is an intermediate data type for xml conversion
type Range struct {
	XMLName xml.Name `xml:"Range"`
	Range   string   `xml:",chardata"`
	Hash    string   `xml:"sha1,attr,omitempty"`
}

// MarshalXML is used to convert BlockRange to a Range object for XML output
func (b *BlockRange) MarshalXML(e *xml.Encoder, start xml.StartElement) error {
	//fmt.Println(e)
	//fmt.Println(start)

	r := ""
	if b.Start == b.End {
		r = fmt.Sprintf("%d", b.Start)
	} else {
		r = fmt.Sprintf("%d-%d", b.Start, b.End)
	}
	bl := Range{Range: r, Hash: b.Hash}
	err := e.Encode(bl)

	return err
}

// UnmarshalXML is used to convert a Range object from XML input into a BlockRange object
func (b *BlockRange) UnmarshalXML(d *xml.Decoder, start xml.StartElement) error {

	r := Range{}
	err := d.DecodeElement(&r, &start)
	if err != nil {
		fmt.Println(err.Error())
		return nil
	}
	b.Hash = r.Hash
	type pair struct {
		First  int64
		Second int64
	}
	//p := []int64{}
	p := pair{}

	n, err := fmt.Sscanf(r.Range, "%d-%d", &p.First, &p.Second)
	if err != nil || n != 2 {

		n, err := fmt.Sscanf(r.Range, "%d", &p.First)
		if err != nil || n != 1 { //more than one block
			fmt.Println("Multiple Blocks:", err.Error())
			return err
		}

		b.Start = p.First
		b.End = p.First

	} else {
		b.Start = p.First
		b.End = p.Second
	}

	return nil
}

// BlockMap represents a list of Block Ranges
type BlockMap struct {
	Range []BlockRange
}

// BMap contains all of the information that needs to be represented in the bmap file
type BMap struct {
	XMLName           xml.Name `xml:"bmap"`
	Version           string   `xml:"version,attr"`
	ImageSize         int64
	BlockSize         int64
	BlocksCount       int64
	MappedBlocksCount int64
	BmapFileSHA1      string
	BlockMap          BlockMap `xml:"BlockMap"`
}

// NewBMap creates a new BMap representation by reading an image file
func NewBMap(filename string, addChecksum bool) (BMap, error) {
	stat := syscall.Stat_t{}
	syscall.Stat(filename, &stat)
	blockSize := stat.Blksize

	fd, err := os.Open(filename)
	defer fd.Close()
	if err != nil {
		return BMap{}, err
	}

	blockMap, total := getBlockMap(fd, blockSize, addChecksum)
	return BMap{Version: "1.3", ImageSize: stat.Size, BlockSize: blockSize, BlocksCount: stat.Size / blockSize, MappedBlocksCount: total, BmapFileSHA1: defaultBMAPHash, BlockMap: BlockMap{blockMap}}, nil
}

// XMLOutput converts a BMap object to it's corresponding XML output
func (b *BMap) XMLOutput() ([]byte, error) {
	output, err := xml.MarshalIndent(b, "", "    ")
	if err != nil {
		fmt.Printf("error: %v\n", err)
	}
	return append([]byte(xml.Header), output...), err
}

// Write generates initial XML output, takes the SHA signature, and writes updated XML to file
func (b *BMap) Write(outputFilename string) error {
	output, err := b.XMLOutput()
	if err != nil {
		return err
	}

	bmapHash := getSHA1Hash(bytes.NewBuffer(output))

	b.BmapFileSHA1 = bmapHash
	output, err = b.XMLOutput()
	if err != nil {
		return err
	}
	outputFile, err := os.Create(outputFilename)
	defer outputFile.Close()
	if err != nil {
		return err
	}

	outputFile.Write(output)

	return err
}

// LoadFromFile reads and converts a bmap file to a BMap object
func LoadFromFile(filename string, verify bool) (BMap, error) {
	b := BMap{}
	data, err := ioutil.ReadFile(filename)
	if err != nil {
		return b, err
	}

	err = xml.Unmarshal(data, &b)
	if err != nil {
		return b, err
	}

	if verify {
		// Verify SHA Hash

		bmapHash := b.BmapFileSHA1

		b.BmapFileSHA1 = defaultBMAPHash
		output, err := b.XMLOutput()
		if err != nil {
			//fmt.Printf("error: %v\n", err)
			return b, err
		}

		if bmapHash != getSHA1Hash(bytes.NewBuffer(output)) {
			return b, errors.New("XML SHA Signature does not match")
		}

		b.BmapFileSHA1 = bmapHash
	}

	return b, err
}

// LoadFromReader reads and converts a bmap in an io.Reader to a BMap object
func LoadFromReader(r io.Reader, verify bool) (BMap, error) {
	b := BMap{}
	data, err := ioutil.ReadAll(r)
	if err != nil {
		return b, err
	}

	err = xml.Unmarshal(data, &b)
	if err != nil {
		return b, err
	}

	if verify {
		// Verify SHA Hash

		bmapHash := b.BmapFileSHA1

		b.BmapFileSHA1 = defaultBMAPHash
		output, err := b.XMLOutput()
		if err != nil {
			//fmt.Printf("error: %v\n", err)
			return b, err
		}

		if bmapHash != getSHA1Hash(bytes.NewBuffer(output)) {
			return b, errors.New("XML SHA Signature does not match")
		}

		b.BmapFileSHA1 = bmapHash
	}

	return b, err
}

// loadInput is designed to handle loading local and remote files
func loadInput(input string) (io.ReadCloser, error) {
	var reader io.ReadCloser

	inFile, err := os.Open(input)
	//defer inFile.Close()
	reader = inFile
	if err != nil { // Not local file. Try to grab remote file
		switch {
		case strings.HasPrefix(input, "http"):
			//fmt.Println("Get URL", input)
			resp, err := http.Get(input)
			if err != nil {
				return reader, errors.New("Unable to download image from " + input)
			}
			reader = resp.Body
			//fmt.Println("Retrieved", resp.ContentLength, "bytes")

		}
		//return reader, err
	}

	return reader, nil
}

// decompressInput "unwraps" - aka decompresses and unarchives image files
func decompressInput(input string, r io.Reader) (io.Reader, error) {
	var reader io.Reader
	var err error

	switch {
	case strings.HasSuffix(input, ".tar.gz"), strings.HasSuffix(input, ".tgz"):

		reader, err = getTarGZReader(r)

	case strings.HasSuffix(input, ".gz"), strings.HasSuffix(input, ".gzip"):

		reader, err = getGZReader(r)

	case strings.HasSuffix(input, ".xz"), strings.HasSuffix(input, ".lzma"):

		reader, err = getLZMAReader(r)

	default:
		reader = r

	}

	return reader, err
}

// Copy copies an input file and uses the BMap data to create a new image file
func (b *BMap) Copy(input string, output string, verify bool) error {

	var reader io.Reader

	r, err := loadInput(input)
	defer r.Close()
	if err != nil {
		return err
	}

	reader, err = decompressInput(input, r)
	if err != nil {
		return err
	}

	outFile, err := os.Create(output)
	defer outFile.Close()
	if err != nil {
		return err
	}

	block, err := isBlockDevice(outFile)
	if err != nil {
		return err
	}

	if !block { // Can't truncate block devices
		err = outFile.Truncate(b.ImageSize)
		if err != nil {
			return err
		}
	}

	place := int64(0)
	for _, block := range b.BlockMap.Range {

		diff := block.Start*b.BlockSize - place
		place = block.Start * b.BlockSize

		_, err := moveReaderForward(reader, diff)
		if err != nil {
			return err
		}
		_, err = outFile.Seek(diff, os.SEEK_CUR)
		if err != nil {
			return err
		}
		length := (block.End - block.Start + 1) * b.BlockSize
		place += length
		if len(block.Hash) != 0 && verify {
			//println("checking hash")
			// Verify hash sum
			//var checksumReader bytes.Buffer
			checksumReader := io.TeeReader(reader, outFile)
			h := sha1.New()
			io.CopyN(h, checksumReader, length)
			if block.Hash != hex.EncodeToString(h.Sum(nil)) {
				return fmt.Errorf("Checksum mismatch for blockrange %d-%d", block.Start, block.End)
			}
		} else {
			written, err := io.CopyN(outFile, reader, length)
			if err != nil {
				return err

			}
			if written != length {
				return errors.New("Unable to copy")

			}
		}

	}

	outFile.Sync()
	return nil
}

// Copy is used to copy images to a destination when a bmap file is unavailable
func Copy(input string, output string) error {

	reader, err := loadInput(input)
	defer reader.Close()
	if err != nil {
		return nil
	}

	seeker, err := decompressInput(input, reader)
	if err != nil {
		return err
	}

	outFile, err := os.Create(output)
	defer outFile.Close()
	if err != nil {
		return err
	}

	io.Copy(outFile, seeker)

	outFile.Sync()

	return nil
}

func getTarReader(reader io.Reader) (io.ReadSeeker, error) {
	tarReader := tar.NewReader(reader)
	tarReader.Next()

	data, err := ioutil.ReadAll(tarReader)
	if err != nil {
		return nil, err
	}
	return bytes.NewReader(data), nil

}

func getTarGZReader(reader io.Reader) (io.Reader, error) {
	gzReader, err := gzip.NewReader(reader)
	if err != nil {
		return nil, err
	}
	tarReader := tar.NewReader(gzReader)
	tarReader.Next()

	return tarReader, nil

}

func getGZReader(reader io.Reader) (io.Reader, error) {
	return gzip.NewReader(reader)
}

func getTarBZReader(reader io.Reader) (io.Reader, error) {
	bzReader := bzip2.NewReader(reader)

	tarReader := tar.NewReader(bzReader)
	tarReader.Next()

	return tarReader, nil

}

func getBZReader(reader io.Reader) (io.Reader, error) {
	bzReader := bzip2.NewReader(reader)

	return bzReader, nil
}

func getLZMAReader(reader io.Reader) (io.Reader, error) {

	/*
		return lzma.NewReader(reader)

	*/
	xzReader := lzma.NewReader(reader)

	return xzReader, nil
}

func getXZReader(reader io.Reader) (io.ReadSeeker, error) {
	//return xz.NewSeekReader(reader)
	return nil, nil
}

// getBlockMap finds all of the ranges of written blocks in a file/image
func getBlockMap(fd *os.File, blockSize int64, addChecksum bool) ([]BlockRange, int64) {
	//fd, _ := os.Open(filename)
	f := fibmap.NewFibmapFile(fd)

	holes := f.SeekDataHole()
	blockMap := make([]BlockRange, len(holes)/2)
	currentBlockRange := BlockRange{}
	mappedBlocks := int64(0)
	for i, v := range holes {
		if i%2 == 0 { //BlockRange found @
			currentBlockRange = BlockRange{Start: v / blockSize}
		} else { // Length of block
			length := v / blockSize
			currentBlockRange.End = currentBlockRange.Start + length - 1

			if addChecksum {
				h := sha1.New()
				fd.Seek(currentBlockRange.Start*blockSize, 0)
				io.CopyN(h, fd, v)
				currentBlockRange.Hash = hex.EncodeToString(h.Sum(nil))
			}

			blockMap[i/2] = currentBlockRange
			mappedBlocks += length
		}
	}
	return blockMap, mappedBlocks
}

// getSHA1Hash returns a hex encoded hash from an io.Reader
func getSHA1Hash(r io.Reader) string {
	h := sha1.New()
	io.Copy(h, r)
	return hex.EncodeToString(h.Sum(nil))
}

// moveReaderForward provides seek capabilities for an io.Reader, even ones that aren't an io.Seeker
func moveReaderForward(r io.Reader, count int64) (int64, error) {

	seeker, ok := r.(io.Seeker)
	if ok {
		return seeker.Seek(count, os.SEEK_CUR)
	} else {
		return io.CopyN(ioutil.Discard, r, count)
	}

}

// isBlockDevice checks if a file descriptor is a block device
func isBlockDevice(fd *os.File) (bool, error) {
	block := false
	s, err := fd.Stat()
	if err != nil {
		return block, err
	}
	block = (s.Mode() & os.ModeDevice) != 0
	return block, nil
}

// CleanBMap cleans spaces out of bmaptools bmap file and updates SHA Hash
func CleanBMap(input string) error {
	data, err := ioutil.ReadFile(input)
	if err != nil {
		return err
	}
	re := regexp.MustCompile(`>\s*([\d\w-]+)\s*<`)
	fixed := re.ReplaceAllString(string(data), ">${1}<")
	b, err := LoadFromReader(bytes.NewReader([]byte(fixed)), false)
	if err != nil {
		return err
	}
	b.BmapFileSHA1 = defaultBMAPHash
	return b.Write(input)

}