/*
Package litprog implements noweb's noweave and notangle functionality.
*/
package main

import (
    "bufio"
    "bytes"
    "flag"
    "fmt"
    "os"
    "path"
    "regexp"
    "sort"
    "strconv"
    "strings"
    "time"
)

const (
    name    = "neweb"
    version = "0.3.1"

    // split into single string to enable self-tangling
    openMark  = "^<" + "<"
    closeMark = ">" + ">="
    end       = "^@"
)

// CodeChunk ...
//
// number is simply the incremented number of code chunks in the document.
// next and prev point to chunk continuations.
type CodeChunk struct {
    number      int
    name        string
    lang        string
    lines       []string
    definitions []string
    next        int
    prev        int
}

// factory function (no method!)
func NewCodeChunk() *CodeChunk {
    return &CodeChunk{}
}

func (c *CodeChunk) Append(line string) {
    c.lines = append(c.lines, line)
}

// DocChunk ...
type DocChunk struct {
    lines []string
}

func NewDocChunk() *DocChunk {
    return &DocChunk{}
}

func (c *DocChunk) Append(line string) {
    c.lines = append(c.lines, line)
}

// ChunkPair ...
type ChunkPair struct {
    docChunk  *DocChunk
    codeChunk *CodeChunk
}

func NewChunkPair(doc *DocChunk, code *CodeChunk) ChunkPair {
    return ChunkPair{doc, code}
}

// Chunkifier ...
type Chunkifier struct {
    chunks []ChunkPair
}

// Konstruktoren sind keine Methoden!
func NewChunkifier() *Chunkifier {
    return &Chunkifier{}
}

// Tokenize parses lines for code or doc chunk definitions.
//
// Returns a list of chunk pairs, codeChunkSequenceMap and codeChunkUseMap.
func (c *Chunkifier) Tokenize() ([]ChunkPair, *map[string][]int, *map[string][]int) {
    codeChunkDefinition := openMark + `([^>]+)` + closeMark
    codeChunkDefinitionExtended := openMark + `([^>]+)` + closeMark + `\s*\((\w+)\)`
    docChunkDefinition := end + `\s*(%def [\w ]+)*\s*.*`

    // create an empty list of chunk pairs
    chunkPairs := make([]ChunkPair, 0)

    // use map
    // "name" is used like < <...> > in chunks [1, 3, 6]
    codeChunkUseMap := make(map[string][]int, 0)

    // sequence/continuations
    // 1 <-> 5 <-> 6
    codeChunkSequenceMap := make(map[string][]int, 0)

    // MustCompile wirft eine panic, wenn der reg. Ausdruck fehlerhaft ist
    reCodeChunkDefinition := regexp.MustCompile(codeChunkDefinition)
    reCodeChunkDefinitionExtended := regexp.MustCompile(codeChunkDefinitionExtended)
    reDocChunkDefinition := regexp.MustCompile(docChunkDefinition)
    reCodeChunkUse := regexp.MustCompile(`<` + `<([^>]+)>` + `>`)

    // prepare our first chunk pair
    currentDocChunk := NewDocChunk()
    currentCodeChunk := NewCodeChunk()
    currentChunkPair := NewChunkPair(currentDocChunk, currentCodeChunk)

    chunkCounter := 0

    // we always start with a doc chunk
    inDocChunk := true

    // used to continue earlier
    continue_immediately := false

    scanner := bufio.NewScanner(os.Stdin)
    for scanner.Scan() {
        line := scanner.Text()

        // check for new code chunk with extended syntax
        match := reCodeChunkDefinitionExtended.FindStringSubmatch(line)
        if match != nil {
            currentCodeChunk.name = match[1]
            currentCodeChunk.lang = match[2]
            inDocChunk = false
        } else {
            // check for old syntax extension: < <code chunk name> >=
            match := reCodeChunkDefinition.FindStringSubmatch(line)

            if match != nil {
                currentCodeChunk.name = match[1]
                inDocChunk = false
            } else {
                // check for doc chunk beginning
                match := reDocChunkDefinition.FindStringSubmatch(line)

                // found doc chunk
                // check for additional information
                if match != nil {
                    if len(match[1]) != 0 {
                        // [1:]: strip of %def
                        currentCodeChunk.definitions = strings.Split(match[1], " ")[1:]
                        // to straigt to next line
                        continue_immediately = true
                    } else {
                        // this line is part of the
                        // new doc string; strip of `@`
                        line = line[1:]
                    }

                    // save pair
                    // and create a new one
                    currentCodeChunk.number = chunkCounter
                    chunkPairs = append(chunkPairs, currentChunkPair)

                    // add number to sequence map
                    chunkIndices, ok := codeChunkSequenceMap[currentCodeChunk.name]
                    if ok != true {
                        // ok, new chunk; we need a
                        // new entry
                        codeChunkSequenceMap[currentCodeChunk.name] = make([]int, 0)
                        codeChunkSequenceMap[currentCodeChunk.name] = append(codeChunkSequenceMap[currentCodeChunk.name], chunkCounter)

                    } else {
                        codeChunkSequenceMap[currentCodeChunk.name] = append(chunkIndices, chunkCounter)
                    }
                    chunkCounter += 1

                    currentDocChunk = NewDocChunk()
                    currentCodeChunk = NewCodeChunk()
                    currentChunkPair = NewChunkPair(currentDocChunk, currentCodeChunk)

                    inDocChunk = true
                }
                if continue_immediately {
                    continue_immediately = false
                    continue
                }

                // are we in a doc chunk or a code chunk?
                if inDocChunk {
                    // save the line
                    currentDocChunk.Append(line)

                } else {
                    // save the line...
                    currentCodeChunk.Append(line)

                    // ...and search for code chunk uses
                    for {
                        match := reCodeChunkUse.FindStringIndex(line)
                        var chunkName, after string

                        if match != nil {
                            // match[0]-match[1]: matched code chunk use
                            // match[1+1]: part after
                            after = line[match[1]:]
                            // split off < < and > >!
                            chunkName = line[match[0]+2 : match[1]-2]

                            _, ok := codeChunkUseMap[chunkName]
                            if ok != true {
                                codeChunkUseMap[chunkName] = make([]int, 0)
                            }
                            codeChunkUseMap[chunkName] = append(codeChunkUseMap[chunkName], chunkCounter)
                            line = after
                        } else {
                            break
                        }
                    }

                }
            }
        }
    }

    // add last pair to list
    chunkPairs = append(chunkPairs, currentChunkPair)

    if err := scanner.Err(); err != nil {
        fmt.Println(os.Stderr, "reading standard input:", err)
    }

    return chunkPairs, &codeChunkSequenceMap, &codeChunkUseMap

}

// stamp is needed because we want unique code chunk indentifiers across
// multiple files.  A nano timestamp should do it.  It must be added to every id
// and link and fenced code id write below.
type MarkdownFormatter struct {
    writer *bufio.Writer
    stamp  int64
}

// factory function (no method!)
func NewMarkdownFormatter() *MarkdownFormatter {
    t := time.Now()
    return &MarkdownFormatter{bufio.NewWriter(os.Stdout), t.UnixNano()}
}

// normalizeLabel replaces all non-link chars with link-acceptable chars.
func (f *MarkdownFormatter) normalizeLabel(line string) string {
    label := strings.Replace(line, " ", "_", -1)
    label = strings.Replace(label, "/", "_", -1)
    label = strings.Replace(label, ".", "_", -1)

    return label
}

// replacement function for ReplaceAllFunc
func (f *MarkdownFormatter) replaceBrackets(match []byte) []byte {
    name := bytes.Replace(match, []byte("[["), []byte(""), -1)
    name = bytes.Replace(name, []byte("]]"), []byte(""), -1)

    linkTarget := f.normalizeLabel(string(name))

    link := "[" + string(name) + "](#" + string(f.stamp) + linkTarget + ")"

    return []byte(link)

}

func (f *MarkdownFormatter) writeCodeChunkIndex(sortedChunkNames []string, sequenceMap *map[string][]int) {
    for _, key := range sortedChunkNames {
        f.writer.WriteString("`" + key + "`\n  ~ ")

        var line string
        for _, number := range (*sequenceMap)[key] {
            line = fmt.Sprintf("<a href=\"#code:%d:%d\">`%d`</a> ", f.stamp, number, number)
            f.writer.WriteString(line)
        }
        f.writer.WriteString("\n\n")
    }
}

func (f *MarkdownFormatter) writeDocChunk(chunk *DocChunk,
    sortedChunkNames []string, sequenceMap *map[string][]int) {
    // inline [[references]]
    reBracketRef := regexp.MustCompile(`\[\[[^>\]]+\]\]`)
    reCodeChunkIndex := regexp.MustCompile(`^\s*{{CODE CHUNK INDEX}}\s*$`)

    for _, w := range chunk.lines {
        match := reBracketRef.FindStringSubmatch(w)
        if match != nil {
            line := reBracketRef.ReplaceAllFunc([]byte(w), f.replaceBrackets)
            f.writer.WriteString(string(line) + "\n")
        } else {
            // check for a line with only {{CODE CHUNK INDEX}}
            match := reCodeChunkIndex.FindStringSubmatch(w)
            if match != nil {
                f.writeCodeChunkIndex(sortedChunkNames, sequenceMap)
            } else {
                f.writer.WriteString(w + "\n")
            }
        }
    }
}

func isIn(value int, array []int) (int, bool) {
    // array is already sorted, so we can use the binary search algorithm
    // from sort.Search
    //
    // i is the index where value in array exists
    i := sort.SearchInts(array, value)
    if i < len(array) && array[i] == value {
        return i, true
    } else {
        return -1, false
    }
}

func (f *MarkdownFormatter) writeCodeChunk(chunk *CodeChunk, sequence []int, useMap *map[string][]int) {
    // is this an empty chunk?
    if len(chunk.lines) != 0 {
        // write fenced code chunk
        //
        // * check if this is the first of a sequence
        // * check for following chunks
        // * check if we are in the middle of a sequence
        //
        pos, ok := isIn(chunk.number, sequence)
        if ok {
            // write used in if available
            formattedChunkIds, ok := f.codeUsedIn(chunk.name, useMap)
            if ok {
                formattedChunkIds = "(" + formattedChunkIds + ")"
            }

            line := fmt.Sprintf("\\noindent\n<a id=\"code:%d:%d\" />`%d ⟨%s⟩≡` %s", f.stamp, chunk.number, chunk.number, chunk.name, formattedChunkIds)
            switch {
            case pos == 0 && len(sequence) == 1:
                // single chunk
                line += "\n\n"
            case pos == 0 && len(sequence) > 1:
                // first position of sequence
                line += fmt.Sprintf(" <a href=\"#code:%d:%d\">`%d ▷`</a>\n\n", f.stamp, sequence[pos+1], sequence[pos+1])
            case pos > 0 && pos < len(sequence)-1:
                // continued chunk
                line += fmt.Sprintf(" <a href=\"#code:%d:%d\">`◁ %d`</a> <a href=\"#code:%d:%d\">`%d ▷`</a>\n\n", f.stamp, sequence[pos-1], sequence[pos-1], f.stamp, sequence[pos+1], sequence[pos+1])
            case pos == len(sequence)-1:
                // last chunk of sequence
                line += fmt.Sprintf(" <a href=\"#code:%d:%d\">`◁ %d`</a>\n\n", f.stamp, sequence[pos-1], sequence[pos-1])
            }
            f.writer.WriteString(line)
        }

        f.writer.WriteString("~~~")
        if len(chunk.lang) != 0 {
            f.writer.WriteString(" {." + chunk.lang)

        } else {
            f.writer.WriteString(" {")
        }
        // add a label
        cleanName := f.normalizeLabel(chunk.name)
        f.writer.WriteString(" #" + cleanName + "}\n")

        //fmt.Println("defs:", v.codeChunk.definitions)

        for _, w := range chunk.lines {
            f.writer.WriteString(w + "\n")
        }

        f.writer.WriteString("~~~\n")
        // write def lines if available
        if len(chunk.definitions) > 0 {
            f.writer.WriteString("<div style=\"font-size: smaller;\">Defines:<br/>&nbsp;\n")
            line := strings.Join(chunk.definitions, ", ")
            f.writer.WriteString(line + ".</div>\n")
        }

        // write used in if available
        formattedChunkIds, ok := f.codeUsedIn(chunk.name, useMap)
        if ok {
            line := "<div style=\"font-size: smaller;\">This code is used in chunk "
            line += formattedChunkIds
            f.writer.WriteString(line + ".</div>\n")
        }
    }
}

func (f *MarkdownFormatter) codeUsedIn(name string, useMap *map[string][]int) (string, bool) {
    // write used in if available
    line := ""
    if len((*useMap)[name]) > 0 {
        uses := make([]string, 0)
        for _, v := range (*useMap)[name] {
            uses = append(uses, fmt.Sprintf("<a href=\"#code:%d:%d\">`%d`</a>", f.stamp, v, v))
        }
        line += strings.Join(uses, ", ")

        return line, true
    } else {
        return line, false
    }

}

// ints2Strings converts a list of ints to a list of strings, usable for
// strings.Join().
func ints2Strings(ints []int) []string {
    strings := make([]string, 0)
    for _, v := range ints {
        strings = append(strings, strconv.Itoa(v))
    }

    fmt.Println(os.Stderr, strings)
    return strings
}

func (f *MarkdownFormatter) Convert(pairs []ChunkPair, sequenceMap, useMap *map[string][]int) {
    // sort keys for index generation
    var sortedChunkNames []string

    for k, _ := range *sequenceMap {
        sortedChunkNames = append(sortedChunkNames, k)
    }
    sort.Strings(sortedChunkNames)

    for _, v := range pairs {
        f.writeDocChunk(v.docChunk, sortedChunkNames, sequenceMap)
        f.writeCodeChunk(v.codeChunk, (*sequenceMap)[v.codeChunk.name], useMap)
    }

    // do not forget to flush! otherwise the output will be incomplete
    f.writer.Flush()
}

// Tokenize creates pairs of (doc, code) chunks from a literate program on Stdin.
func Tokenize() ([]ChunkPair, *map[string][]int, *map[string][]int) {
    c := NewChunkifier()
    return c.Tokenize()
}

// Weave converts the input on Stdin to pandoc Markdown and writes the result
// to Stdout.
func Weave() {
    formatter := NewMarkdownFormatter()

    pairs, sequenceMap, useMap := Tokenize()
    //fmt.Println(os.Stderr, "weave:", useMap)
    formatter.Convert(pairs, sequenceMap, useMap)
}

// FlattenCodeChunks flattens the string list of a code chunk into a single string.
func flattenCodeChunkLinesToString(pairs []ChunkPair) map[string]string {
    codeChunkMap := make(map[string]string)

    for _, v := range pairs {
        codeChunk := v.codeChunk

        // check if chunk is already in map
        _, ok := codeChunkMap[codeChunk.name]
        if ok != true {
            codeChunkMap[codeChunk.name] = ""
        }

        // Join(list_of_strings, sep) -> string
        // and add the \n
        for _, line := range codeChunk.lines {
            codeChunkMap[codeChunk.name] = codeChunkMap[codeChunk.name] + line + "\n"
        }

        // strip off last newline, but check for empty chunks
        length := len(codeChunkMap[codeChunk.name])
        if length > 0 {
            codeChunkMap[codeChunk.name] = codeChunkMap[codeChunk.name][:length-1]
        }
    }

    return codeChunkMap
}

// Expand replaces chunk code uses.
func expand(startChunkName string, codeChunkMap map[string]string, indent string) string {
    reCodeChunkUse := regexp.MustCompile(`<` + `<([^>]+)>` + `>`)
    chunkName := ""

    expandedLine := ""

    // first check if we have a chunk with the given name
    chunkLines, ok := codeChunkMap[startChunkName]
    if ok != true {
        fmt.Println(os.Stderr, "Cannot expand unknown chunk name: ", chunkName)
        os.Exit(1)
    } else {

        // loop as long as there are expandable code chunk uses
        // in chunkLines
        // already expanded chunks and non-chunk parts are in
        // expandedLine
        //
        // use a flag to indicate that there might be more to expand:
        finished := 2
        for {
            
            fmt.Println(os.Stderr, "loop starts with finished: ", finished)
            if finished == 0 {
                break
            }
            match := reCodeChunkUse.FindStringIndex(chunkLines)
            var after string

            if match != nil {
                // something found → set flag
                finished = 1

                // 0-match[0-1]: part before
                // match[0]-match[1]: matched code chunk use
                // match[1+1]: part after
                //
                // the first part is not expandable, so store
                // it and omit a new unnecessary findstring
                expandedLine += chunkLines[0:match[0]]

                after = chunkLines[match[1]:]
                // split off < < and > >!
                chunkName = chunkLines[match[0]+2 : match[1]-2]

                // check for escaped chunk: @<<
                if match[0] > 0 && string(chunkLines[match[0]-1]) == "@" {
                    // add the escaped chunk name to
                    // result without the @
                    expandedLine = expandedLine[0:len(expandedLine)-1] + chunkLines[match[0]:match[1]]
                } else {
                    // try to replace the code chunk
                    _, ok := codeChunkMap[chunkName]
                    if ok != true {
                        fmt.Println(os.Stderr, "Cannot find code chunk named", chunkName)
                        os.Exit(1)
                    }
                    //chunkLines = before + codeChunkMap[chunkName] + after
                    expandedLine = expandedLine + codeChunkMap[chunkName]
                }
                // the after part may contain more expandable
                // chunks
                chunkLines = after
            } else {
                // do not forget to add the unexandable rest
                // to the result
                expandedLine += chunkLines

                // check finished status
                if finished == 2 {
                    // nothing found to expand (lines with
                    // no code chunks): exit
                    break
                } else {
                    // something was expaned, but there
                    // may be more: stay in for loop
                    chunkLines = expandedLine
                    finished = 2
                }
            }
        }

    }

    return expandedLine
}

// Tangle extracts the sourcecode from the literate program on Stdin, starting
// with start_chunk.
func Tangle(startChunk string) {
    pairs, _, _ := Tokenize()

    codeChunkMap := flattenCodeChunkLinesToString(pairs)
    lines := expand(startChunk, codeChunkMap, "")
    writer := bufio.NewWriter(os.Stdout)

    writer.WriteString(lines)
    writer.Flush()
}

func main() {
    // add -v: version
    versionFlag := flag.Bool("v", false, "print version")

    // dispatch tangle and weave
    switch prg := path.Base(os.Args[0]); prg {
    case "neweave":
        flag.Usage = func() {
            fmt.Printf("Usage of %s:\n", os.Args[0])
            fmt.Printf("  %s < x.nw > x.md\n", os.Args[0])
        }
        flag.Parse()

        if *versionFlag {
            fmt.Printf("%s %s\n", name, version)
            os.Exit(0)
        } else {

            Weave()
        }
    case "newtangle":
        var tangle_start_chunk string
        flag.StringVar(&tangle_start_chunk, "R", "*", "chunk name to start with")
        flag.Parse()

        if *versionFlag {
            fmt.Printf("%s %s\n", name, version)
            os.Exit(0)
        } else {
            Tangle(tangle_start_chunk)
        }
    default:
        fmt.Fprintf(os.Stderr, "unknown program name\n")
        os.Exit(-2)
    }
}

$ pandoc -o neweb.pdf --filter pandoc-citeproc --latex-engine=xelatex neweb.nw

./
nw
pandoc -o neweb.pdf --filter pandoc-citeproc --latex-engine=xelatex neweb.nw