<commit_before>package quadtree

import (
	"math/rand"
	"reflect"
	"sort"
	"testing"

	"github.com/paulmach/orb"
	"github.com/paulmach/orb/planar"
)

func TestNew(t *testing.T) {
	bound := orb.Bound{Min: orb.Point{0, 2}, Max: orb.Point{1, 3}}
	qt := New(bound)

	if !qt.Bound().Equal(bound) {
		t.Errorf("should use provided bound, got %v", qt.Bound())
	}
}

func TestQuadtreeAdd(t *testing.T) {
	p := orb.Point{}
	qt := New(orb.Bound{Min: orb.Point{0, 0}, Max: orb.Point{1, 1}})
	for i := 0; i < 10; i++ {
		// should be able to insert the same point over and over.
		qt.Add(p)
	}
}

func TestQuadtreeRemove(t *testing.T) {
	r := rand.New(rand.NewSource(42))

	qt := New(orb.Bound{Min: orb.Point{0, 0}, Max: orb.Point{1, 1}})
	mp := orb.MultiPoint{}
	for i := 0; i < 1000; i++ {
		mp = append(mp, orb.Point{r.Float64(), r.Float64()})
		qt.Add(mp[i])
	}

	for i := 0; i < 1000; i += 3 {
		qt.Remove(mp[i], nil)
		mp[i] = orb.Point{-10000, -10000}
	}

	// make sure finding still works for 1000 random points
	for i := 0; i < 1000; i++ {
		p := orb.Point{r.Float64(), r.Float64()}

		f := qt.Find(p)
		_, j := planar.DistanceFromWithIndex(mp, p)

		if e := mp[j]; !e.Equal(f.Point()) {
			t.Errorf("index: %d, unexpected point %v != %v", i, e, f.Point())
		}
	}
}

func TestQuadtreeFind(t *testing.T) {
	points := orb.MultiPoint{}
	dim := 17

	for i := 0; i < dim*dim; i++ {
		points = append(points, orb.Point{float64(i % dim), float64(i / dim)})
	}

	qt := New(points.Bound())
	for _, p := range points {
		qt.Add(p)
	}

	cases := []struct {
		point    orb.Point
		expected orb.Point
	}{
		{point: orb.Point{0.1, 0.1}, expected: orb.Point{0, 0}},
		{point: orb.Point{3.1, 2.9}, expected: orb.Point{3, 3}},
		{point: orb.Point{7.1, 7.1}, expected: orb.Point{7, 7}},
		{point: orb.Point{0.1, 15.9}, expected: orb.Point{0, 16}},
		{point: orb.Point{15.9, 15.9}, expected: orb.Point{16, 16}},
	}

	for i, tc := range cases {
		if v := qt.Find(tc.point); !v.Point().Equal(tc.expected) {
			t.Errorf("incorrect point on %d, got %v", i, v)
		}
	}
}

func TestQuadtreeFind_Random(t *testing.T) {
	r := rand.New(rand.NewSource(42))

	qt := New(orb.Bound{Min: orb.Point{0, 0}, Max: orb.Point{1, 1}})
	mp := orb.MultiPoint{}
	for i := 0; i < 1000; i++ {
		mp = append(mp, orb.Point{r.Float64(), r.Float64()})
		qt.Add(mp[i])
	}

	for i := 0; i < 1000; i++ {
		p := orb.Point{r.Float64(), r.Float64()}

		f := qt.Find(p)
		_, j := planar.DistanceFromWithIndex(mp, p)

		if e := mp[j]; !e.Equal(f.Point()) {
			t.Errorf("index: %d, unexpected point %v != %v", i, e, f.Point())
		}
	}
}

func TestQuadtreeMatching(t *testing.T) {
	type dataPointer struct {
		orb.Pointer
		visible bool
	}

	qt := New(orb.Bound{Min: orb.Point{0, 0}, Max: orb.Point{1, 1}})
	qt.Add(dataPointer{orb.Point{0, 0}, false})
	qt.Add(dataPointer{orb.Point{1, 1}, true})

	cases := []struct {
		name     string
		filter   FilterFunc
		point    orb.Point
		expected orb.Point
	}{
		{
			name:     "no filtred",
			point:    orb.Point{0.1, 0.1},
			expected: orb.Point{0, 0},
		},
		{
			name:     "with filter",
			filter:   func(p orb.Pointer) bool { return p.(dataPointer).visible },
			point:    orb.Point{0.1, 0.1},
			expected: orb.Point{1, 1},
		},
		{
			name:     "match none filter",
			filter:   func(p orb.Pointer) bool { return false },
			point:    orb.Point{0.1, 0.1},
			expected: orb.Point{},
		},
	}

	for _, tc := range cases {
		t.Run(tc.name, func(t *testing.T) {
			v := qt.Matching(tc.point, tc.filter)
			if !v.Point().Equal(tc.expected) {
				t.Errorf("incorrect point %v != %v", v, tc.expected)
			}
		})
	}
}

func TestQuadtreeKNearest(t *testing.T) {
	type dataPointer struct {
		orb.Pointer
		visible bool
	}

	q := New(orb.Bound{Max: orb.Point{5, 5}})
	q.Add(dataPointer{orb.Point{0, 0}, false})
	q.Add(dataPointer{orb.Point{1, 1}, true})
	q.Add(dataPointer{orb.Point{2, 2}, false})
	q.Add(dataPointer{orb.Point{3, 3}, true})
	q.Add(dataPointer{orb.Point{4, 4}, false})
	q.Add(dataPointer{orb.Point{5, 5}, true})

	filters := map[bool]FilterFunc{
		false: nil,
		true:  func(p orb.Pointer) bool { return p.(dataPointer).visible },
	}

	cases := []struct {
		name     string
		filtered bool
		point    orb.Point
		expected []orb.Point
	}{
		{
			name:     "unfiltered",
			filtered: false,
			point:    orb.Point{0.1, 0.1},
			expected: []orb.Point{{0, 0}, {1, 1}},
		},
		{
			name:     "filtered",
			filtered: true,
			point:    orb.Point{0.1, 0.1},
			expected: []orb.Point{{1, 1}, {3, 3}},
		},
	}

	for _, tc := range cases {
		t.Run(tc.name, func(t *testing.T) {
			if !tc.filtered {
				v := q.KNearest(nil, tc.point, 2)
				if len(v) != len(tc.expected) {
					t.Errorf("incorrect response length: %d != %d", len(v), len(tc.expected))
				}
			}

			v := q.KNearestMatching(nil, tc.point, 2, filters[tc.filtered])
			if len(v) != len(tc.expected) {
				t.Errorf("incorrect response length: %d != %d", len(v), len(tc.expected))
			}

			result := make([]orb.Point, 0)
			for _, p := range v {
				result = append(result, p.Point())
			}

			sort.Slice(result, func(i, j int) bool {
				return result[i][0] < result[j][0]
			})

			sort.Slice(tc.expected, func(i, j int) bool {
				return tc.expected[i][0] < tc.expected[j][0]
			})

			if !reflect.DeepEqual(result, tc.expected) {
				t.Log(result)
				t.Log(tc.expected)
				t.Errorf("incorrect results")
			}
		})
	}
}

func TestQuadtreeKNearest_DistanceLimit(t *testing.T) {
	type dataPointer struct {
		orb.Pointer
		visible bool
	}

	q := New(orb.Bound{Max: orb.Point{5, 5}})
	q.Add(dataPointer{orb.Point{0, 0}, false})
	q.Add(dataPointer{orb.Point{1, 1}, true})
	q.Add(dataPointer{orb.Point{2, 2}, false})
	q.Add(dataPointer{orb.Point{3, 3}, true})
	q.Add(dataPointer{orb.Point{4, 4}, false})
	q.Add(dataPointer{orb.Point{5, 5}, true})

	filters := map[bool]FilterFunc{
		false: nil,
		true:  func(p orb.Pointer) bool { return p.(dataPointer).visible },
	}

	cases := []struct {
		name     string
		filtered bool
		distance float64
		point    orb.Point
		expected []orb.Point
	}{
		{
			name:     "filtered",
			filtered: true,
			distance: 5,
			point:    orb.Point{0.1, 0.1},
			expected: []orb.Point{{1, 1}, {3, 3}},
		},
		{
			name:     "unfiltered",
			filtered: false,
			distance: 1,
			point:    orb.Point{0.1, 0.1},
			expected: []orb.Point{{0, 0}},
		},
	}

	var v []orb.Pointer
	for _, tc := range cases {
		t.Run(tc.name, func(t *testing.T) {
			v = q.KNearestMatching(v, tc.point, 5, filters[tc.filtered], tc.distance)
			if len(v) != len(tc.expected) {
				t.Errorf("incorrect response length: %d != %d", len(v), len(tc.expected))
			}

			result := make([]orb.Point, 0)
			for _, p := range v {
				result = append(result, p.Point())
			}

			sort.Slice(result, func(i, j int) bool {
				return result[i][0] < result[j][0]
			})

			sort.Slice(tc.expected, func(i, j int) bool {
				return tc.expected[i][0] < tc.expected[j][0]
			})

			if !reflect.DeepEqual(result, tc.expected) {
				t.Log(result)
				t.Log(tc.expected)
				t.Errorf("incorrect results")
			}
		})
	}
}

func TestQuadtreeInBoundMatching(t *testing.T) {
	type dataPointer struct {
		orb.Pointer
		visible bool
	}

	q := New(orb.Bound{Max: orb.Point{5, 5}})
	q.Add(dataPointer{orb.Point{0, 0}, false})
	q.Add(dataPointer{orb.Point{1, 1}, true})
	q.Add(dataPointer{orb.Point{2, 2}, false})
	q.Add(dataPointer{orb.Point{3, 3}, true})
	q.Add(dataPointer{orb.Point{4, 4}, false})
	q.Add(dataPointer{orb.Point{5, 5}, true})

	filters := map[bool]FilterFunc{
		false: nil,
		true:  func(p orb.Pointer) bool { return p.(dataPointer).visible },
	}

	cases := []struct {
		name     string
		filtered bool
		expected []orb.Point
	}{
		{
			name:     "unfiltered",
			filtered: false,
			expected: []orb.Point{{0, 0}, {1, 1}, {2, 2}},
		},
		{
			name:     "filtered",
			filtered: true,
			expected: []orb.Point{{1, 1}},
		},
	}

	bound := orb.Bound{Min: orb.Point{0, 0}, Max: orb.Point{2, 2}}

	var v []orb.Pointer
	for _, tc := range cases {
		t.Run(tc.name, func(t *testing.T) {
			v = q.InBoundMatching(v, bound, filters[tc.filtered])
			if len(v) != len(tc.expected) {
				t.Errorf("incorrect response length: %d != %d", len(v), len(tc.expected))
			}

			result := make([]orb.Point, 0)
			for _, p := range v {
				result = append(result, p.Point())
			}

			sort.Slice(result, func(i, j int) bool {
				return result[i][0] < result[j][0]
			})

			sort.Slice(tc.expected, func(i, j int) bool {
				return tc.expected[i][0] < tc.expected[j][0]
			})

			if !reflect.DeepEqual(result, tc.expected) {
				t.Log(result)
				t.Log(tc.expected)
				t.Errorf("incorrect results")
			}
		})
	}

}

func TestQuadtreeInBound_Random(t *testing.T) {
	r := rand.New(rand.NewSource(43))

	qt := New(orb.Bound{Min: orb.Point{0, 0}, Max: orb.Point{1, 1}})
	mp := orb.MultiPoint{}
	for i := 0; i < 1000; i++ {
		mp = append(mp, orb.Point{r.Float64(), r.Float64()})
		qt.Add(mp[i])
	}

	for i := 0; i < 1000; i++ {
		p := orb.Point{r.Float64(), r.Float64()}

		b := orb.Bound{Min: p, Max: p}
		b = b.Pad(0.1)
		ps := qt.InBound(nil, b)

		// find the right answer brute force
		var list []orb.Pointer
		for _, p := range mp {
			if b.Contains(p) {
				list = append(list, p)
			}
		}

		if len(list) != len(ps) {
			t.Errorf("index: %d, lengths not equal %v != %v", i, len(list), len(ps))
		}
	}
}
<commit_msg>update TestQuadtreeMatching to new behavior of returning nil<commit_after>package quadtree

import (
	"math/rand"
	"reflect"
	"sort"
	"testing"

	"github.com/paulmach/orb"
	"github.com/paulmach/orb/planar"
)

func TestNew(t *testing.T) {
	bound := orb.Bound{Min: orb.Point{0, 2}, Max: orb.Point{1, 3}}
	qt := New(bound)

	if !qt.Bound().Equal(bound) {
		t.Errorf("should use provided bound, got %v", qt.Bound())
	}
}

func TestQuadtreeAdd(t *testing.T) {
	p := orb.Point{}
	qt := New(orb.Bound{Min: orb.Point{0, 0}, Max: orb.Point{1, 1}})
	for i := 0; i < 10; i++ {
		// should be able to insert the same point over and over.
		qt.Add(p)
	}
}

func TestQuadtreeRemove(t *testing.T) {
	r := rand.New(rand.NewSource(42))

	qt := New(orb.Bound{Min: orb.Point{0, 0}, Max: orb.Point{1, 1}})
	mp := orb.MultiPoint{}
	for i := 0; i < 1000; i++ {
		mp = append(mp, orb.Point{r.Float64(), r.Float64()})
		qt.Add(mp[i])
	}

	for i := 0; i < 1000; i += 3 {
		qt.Remove(mp[i], nil)
		mp[i] = orb.Point{-10000, -10000}
	}

	// make sure finding still works for 1000 random points
	for i := 0; i < 1000; i++ {
		p := orb.Point{r.Float64(), r.Float64()}

		f := qt.Find(p)
		_, j := planar.DistanceFromWithIndex(mp, p)

		if e := mp[j]; !e.Equal(f.Point()) {
			t.Errorf("index: %d, unexpected point %v != %v", i, e, f.Point())
		}
	}
}

func TestQuadtreeFind(t *testing.T) {
	points := orb.MultiPoint{}
	dim := 17

	for i := 0; i < dim*dim; i++ {
		points = append(points, orb.Point{float64(i % dim), float64(i / dim)})
	}

	qt := New(points.Bound())
	for _, p := range points {
		qt.Add(p)
	}

	cases := []struct {
		point    orb.Point
		expected orb.Point
	}{
		{point: orb.Point{0.1, 0.1}, expected: orb.Point{0, 0}},
		{point: orb.Point{3.1, 2.9}, expected: orb.Point{3, 3}},
		{point: orb.Point{7.1, 7.1}, expected: orb.Point{7, 7}},
		{point: orb.Point{0.1, 15.9}, expected: orb.Point{0, 16}},
		{point: orb.Point{15.9, 15.9}, expected: orb.Point{16, 16}},
	}

	for i, tc := range cases {
		if v := qt.Find(tc.point); !v.Point().Equal(tc.expected) {
			t.Errorf("incorrect point on %d, got %v", i, v)
		}
	}
}

func TestQuadtreeFind_Random(t *testing.T) {
	r := rand.New(rand.NewSource(42))

	qt := New(orb.Bound{Min: orb.Point{0, 0}, Max: orb.Point{1, 1}})
	mp := orb.MultiPoint{}
	for i := 0; i < 1000; i++ {
		mp = append(mp, orb.Point{r.Float64(), r.Float64()})
		qt.Add(mp[i])
	}

	for i := 0; i < 1000; i++ {
		p := orb.Point{r.Float64(), r.Float64()}

		f := qt.Find(p)
		_, j := planar.DistanceFromWithIndex(mp, p)

		if e := mp[j]; !e.Equal(f.Point()) {
			t.Errorf("index: %d, unexpected point %v != %v", i, e, f.Point())
		}
	}
}

func TestQuadtreeMatching(t *testing.T) {
	type dataPointer struct {
		orb.Pointer
		visible bool
	}

	qt := New(orb.Bound{Min: orb.Point{0, 0}, Max: orb.Point{1, 1}})
	qt.Add(dataPointer{orb.Point{0, 0}, false})
	qt.Add(dataPointer{orb.Point{1, 1}, true})

	cases := []struct {
		name     string
		filter   FilterFunc
		point    orb.Point
		expected orb.Pointer
	}{
		{
			name:     "no filtred",
			point:    orb.Point{0.1, 0.1},
			expected: orb.Point{0, 0},
		},
		{
			name:     "with filter",
			filter:   func(p orb.Pointer) bool { return p.(dataPointer).visible },
			point:    orb.Point{0.1, 0.1},
			expected: orb.Point{1, 1},
		},
		{
			name:     "match none filter",
			filter:   func(p orb.Pointer) bool { return false },
			point:    orb.Point{0.1, 0.1},
			expected: nil,
		},
	}

	for _, tc := range cases {
		t.Run(tc.name, func(t *testing.T) {
			v := qt.Matching(tc.point, tc.filter)

			// case 1: exact match, important for testing `nil`
			if v == tc.expected {
				return
			}

			// case 2: match on returned orb.Point value
			if !v.Point().Equal(tc.expected.Point()) {
				t.Errorf("incorrect point %v != %v", v, tc.expected)
			}
		})
	}
}

func TestQuadtreeKNearest(t *testing.T) {
	type dataPointer struct {
		orb.Pointer
		visible bool
	}

	q := New(orb.Bound{Max: orb.Point{5, 5}})
	q.Add(dataPointer{orb.Point{0, 0}, false})
	q.Add(dataPointer{orb.Point{1, 1}, true})
	q.Add(dataPointer{orb.Point{2, 2}, false})
	q.Add(dataPointer{orb.Point{3, 3}, true})
	q.Add(dataPointer{orb.Point{4, 4}, false})
	q.Add(dataPointer{orb.Point{5, 5}, true})

	filters := map[bool]FilterFunc{
		false: nil,
		true:  func(p orb.Pointer) bool { return p.(dataPointer).visible },
	}

	cases := []struct {
		name     string
		filtered bool
		point    orb.Point
		expected []orb.Point
	}{
		{
			name:     "unfiltered",
			filtered: false,
			point:    orb.Point{0.1, 0.1},
			expected: []orb.Point{{0, 0}, {1, 1}},
		},
		{
			name:     "filtered",
			filtered: true,
			point:    orb.Point{0.1, 0.1},
			expected: []orb.Point{{1, 1}, {3, 3}},
		},
	}

	for _, tc := range cases {
		t.Run(tc.name, func(t *testing.T) {
			if !tc.filtered {
				v := q.KNearest(nil, tc.point, 2)
				if len(v) != len(tc.expected) {
					t.Errorf("incorrect response length: %d != %d", len(v), len(tc.expected))
				}
			}

			v := q.KNearestMatching(nil, tc.point, 2, filters[tc.filtered])
			if len(v) != len(tc.expected) {
				t.Errorf("incorrect response length: %d != %d", len(v), len(tc.expected))
			}

			result := make([]orb.Point, 0)
			for _, p := range v {
				result = append(result, p.Point())
			}

			sort.Slice(result, func(i, j int) bool {
				return result[i][0] < result[j][0]
			})

			sort.Slice(tc.expected, func(i, j int) bool {
				return tc.expected[i][0] < tc.expected[j][0]
			})

			if !reflect.DeepEqual(result, tc.expected) {
				t.Log(result)
				t.Log(tc.expected)
				t.Errorf("incorrect results")
			}
		})
	}
}

func TestQuadtreeKNearest_DistanceLimit(t *testing.T) {
	type dataPointer struct {
		orb.Pointer
		visible bool
	}

	q := New(orb.Bound{Max: orb.Point{5, 5}})
	q.Add(dataPointer{orb.Point{0, 0}, false})
	q.Add(dataPointer{orb.Point{1, 1}, true})
	q.Add(dataPointer{orb.Point{2, 2}, false})
	q.Add(dataPointer{orb.Point{3, 3}, true})
	q.Add(dataPointer{orb.Point{4, 4}, false})
	q.Add(dataPointer{orb.Point{5, 5}, true})

	filters := map[bool]FilterFunc{
		false: nil,
		true:  func(p orb.Pointer) bool { return p.(dataPointer).visible },
	}

	cases := []struct {
		name     string
		filtered bool
		distance float64
		point    orb.Point
		expected []orb.Point
	}{
		{
			name:     "filtered",
			filtered: true,
			distance: 5,
			point:    orb.Point{0.1, 0.1},
			expected: []orb.Point{{1, 1}, {3, 3}},
		},
		{
			name:     "unfiltered",
			filtered: false,
			distance: 1,
			point:    orb.Point{0.1, 0.1},
			expected: []orb.Point{{0, 0}},
		},
	}

	var v []orb.Pointer
	for _, tc := range cases {
		t.Run(tc.name, func(t *testing.T) {
			v = q.KNearestMatching(v, tc.point, 5, filters[tc.filtered], tc.distance)
			if len(v) != len(tc.expected) {
				t.Errorf("incorrect response length: %d != %d", len(v), len(tc.expected))
			}

			result := make([]orb.Point, 0)
			for _, p := range v {
				result = append(result, p.Point())
			}

			sort.Slice(result, func(i, j int) bool {
				return result[i][0] < result[j][0]
			})

			sort.Slice(tc.expected, func(i, j int) bool {
				return tc.expected[i][0] < tc.expected[j][0]
			})

			if !reflect.DeepEqual(result, tc.expected) {
				t.Log(result)
				t.Log(tc.expected)
				t.Errorf("incorrect results")
			}
		})
	}
}

func TestQuadtreeInBoundMatching(t *testing.T) {
	type dataPointer struct {
		orb.Pointer
		visible bool
	}

	q := New(orb.Bound{Max: orb.Point{5, 5}})
	q.Add(dataPointer{orb.Point{0, 0}, false})
	q.Add(dataPointer{orb.Point{1, 1}, true})
	q.Add(dataPointer{orb.Point{2, 2}, false})
	q.Add(dataPointer{orb.Point{3, 3}, true})
	q.Add(dataPointer{orb.Point{4, 4}, false})
	q.Add(dataPointer{orb.Point{5, 5}, true})

	filters := map[bool]FilterFunc{
		false: nil,
		true:  func(p orb.Pointer) bool { return p.(dataPointer).visible },
	}

	cases := []struct {
		name     string
		filtered bool
		expected []orb.Point
	}{
		{
			name:     "unfiltered",
			filtered: false,
			expected: []orb.Point{{0, 0}, {1, 1}, {2, 2}},
		},
		{
			name:     "filtered",
			filtered: true,
			expected: []orb.Point{{1, 1}},
		},
	}

	bound := orb.Bound{Min: orb.Point{0, 0}, Max: orb.Point{2, 2}}

	var v []orb.Pointer
	for _, tc := range cases {
		t.Run(tc.name, func(t *testing.T) {
			v = q.InBoundMatching(v, bound, filters[tc.filtered])
			if len(v) != len(tc.expected) {
				t.Errorf("incorrect response length: %d != %d", len(v), len(tc.expected))
			}

			result := make([]orb.Point, 0)
			for _, p := range v {
				result = append(result, p.Point())
			}

			sort.Slice(result, func(i, j int) bool {
				return result[i][0] < result[j][0]
			})

			sort.Slice(tc.expected, func(i, j int) bool {
				return tc.expected[i][0] < tc.expected[j][0]
			})

			if !reflect.DeepEqual(result, tc.expected) {
				t.Log(result)
				t.Log(tc.expected)
				t.Errorf("incorrect results")
			}
		})
	}

}

func TestQuadtreeInBound_Random(t *testing.T) {
	r := rand.New(rand.NewSource(43))

	qt := New(orb.Bound{Min: orb.Point{0, 0}, Max: orb.Point{1, 1}})
	mp := orb.MultiPoint{}
	for i := 0; i < 1000; i++ {
		mp = append(mp, orb.Point{r.Float64(), r.Float64()})
		qt.Add(mp[i])
	}

	for i := 0; i < 1000; i++ {
		p := orb.Point{r.Float64(), r.Float64()}

		b := orb.Bound{Min: p, Max: p}
		b = b.Pad(0.1)
		ps := qt.InBound(nil, b)

		// find the right answer brute force
		var list []orb.Pointer
		for _, p := range mp {
			if b.Contains(p) {
				list = append(list, p)
			}
		}

		if len(list) != len(ps) {
			t.Errorf("index: %d, lengths not equal %v != %v", i, len(list), len(ps))
		}
	}
}
<|endoftext|>


<commit_before>package db

/*
Stores the database functions related to tasks like
GetTaskByID(id int)
GetTasks(status string)
DeleteAll()
*/

import (
	"database/sql"
	"log"
	"strings"
	"time"

	_ "github.com/mattn/go-sqlite3" //we want to use sqlite natively
	md "github.com/shurcooL/github_flavored_markdown"
	"github.com/thewhitetulip/Tasks/types"
)

var database Database
var taskStatus map[string]int
var err error

//Database encapsulates database
type Database struct {
	db *sql.DB
}

//Begins a transaction
func (db Database) begin() (tx *sql.Tx) {
	tx, err := db.db.Begin()
	if err != nil {
		log.Println(err)
		return nil
	}
	return tx
}

func (db Database) prepare(q string) (stmt *sql.Stmt) {
	stmt, err := db.db.Prepare(q)
	if err != nil {
		log.Println(err)
		return nil
	}
	return stmt
}

func (db Database) query(q string, args ...interface{}) (rows *sql.Rows) {
	rows, err := db.db.Query(q, args...)
	if err != nil {
		log.Println(err)
		return nil
	}
	return rows
}

func init() {
	database.db, err = sql.Open("sqlite3", "./tasks.db")
	taskStatus = map[string]int{"COMPLETE": 1, "PENDING": 2, "DELETED": 3}
	if err != nil {
		log.Fatal(err)
	}
}

//Close function closes this database connection
func Close() {
	database.db.Close()
}

//GetTasks retrieves all the tasks depending on the
//status pending or trashed or completed
func GetTasks(username, status, category string) (types.Context, error) {
	log.Println("getting tasks for ", status)
	var tasks []types.Task
	var task types.Task
	var TaskCreated time.Time
	var context types.Context
	var getTaskSQL string
	var rows *sql.Rows

	comments, err := GetComments(username)

	if err != nil {
		return context, err
	}

	basicSQL := "select t.id, title, content, created_date, priority, c.name from task t, category c, status s, user u where u.username=? and s.id=t.task_status_id and c.id=t.cat_id and u.id=t.user_id"
	if category == "" {
		switch status {
		case "pending":
			getTaskSQL = basicSQL + " and s.status='PENDING'"
		case "deleted":
			getTaskSQL = basicSQL + " and s.status='DELETED' "
		case "completed":
			getTaskSQL = basicSQL + " and s.status='COMPLETE'"
		}

		basicSQL += " order by priority desc, created_date asc"
		rows = database.query(getTaskSQL, username)
	} else {
		status = category
		getTaskSQL = basicSQL + " and name = ?  and  s.status!='DELETED'  order by priority desc, created_date asc, finish_date asc"
		rows, err = database.db.Query(getTaskSQL, username, category)
		log.Print(getTaskSQL)

		if err != nil {
			log.Println("tasks.go: something went wrong while getting query fetch tasks by category")
		}
	}

	defer rows.Close()
	for rows.Next() {
		task = types.Task{}

		err = rows.Scan(&task.Id, &task.Title, &task.Content, &TaskCreated, &task.Priority, &task.Category)

		task.Content = string(md.Markdown([]byte(task.Content)))
		// TaskContent = strings.Replace(TaskContent, "\n", "
", -1)
		if err != nil {
			log.Println(err)
		}

		if comments[task.Id] != nil {
			task.Comments = comments[task.Id]
		}

		TaskCreated = TaskCreated.Local()
		if task.Priority != "1" { // if priority is not 1 then calculate, else why bother?
			CurrentTime := time.Now().Local()
			diff := CurrentTime.Sub(TaskCreated).Hours()
			if diff > 168 {
				task.IsOverdue = true // If one week then overdue by default
			}
		}
		task.Created = TaskCreated.Format("Jan 2 2006")

		tasks = append(tasks, task)
	}
	context = types.Context{Tasks: tasks, Navigation: status}
	return context, nil
}

//GetTaskByID function gets the tasks from the ID passed to the function, used to populate EditTask
func GetTaskByID(username string, id int) (types.Context, error) {
	var tasks []types.Task
	var task types.Task

	getTaskSQL := "select t.id, t.title, t.content, t.priority, c.name from task t join user u left outer join category c  where c.id = t.cat_id and t.id=? and u.username=?"

	rows := database.query(getTaskSQL, id, username)
	defer rows.Close()
	if rows.Next() {
		err := rows.Scan(&task.Id, &task.Title, &task.Content, &task.Priority, &task.Category)
		if err != nil {
			log.Println(err)
			//send email to respective people
		}
	}
	tasks = append(tasks, task)
	context := types.Context{Tasks: tasks, Navigation: "edit"}
	return context, nil
}

//TrashTask is used to delete the task
func TrashTask(username string, id int) error {
	err := taskQuery("update task set task_status_id=?,last_modified_at=datetime() where user_id=(select id from user where username=?) and id=?", taskStatus["DELETED"], username, id)
	return err
}

//CompleteTask  is used to mark tasks as complete
func CompleteTask(username string, id int) error {
	err := taskQuery("update task set task_status_id=?, finish_date=datetime(),last_modified_at=datetime() where id=? and user_id=(select id from user where username=?) ", taskStatus["COMPLETE"], id, username)
	return err
}

//DeleteAll is used to empty the trash
func DeleteAll(username string) error {
	err := taskQuery("delete from task where task_status_id=? where user_id=(select id from user where username=?)", taskStatus["DELETED"], username)
	return err
}

//RestoreTask is used to restore tasks from the Trash
func RestoreTask(username string, id int) error {
	err := taskQuery("update task set task_status_id=?,last_modified_at=datetime(),finish_date=null where id=? and user_id=(select id from user where username=?)", taskStatus["PENDING"], id, username)
	return err
}

//RestoreTaskFromComplete is used to restore tasks from the Trash
func RestoreTaskFromComplete(username string, id int) error {
	err := taskQuery("update task set finish_date=null,last_modified_at=datetime(), task_status_id=? where id=? and user_id=(select id from user where username=?)", taskStatus["PENDING"], id, username)
	return err
}

//DeleteTask is used to delete the task from the database
func DeleteTask(username string, id int) error {
	err := taskQuery("delete from task where id = ? and user_id=(select id from user where username=?)", id, username)
	return err
}

//AddTask is used to add the task in the database
func AddTask(title, content, category string, taskPriority int, username string) error {
	log.Println("AddTask: started function")
	var err error
	userID, err := GetUserID(username)
	if err != nil {
		return err
	}

	if category == "" {
		err = taskQuery("insert into task(title, content, priority, task_status_id, created_date, last_modified_at, user_id) values(?,?,?,?,datetime(), datetime(),?)", title, content, taskPriority, taskStatus["PENDING"], userID)
	} else {
		categoryID := GetCategoryByName(username, category)
		err = taskQuery("insert into task(title, content, priority, created_date, last_modified_at, cat_id, task_status_id, user_id) values(?,?,?,datetime(), datetime(), ?,?,?)", title, content, taskPriority, categoryID, taskStatus["PENDING"], userID)
	}
	return err
}

//GetCategoryIDByName will return the category ID for the category, used in the edit task
//function where we need to be able to update the categoryID of the task
func GetCategoryIDByName(username string, category string) int {
	var categoryID int
	getTaskSQL := "select c.id from category c , user u where u.id = c.user_id and name=? and u.username=?"

	rows := database.query(getTaskSQL, category, username)
	defer rows.Close()
	if rows.Next() {
		err := rows.Scan(&categoryID)
		if err != nil {
			log.Println(err)
			//send email to respective people
		}
	}

	return categoryID
}

//UpdateTask is used to update the tasks in the database
func UpdateTask(id int, title, content, category string, priority int, username string) error {
	categoryID := GetCategoryIDByName(username, category)
	userID, err := GetUserID(username)
	if err != nil {
		return err
	}
	err = taskQuery("update task set title=?, content=?, cat_id=?, priority = ? where id=? and user_id=?", title, content, categoryID, priority, id, userID)
	return err
}

//taskQuery encapsulates running multiple queries which don't do much things
func taskQuery(sql string, args ...interface{}) error {
	log.Print("inside task query")
	SQL := database.prepare(sql)
	tx := database.begin()
	_, err = tx.Stmt(SQL).Exec(args...)
	if err != nil {
		log.Println("taskQuery: ", err)
		tx.Rollback()
	} else {
		err = tx.Commit()
		if err != nil {
			log.Println(err)
			return err
		}
		log.Println("Commit successful")
	}
	return err
}

//SearchTask is used to return the search results depending on the query
func SearchTask(username, query string) (types.Context, error) {
	var tasks []types.Task
	var task types.Task
	var TaskCreated time.Time
	var context types.Context

	comments, err := GetComments(username)
	if err != nil {
		log.Println("SearchTask: something went wrong in finding comments")
	}

	userID, err := GetUserID(username)
	if err != nil {
		return context, err
	}

	stmt := "select t.id, title, content, created_date, priority, c.name from task t, category c where t.user_id=? and c.id = t.cat_id and (title like '%" + query + "%' or content like '%" + query + "%') order by created_date desc"

	rows := database.query(stmt, userID, query, query)
	defer rows.Close()
	for rows.Next() {
		err := rows.Scan(&task.Id, &task.Title, &task.Content, &TaskCreated, &task.Priority, &task.Category)
		if err != nil {
			log.Println(err)
		}

		if comments[task.Id] != nil {
			task.Comments = comments[task.Id]
		}

		task.Title = strings.Replace(task.Title, query, ""+query+"", -1)
		task.Content = strings.Replace(task.Content, query, ""+query+"", -1)
		task.Content = string(md.Markdown([]byte(task.Content)))

		TaskCreated = TaskCreated.Local()
		CurrentTime := time.Now().Local()
		week := TaskCreated.AddDate(0, 0, 7)

		if (week.String() < CurrentTime.String()) && (task.Priority != "1") {
			task.IsOverdue = true // If one week then overdue by default
		}
		task.Created = TaskCreated.Format("Jan 2 2006")

		tasks = append(tasks, task)
	}
	context = types.Context{Tasks: tasks, Search: query, Navigation: "search"}
	return context, nil
}

//GetComments is used to get comments, all of them.
//We do not want 100 different pages to show tasks, we want to use as few pages as possible
//so we are going to populate everything on the damn home pages
func GetComments(username string) (map[int][]types.Comment, error) {
	commentMap := make(map[int][]types.Comment)

	var taskID int
	var comment types.Comment
	var created time.Time

	userID, err := GetUserID(username)
	if err != nil {
		return commentMap, err
	}
	stmt := "select c.id, c.taskID, c.content, c.created from comments c, task t where t.id=c.taskID and c.user_id=?;"
	rows := database.query(stmt, userID)

	defer rows.Close()
	for rows.Next() {
		err := rows.Scan(&comment.ID, &taskID, &comment.Content, &created)
		if err != nil {
			return commentMap, err
		}
		// comment.Content = string(md.Markdown([]byte(comment.Content))) ## have to fix the
 issue markdown support
		created = created.Local()
		comment.Created = created.Format("Jan 2 2006 15:04:05")
		commentMap[taskID] = append(commentMap[taskID], comment)
	}
	return commentMap, nil
}

//AddComments will be used to add comments in the database
func AddComments(username string, id int, comment string) error {
	userID, err := GetUserID(username)
	if err != nil {
		return err
	}
	stmt := "insert into comments(taskID, content, created, user_id) values (?,?,datetime(),?)"
	err = taskQuery(stmt, id, comment, userID)

	if err != nil {
		return err
	}

	log.Println("added comment to task ID ", id)

	return nil
}
<commit_msg>fixed a bug which didn't show correct entires on the home page<commit_after>package db

/*
Stores the database functions related to tasks like
GetTaskByID(id int)
GetTasks(status string)
DeleteAll()
*/

import (
	"database/sql"
	"log"
	"strings"
	"time"

	_ "github.com/mattn/go-sqlite3" //we want to use sqlite natively
	md "github.com/shurcooL/github_flavored_markdown"
	"github.com/thewhitetulip/Tasks/types"
)

var database Database
var taskStatus map[string]int
var err error

//Database encapsulates database
type Database struct {
	db *sql.DB
}

//Begins a transaction
func (db Database) begin() (tx *sql.Tx) {
	tx, err := db.db.Begin()
	if err != nil {
		log.Println(err)
		return nil
	}
	return tx
}

func (db Database) prepare(q string) (stmt *sql.Stmt) {
	stmt, err := db.db.Prepare(q)
	if err != nil {
		log.Println(err)
		return nil
	}
	return stmt
}

func (db Database) query(q string, args ...interface{}) (rows *sql.Rows) {
	rows, err := db.db.Query(q, args...)
	if err != nil {
		log.Println(err)
		return nil
	}
	return rows
}

func init() {
	database.db, err = sql.Open("sqlite3", "./tasks.db")
	taskStatus = map[string]int{"COMPLETE": 1, "PENDING": 2, "DELETED": 3}
	if err != nil {
		log.Fatal(err)
	}
}

//Close function closes this database connection
func Close() {
	database.db.Close()
}

//GetTasks retrieves all the tasks depending on the
//status pending or trashed or completed
func GetTasks(username, status, category string) (types.Context, error) {
	log.Println("getting tasks for ", status)
	var tasks []types.Task
	var task types.Task
	var TaskCreated time.Time
	var context types.Context
	var getTaskSQL string
	var rows *sql.Rows

	comments, err := GetComments(username)

	if err != nil {
		return context, err
	}

	basicSQL := "select t.id, title, content, created_date, priority, c.name from task t, category c, status s, user u where u.username=? and s.id=t.task_status_id and c.id=t.cat_id and u.id=t.user_id"
	if category == "" {
		switch status {
		case "pending":
			getTaskSQL = basicSQL + " and s.status='PENDING'"
		case "deleted":
			getTaskSQL = basicSQL + " and s.status='DELETED' "
		case "completed":
			getTaskSQL = basicSQL + " and s.status='COMPLETE'"
		}

		basicSQL += " order by priority desc, created_date asc"
		rows = database.query(getTaskSQL, username)
	} else {
		status = category
		getTaskSQL = basicSQL + " and name = ?  and  s.status='PENDING'  order by priority desc, created_date asc, finish_date asc"
		rows, err = database.db.Query(getTaskSQL, username, category)
		log.Print(getTaskSQL)

		if err != nil {
			log.Println("tasks.go: something went wrong while getting query fetch tasks by category")
		}
	}

	defer rows.Close()
	for rows.Next() {
		task = types.Task{}

		err = rows.Scan(&task.Id, &task.Title, &task.Content, &TaskCreated, &task.Priority, &task.Category)

		task.Content = string(md.Markdown([]byte(task.Content)))
		// TaskContent = strings.Replace(TaskContent, "\n", "
", -1)
		if err != nil {
			log.Println(err)
		}

		if comments[task.Id] != nil {
			task.Comments = comments[task.Id]
		}

		TaskCreated = TaskCreated.Local()
		if task.Priority != "1" { // if priority is not 1 then calculate, else why bother?
			CurrentTime := time.Now().Local()
			diff := CurrentTime.Sub(TaskCreated).Hours()
			if diff > 168 {
				task.IsOverdue = true // If one week then overdue by default
			}
		}
		task.Created = TaskCreated.Format("Jan 2 2006")

		tasks = append(tasks, task)
	}
	context = types.Context{Tasks: tasks, Navigation: status}
	return context, nil
}

//GetTaskByID function gets the tasks from the ID passed to the function, used to populate EditTask
func GetTaskByID(username string, id int) (types.Context, error) {
	var tasks []types.Task
	var task types.Task

	getTaskSQL := "select t.id, t.title, t.content, t.priority, c.name from task t join user u left outer join category c  where c.id = t.cat_id and t.id=? and u.username=?"

	rows := database.query(getTaskSQL, id, username)
	defer rows.Close()
	if rows.Next() {
		err := rows.Scan(&task.Id, &task.Title, &task.Content, &task.Priority, &task.Category)
		if err != nil {
			log.Println(err)
			//send email to respective people
		}
	}
	tasks = append(tasks, task)
	context := types.Context{Tasks: tasks, Navigation: "edit"}
	return context, nil
}

//TrashTask is used to delete the task
func TrashTask(username string, id int) error {
	err := taskQuery("update task set task_status_id=?,last_modified_at=datetime() where user_id=(select id from user where username=?) and id=?", taskStatus["DELETED"], username, id)
	return err
}

//CompleteTask  is used to mark tasks as complete
func CompleteTask(username string, id int) error {
	err := taskQuery("update task set task_status_id=?, finish_date=datetime(),last_modified_at=datetime() where id=? and user_id=(select id from user where username=?) ", taskStatus["COMPLETE"], id, username)
	return err
}

//DeleteAll is used to empty the trash
func DeleteAll(username string) error {
	err := taskQuery("delete from task where task_status_id=? where user_id=(select id from user where username=?)", taskStatus["DELETED"], username)
	return err
}

//RestoreTask is used to restore tasks from the Trash
func RestoreTask(username string, id int) error {
	err := taskQuery("update task set task_status_id=?,last_modified_at=datetime(),finish_date=null where id=? and user_id=(select id from user where username=?)", taskStatus["PENDING"], id, username)
	return err
}

//RestoreTaskFromComplete is used to restore tasks from the Trash
func RestoreTaskFromComplete(username string, id int) error {
	err := taskQuery("update task set finish_date=null,last_modified_at=datetime(), task_status_id=? where id=? and user_id=(select id from user where username=?)", taskStatus["PENDING"], id, username)
	return err
}

//DeleteTask is used to delete the task from the database
func DeleteTask(username string, id int) error {
	err := taskQuery("delete from task where id = ? and user_id=(select id from user where username=?)", id, username)
	return err
}

//AddTask is used to add the task in the database
func AddTask(title, content, category string, taskPriority int, username string) error {
	log.Println("AddTask: started function")
	var err error
	userID, err := GetUserID(username)
	if err != nil {
		return err
	}

	if category == "" {
		err = taskQuery("insert into task(title, content, priority, task_status_id, created_date, last_modified_at, user_id) values(?,?,?,?,datetime(), datetime(),?)", title, content, taskPriority, taskStatus["PENDING"], userID)
	} else {
		categoryID := GetCategoryByName(username, category)
		err = taskQuery("insert into task(title, content, priority, created_date, last_modified_at, cat_id, task_status_id, user_id) values(?,?,?,datetime(), datetime(), ?,?,?)", title, content, taskPriority, categoryID, taskStatus["PENDING"], userID)
	}
	return err
}

//GetCategoryIDByName will return the category ID for the category, used in the edit task
//function where we need to be able to update the categoryID of the task
func GetCategoryIDByName(username string, category string) int {
	var categoryID int
	getTaskSQL := "select c.id from category c , user u where u.id = c.user_id and name=? and u.username=?"

	rows := database.query(getTaskSQL, category, username)
	defer rows.Close()
	if rows.Next() {
		err := rows.Scan(&categoryID)
		if err != nil {
			log.Println(err)
			//send email to respective people
		}
	}

	return categoryID
}

//UpdateTask is used to update the tasks in the database
func UpdateTask(id int, title, content, category string, priority int, username string) error {
	categoryID := GetCategoryIDByName(username, category)
	userID, err := GetUserID(username)
	if err != nil {
		return err
	}
	err = taskQuery("update task set title=?, content=?, cat_id=?, priority = ? where id=? and user_id=?", title, content, categoryID, priority, id, userID)
	return err
}

//taskQuery encapsulates running multiple queries which don't do much things
func taskQuery(sql string, args ...interface{}) error {
	log.Print("inside task query")
	SQL := database.prepare(sql)
	tx := database.begin()
	_, err = tx.Stmt(SQL).Exec(args...)
	if err != nil {
		log.Println("taskQuery: ", err)
		tx.Rollback()
	} else {
		err = tx.Commit()
		if err != nil {
			log.Println(err)
			return err
		}
		log.Println("Commit successful")
	}
	return err
}

//SearchTask is used to return the search results depending on the query
func SearchTask(username, query string) (types.Context, error) {
	var tasks []types.Task
	var task types.Task
	var TaskCreated time.Time
	var context types.Context

	comments, err := GetComments(username)
	if err != nil {
		log.Println("SearchTask: something went wrong in finding comments")
	}

	userID, err := GetUserID(username)
	if err != nil {
		return context, err
	}

	stmt := "select t.id, title, content, created_date, priority, c.name from task t, category c where t.user_id=? and c.id = t.cat_id and (title like '%" + query + "%' or content like '%" + query + "%') order by created_date desc"

	rows := database.query(stmt, userID, query, query)
	defer rows.Close()
	for rows.Next() {
		err := rows.Scan(&task.Id, &task.Title, &task.Content, &TaskCreated, &task.Priority, &task.Category)
		if err != nil {
			log.Println(err)
		}

		if comments[task.Id] != nil {
			task.Comments = comments[task.Id]
		}

		task.Title = strings.Replace(task.Title, query, ""+query+"", -1)
		task.Content = strings.Replace(task.Content, query, ""+query+"", -1)
		task.Content = string(md.Markdown([]byte(task.Content)))

		TaskCreated = TaskCreated.Local()
		CurrentTime := time.Now().Local()
		week := TaskCreated.AddDate(0, 0, 7)

		if (week.String() < CurrentTime.String()) && (task.Priority != "1") {
			task.IsOverdue = true // If one week then overdue by default
		}
		task.Created = TaskCreated.Format("Jan 2 2006")

		tasks = append(tasks, task)
	}
	context = types.Context{Tasks: tasks, Search: query, Navigation: "search"}
	return context, nil
}

//GetComments is used to get comments, all of them.
//We do not want 100 different pages to show tasks, we want to use as few pages as possible
//so we are going to populate everything on the damn home pages
func GetComments(username string) (map[int][]types.Comment, error) {
	commentMap := make(map[int][]types.Comment)

	var taskID int
	var comment types.Comment
	var created time.Time

	userID, err := GetUserID(username)
	if err != nil {
		return commentMap, err
	}
	stmt := "select c.id, c.taskID, c.content, c.created from comments c, task t where t.id=c.taskID and c.user_id=?;"
	rows := database.query(stmt, userID)

	defer rows.Close()
	for rows.Next() {
		err := rows.Scan(&comment.ID, &taskID, &comment.Content, &created)
		if err != nil {
			return commentMap, err
		}
		// comment.Content = string(md.Markdown([]byte(comment.Content))) ## have to fix the
 issue markdown support
		created = created.Local()
		comment.Created = created.Format("Jan 2 2006 15:04:05")
		commentMap[taskID] = append(commentMap[taskID], comment)
	}
	return commentMap, nil
}

//AddComments will be used to add comments in the database
func AddComments(username string, id int, comment string) error {
	userID, err := GetUserID(username)
	if err != nil {
		return err
	}
	stmt := "insert into comments(taskID, content, created, user_id) values (?,?,datetime(),?)"
	err = taskQuery(stmt, id, comment, userID)

	if err != nil {
		return err
	}

	log.Println("added comment to task ID ", id)

	return nil
}
<|endoftext|>


<commit_before>// query_benchmarker speed tests servers using requests from stdin.
//
// It reads encoded Query objects from stdin, and makes concurrent requests
// to the provided HTTP endpoint. This program has no knowledge of the
// internals of the endpoint.
//
// TODO(rw): On my machine, this only decodes 700k/sec messages from stdin.
package main

import (
	"bufio"
	"encoding/gob"
	"flag"
	"fmt"
	"io"
	"log"
	"os"
	"runtime/pprof"
	"sort"
	"time"
	"sync"
)

// Program option vars:
var (
	daemonUrl     string
	workers       int
	debug         int
	prettyPrintResponses         bool
	limit         int64
	printInterval int64
	memProfile    string
)

// Global vars:
var (
	queryPool    sync.Pool
	queryChan    chan *Query
	statPool     sync.Pool
	statChan     chan *Stat
	workersGroup sync.WaitGroup
	statGroup    sync.WaitGroup
)

// Parse args:
func init() {
	flag.StringVar(&daemonUrl, "url", "http://localhost:8086", "Daemon URL.")
	flag.IntVar(&workers, "workers", 1, "Number of concurrent requests to make.")
	flag.IntVar(&debug, "debug", 0, "Whether to print debug messages.")
	flag.Int64Var(&limit, "limit", -1, "Limit the number of queries to send.")
	flag.Int64Var(&printInterval, "print-interval", 100, "Print timing stats to stderr after this many queries (0 to disable)")
	flag.BoolVar(&prettyPrintResponses, "print-responses", false, "Pretty print JSON response bodies (for correctness checking) (default false).")
	flag.StringVar(&memProfile, "memprofile", "", "Write a memory profile to this file.")

	flag.Parse()
}

func main() {
	// Make pools to minimize heap usage:
	queryPool = sync.Pool{
		New: func() interface{} {
			return &Query{
				HumanLabel:       make([]byte, 0, 1024),
				HumanDescription: make([]byte, 0, 1024),
				Method:           make([]byte, 0, 1024),
				Path:             make([]byte, 0, 1024),
				Body:             make([]byte, 0, 1024),
			}
		},
	}

	statPool = sync.Pool{
		New: func() interface{} {
			return &Stat{
				Label: make([]byte, 0, 1024),
				Value: 0.0,
			}
		},
	}

	// Make data and control channels:
	queryChan = make(chan *Query, workers)
	statChan = make(chan *Stat, workers)

	// Launch the stats processor:
	statGroup.Add(1)
	go processStats()

	// Launch the query processors:
	for i := 0; i < workers; i++ {
		workersGroup.Add(1)
		w := NewHTTPClient(daemonUrl, debug)
		go processQueries(w)
	}

	// Read in jobs, closing the job channel when done:
	input := bufio.NewReaderSize(os.Stdin, 1<<20)
	wallStart := time.Now()
	scan(input)
	close(queryChan)

	// Block for workers to finish sending requests, closing the stats
	// channel when done:
	workersGroup.Wait()
	close(statChan)

	// Wait on the stat collector to finish (and print its results):
	statGroup.Wait()

	wallEnd := time.Now()
	wallTook := wallEnd.Sub(wallStart)
	_, err := fmt.Printf("wall clock time: %fsec\n", float64(wallTook.Nanoseconds()) / 1e9)
	if err != nil {
		log.Fatal(err)
	}

	// (Optional) create a memory profile:
	if memProfile != "" {
		f, err := os.Create(memProfile)
		if err != nil {
			log.Fatal(err)
		}
		pprof.WriteHeapProfile(f)
		f.Close()
	}
}

// scan reads encoded Queries and places them onto the workqueue.
func scan(r io.Reader) {
	dec := gob.NewDecoder(r)

	n := int64(0)
	for {
		if limit >= 0 && n >= limit {
			break
		}

		q := queryPool.Get().(*Query)
		err := dec.Decode(q)
		if err == io.EOF {
			break
		}
		if err != nil {
			log.Fatal(err)
		}

		q.ID = n

		queryChan <- q

		n++

	}
}

// processQueries reads byte buffers from queryChan and writes them to the
// target server, while tracking latency.
func processQueries(w *HTTPClient) {
	opts := &HTTPClientDoOptions{
		Debug: debug,
		PrettyPrintResponses: prettyPrintResponses,
	}
	for q := range queryChan {
		lag, err := w.Do(q, opts)

		stat := statPool.Get().(*Stat)
		stat.Init(q.HumanLabel, lag)
		statChan <- stat

		queryPool.Put(q)
		if err != nil {
			log.Fatalf("Error during request: %s\n", err.Error())
		}
	}
	workersGroup.Done()
}

// processStats collects latency results, aggregating them into summary
// statistics. Optionally, they are printed to stderr at regular intervals.
func processStats() {
	const allQueriesLabel = "all queries           "
	statMapping := map[string]*StatGroup{
		allQueriesLabel: &StatGroup{},
	}

	i := int64(0)
	for stat := range statChan {
		if _, ok := statMapping[string(stat.Label)]; !ok {
			statMapping[string(stat.Label)] = &StatGroup{}
		}

		statMapping[allQueriesLabel].Push(stat.Value)
		statMapping[string(stat.Label)].Push(stat.Value)

		statPool.Put(stat)

		i++

		// print stats to stderr (if printInterval is greater than zero):
		if printInterval > 0 && i > 0 && i%printInterval == 0 && (i < limit || limit < 0) {
			_, err := fmt.Fprintf(os.Stderr, "after %d queries with %d workers:\n", i, workers)
			if err != nil {
				log.Fatal(err)
			}
			fprintStats(os.Stderr, statMapping)
			_, err = fmt.Fprintf(os.Stderr, "\n")
			if err != nil {
				log.Fatal(err)
			}
		}
	}

	// the final stats output goes to stdout:
	_, err := fmt.Printf("run complete after %d queries with %d workers:\n", i, workers)
	if err != nil {
		log.Fatal(err)
	}
	fprintStats(os.Stdout, statMapping)
	statGroup.Done()
}

// fprintStats pretty-prints stats to the given writer.
func fprintStats(w io.Writer, statGroups map[string]*StatGroup) {
	keys := make([]string, 0, len(statGroups))
	for k := range statGroups {
		keys = append(keys, k)
	}
	sort.Strings(keys)
	for _, k := range keys {
		v := statGroups[k]
		minRate := 1e3 / v.Min
		meanRate := 1e3 / v.Mean
		maxRate := 1e3 / v.Max
		_, err := fmt.Fprintf(w, "%s: min: %8.2fms (%7.2f/sec), mean: %8.2fms (%7.2f/sec), max: %7.2fms (%6.2f/sec), count: %8d, sum: %5.1fsec \n", k, v.Min, minRate, v.Mean, meanRate, v.Max, maxRate, v.Count, v.Sum/1e3)
		if err != nil {
			log.Fatal(err)
		}
	}

}
<commit_msg>Make summary stats easier to read by padding their labels.<commit_after>// query_benchmarker speed tests servers using requests from stdin.
//
// It reads encoded Query objects from stdin, and makes concurrent requests
// to the provided HTTP endpoint. This program has no knowledge of the
// internals of the endpoint.
//
// TODO(rw): On my machine, this only decodes 700k/sec messages from stdin.
package main

import (
	"bufio"
	"encoding/gob"
	"flag"
	"fmt"
	"io"
	"log"
	"os"
	"runtime/pprof"
	"sort"
	"sync"
	"time"
)

// Program option vars:
var (
	daemonUrl            string
	workers              int
	debug                int
	prettyPrintResponses bool
	limit                int64
	printInterval        int64
	memProfile           string
)

// Global vars:
var (
	queryPool    sync.Pool
	queryChan    chan *Query
	statPool     sync.Pool
	statChan     chan *Stat
	workersGroup sync.WaitGroup
	statGroup    sync.WaitGroup
)

// Parse args:
func init() {
	flag.StringVar(&daemonUrl, "url", "http://localhost:8086", "Daemon URL.")
	flag.IntVar(&workers, "workers", 1, "Number of concurrent requests to make.")
	flag.IntVar(&debug, "debug", 0, "Whether to print debug messages.")
	flag.Int64Var(&limit, "limit", -1, "Limit the number of queries to send.")
	flag.Int64Var(&printInterval, "print-interval", 100, "Print timing stats to stderr after this many queries (0 to disable)")
	flag.BoolVar(&prettyPrintResponses, "print-responses", false, "Pretty print JSON response bodies (for correctness checking) (default false).")
	flag.StringVar(&memProfile, "memprofile", "", "Write a memory profile to this file.")

	flag.Parse()
}

func main() {
	// Make pools to minimize heap usage:
	queryPool = sync.Pool{
		New: func() interface{} {
			return &Query{
				HumanLabel:       make([]byte, 0, 1024),
				HumanDescription: make([]byte, 0, 1024),
				Method:           make([]byte, 0, 1024),
				Path:             make([]byte, 0, 1024),
				Body:             make([]byte, 0, 1024),
			}
		},
	}

	statPool = sync.Pool{
		New: func() interface{} {
			return &Stat{
				Label: make([]byte, 0, 1024),
				Value: 0.0,
			}
		},
	}

	// Make data and control channels:
	queryChan = make(chan *Query, workers)
	statChan = make(chan *Stat, workers)

	// Launch the stats processor:
	statGroup.Add(1)
	go processStats()

	// Launch the query processors:
	for i := 0; i < workers; i++ {
		workersGroup.Add(1)
		w := NewHTTPClient(daemonUrl, debug)
		go processQueries(w)
	}

	// Read in jobs, closing the job channel when done:
	input := bufio.NewReaderSize(os.Stdin, 1<<20)
	wallStart := time.Now()
	scan(input)
	close(queryChan)

	// Block for workers to finish sending requests, closing the stats
	// channel when done:
	workersGroup.Wait()
	close(statChan)

	// Wait on the stat collector to finish (and print its results):
	statGroup.Wait()

	wallEnd := time.Now()
	wallTook := wallEnd.Sub(wallStart)
	_, err := fmt.Printf("wall clock time: %fsec\n", float64(wallTook.Nanoseconds())/1e9)
	if err != nil {
		log.Fatal(err)
	}

	// (Optional) create a memory profile:
	if memProfile != "" {
		f, err := os.Create(memProfile)
		if err != nil {
			log.Fatal(err)
		}
		pprof.WriteHeapProfile(f)
		f.Close()
	}
}

// scan reads encoded Queries and places them onto the workqueue.
func scan(r io.Reader) {
	dec := gob.NewDecoder(r)

	n := int64(0)
	for {
		if limit >= 0 && n >= limit {
			break
		}

		q := queryPool.Get().(*Query)
		err := dec.Decode(q)
		if err == io.EOF {
			break
		}
		if err != nil {
			log.Fatal(err)
		}

		q.ID = n

		queryChan <- q

		n++

	}
}

// processQueries reads byte buffers from queryChan and writes them to the
// target server, while tracking latency.
func processQueries(w *HTTPClient) {
	opts := &HTTPClientDoOptions{
		Debug:                debug,
		PrettyPrintResponses: prettyPrintResponses,
	}
	for q := range queryChan {
		lag, err := w.Do(q, opts)

		stat := statPool.Get().(*Stat)
		stat.Init(q.HumanLabel, lag)
		statChan <- stat

		queryPool.Put(q)
		if err != nil {
			log.Fatalf("Error during request: %s\n", err.Error())
		}
	}
	workersGroup.Done()
}

// processStats collects latency results, aggregating them into summary
// statistics. Optionally, they are printed to stderr at regular intervals.
func processStats() {
	const allQueriesLabel = "all queries"
	statMapping := map[string]*StatGroup{
		allQueriesLabel: &StatGroup{},
	}

	i := int64(0)
	for stat := range statChan {
		if _, ok := statMapping[string(stat.Label)]; !ok {
			statMapping[string(stat.Label)] = &StatGroup{}
		}

		statMapping[allQueriesLabel].Push(stat.Value)
		statMapping[string(stat.Label)].Push(stat.Value)

		statPool.Put(stat)

		i++

		// print stats to stderr (if printInterval is greater than zero):
		if printInterval > 0 && i > 0 && i%printInterval == 0 && (i < limit || limit < 0) {
			_, err := fmt.Fprintf(os.Stderr, "after %d queries with %d workers:\n", i, workers)
			if err != nil {
				log.Fatal(err)
			}
			fprintStats(os.Stderr, statMapping)
			_, err = fmt.Fprintf(os.Stderr, "\n")
			if err != nil {
				log.Fatal(err)
			}
		}
	}

	// the final stats output goes to stdout:
	_, err := fmt.Printf("run complete after %d queries with %d workers:\n", i, workers)
	if err != nil {
		log.Fatal(err)
	}
	fprintStats(os.Stdout, statMapping)
	statGroup.Done()
}

// fprintStats pretty-prints stats to the given writer.
func fprintStats(w io.Writer, statGroups map[string]*StatGroup) {
	maxKeyLength := 0
	keys := make([]string, 0, len(statGroups))
	for k := range statGroups {
		if len(k) > maxKeyLength {
			maxKeyLength = len(k)
		}
		keys = append(keys, k)
	}
	sort.Strings(keys)
	for _, k := range keys {
		v := statGroups[k]
		minRate := 1e3 / v.Min
		meanRate := 1e3 / v.Mean
		maxRate := 1e3 / v.Max
		paddedKey := fmt.Sprintf("%s", k)
		for len(paddedKey) < maxKeyLength {
			paddedKey += " "
		}
		_, err := fmt.Fprintf(w, "%s : min: %8.2fms (%7.2f/sec), mean: %8.2fms (%7.2f/sec), max: %7.2fms (%6.2f/sec), count: %8d, sum: %5.1fsec \n", paddedKey, v.Min, minRate, v.Mean, meanRate, v.Max, maxRate, v.Count, v.Sum/1e3)
		if err != nil {
			log.Fatal(err)
		}
	}

}
<|endoftext|>


<commit_before>package azlyrics

import (
	"bytes"
	"errors"
	"github.com/alewmoose/show-lyrics/songinfo"
	"net/http"
	"strings"
	"testing"
)

type test struct {
	si   songinfo.SongInfo
	err  error
	text []byte
}

var tests = [...]test{
	{
		si: songinfo.SongInfo{
			Artist: "Mastodon",
			Title:  "Asleep In The Deep",
		},
		err: nil,
		text: []byte(strings.Join([]string{
			"The moment you walked in the room, my friend",
			"the demons, they all went away",
			"be careful, they're only asleep for a while",
			"pretending there's nothing to say",
			"",
			"Throw salt in all the corners here",
			"make sure you watch him leave",
			"",
			"Build up the walls around this house",
			"and dig out the rot in the floor",
			"block out the entrance with brick and stone",
			"and mortar that's made from coal",
			"",
			"Crawl into this hole I've made",
			"transform these feelings of fear",
			"",
			"I'm on fire",
			"say you'll remember her voice",
			"and I can't get you out of my mind",
			"",
			"Loose lips have fallen on deaf ears",
			"loose lips have fallen on blind eyes",
			"",
			"An ocean of sorrow surrounds this home",
			"I hope that we make it to shore",
			"as time chips away at the fortress walls",
			"it seems that we weathered the storm",
			"",
			"The sun begins to show itself",
			"revealing victory",
			"",
			"I'm on fire",
			"say you'll remember her voice",
			"and I can't get you out of my mind",
		}, "\n")),
	},
	{
		si: songinfo.SongInfo{
			Artist: "naosehntaoshftrdru",
			Title:  "aosehntaoshftrdrutn",
		},
		err:  errors.New("404 Not Found"),
		text: []byte{},
	},
}

func TestFetch(t *testing.T) {
	const fmt = "Failed test #%d:\ngot:\n\"%s\"\nexpected:\n\"%s\"\n"
	client := &http.Client{}

	for i, fetchTest := range tests {
		text, err := Fetch(client, &fetchTest.si)

		if !bytes.Equal(text, fetchTest.text) {
			t.Errorf(fmt, i+1, string(text), string(fetchTest.text))
		}

		var gotErr, expErr string
		if err != nil {
			gotErr = err.Error()
		}
		if fetchTest.err != nil {
			expErr = fetchTest.err.Error()
		}
		if gotErr != expErr {
			t.Errorf(fmt, i+1, gotErr, expErr)
		}
	}
}
<commit_msg>Add test<commit_after>package azlyrics

import (
	"bytes"
	"errors"
	"github.com/alewmoose/show-lyrics/songinfo"
	"net/http"
	"strings"
	"testing"
)

type test struct {
	si   songinfo.SongInfo
	err  error
	text []byte
}

var tests = [...]test{
	{
		si: songinfo.SongInfo{
			Artist: "Mastodon",
			Title:  "Asleep In The Deep",
		},
		err: nil,
		text: []byte(strings.Join([]string{
			"The moment you walked in the room, my friend",
			"the demons, they all went away",
			"be careful, they're only asleep for a while",
			"pretending there's nothing to say",
			"",
			"Throw salt in all the corners here",
			"make sure you watch him leave",
			"",
			"Build up the walls around this house",
			"and dig out the rot in the floor",
			"block out the entrance with brick and stone",
			"and mortar that's made from coal",
			"",
			"Crawl into this hole I've made",
			"transform these feelings of fear",
			"",
			"I'm on fire",
			"say you'll remember her voice",
			"and I can't get you out of my mind",
			"",
			"Loose lips have fallen on deaf ears",
			"loose lips have fallen on blind eyes",
			"",
			"An ocean of sorrow surrounds this home",
			"I hope that we make it to shore",
			"as time chips away at the fortress walls",
			"it seems that we weathered the storm",
			"",
			"The sun begins to show itself",
			"revealing victory",
			"",
			"I'm on fire",
			"say you'll remember her voice",
			"and I can't get you out of my mind",
		}, "\n")),
	},
	{
		si: songinfo.SongInfo{
			Artist: "A Perfect circle",
			Title:  "Breña",
		},
		err: nil,
		text: []byte(strings.Join([]string{
			"My reflection",
			"Wraps and pulls me under",
			"healing waters to be",
			"Bathed in Brena",
			"",
			"Guides me",
			"Safely in",
			"Worlds I've never been to",
			"Heal me",
			"Heal me",
			"My dear Brena",
			"",
			"So vulnerable",
			"But it's alright",
			"",
			"Heal me",
			"Heal me",
			"My dear Brena",
			"",
			"Show me lonely and",
			"Show me openings",
			"To lead me closer to you",
			"My dear Brena",
			"",
			"(Feeling so) vulnerable",
			"But it's alright",
			"",
			"Opening to... heal...",
			"Opening to... heal...",
			"Heal.. Heal.. Heal...",
			"",
			"Heal me",
		}, "\n")),
	},
	{
		si: songinfo.SongInfo{
			Artist: "naosehntaoshftrdru",
			Title:  "aosehntaoshftrdrutn",
		},
		err:  errors.New("404 Not Found"),
		text: []byte{},
	},
}

func TestFetch(t *testing.T) {
	const fmt = "Failed test #%d:\ngot:\n\"%s\"\nexpected:\n\"%s\"\n"
	client := &http.Client{}

	for i, fetchTest := range tests {
		text, err := Fetch(client, &fetchTest.si)

		if !bytes.Equal(text, fetchTest.text) {
			t.Errorf(fmt, i+1, string(text), string(fetchTest.text))
		}

		var gotErr, expErr string
		if err != nil {
			gotErr = err.Error()
		}
		if fetchTest.err != nil {
			expErr = fetchTest.err.Error()
		}
		if gotErr != expErr {
			t.Errorf(fmt, i+1, gotErr, expErr)
		}
	}
}
<|endoftext|>


<commit_before>package kafka

import (
	"bufio"
	"bytes"
	"fmt"
	"io"
	"time"
)

// Message is a data structure representing kafka messages.
type Message struct {
	// Topic indicates which topic this message was consumed from via Reader.
	//
	// When being used with Writer, this can be used to configured the topic if
	// not already specified on the writer itself.
	Topic string

	// Partition is read-only and MUST NOT be set when writing messages
	Partition     int
	Offset        int64
	HighWaterMark int64
	Key           []byte
	Value         []byte
	Headers       []Header

	// If not set at the creation, Time will be automatically set when
	// writing the message.
	Time time.Time
}

func (msg Message) message(cw *crc32Writer) message {
	m := message{
		MagicByte: 1,
		Key:       msg.Key,
		Value:     msg.Value,
		Timestamp: timestamp(msg.Time),
	}
	if cw != nil {
		m.CRC = m.crc32(cw)
	}
	return m
}

const timestampSize = 8

func (msg *Message) size() int32 {
	return 4 + 1 + 1 + sizeofBytes(msg.Key) + sizeofBytes(msg.Value) + timestampSize
}

type message struct {
	CRC        int32
	MagicByte  int8
	Attributes int8
	Timestamp  int64
	Key        []byte
	Value      []byte
}

func (m message) crc32(cw *crc32Writer) int32 {
	cw.crc32 = 0
	cw.writeInt8(m.MagicByte)
	cw.writeInt8(m.Attributes)
	if m.MagicByte != 0 {
		cw.writeInt64(m.Timestamp)
	}
	cw.writeBytes(m.Key)
	cw.writeBytes(m.Value)
	return int32(cw.crc32)
}

func (m message) size() int32 {
	size := 4 + 1 + 1 + sizeofBytes(m.Key) + sizeofBytes(m.Value)
	if m.MagicByte != 0 {
		size += timestampSize
	}
	return size
}

func (m message) writeTo(wb *writeBuffer) {
	wb.writeInt32(m.CRC)
	wb.writeInt8(m.MagicByte)
	wb.writeInt8(m.Attributes)
	if m.MagicByte != 0 {
		wb.writeInt64(m.Timestamp)
	}
	wb.writeBytes(m.Key)
	wb.writeBytes(m.Value)
}

type messageSetItem struct {
	Offset      int64
	MessageSize int32
	Message     message
}

func (m messageSetItem) size() int32 {
	return 8 + 4 + m.Message.size()
}

func (m messageSetItem) writeTo(wb *writeBuffer) {
	wb.writeInt64(m.Offset)
	wb.writeInt32(m.MessageSize)
	m.Message.writeTo(wb)
}

type messageSet []messageSetItem

func (s messageSet) size() (size int32) {
	for _, m := range s {
		size += m.size()
	}
	return
}

func (s messageSet) writeTo(wb *writeBuffer) {
	for _, m := range s {
		m.writeTo(wb)
	}
}

type messageSetReader struct {
	empty   bool
	version int
	v1      messageSetReaderV1
	v2      messageSetReaderV2
}

func (r *messageSetReader) readMessage(min int64,
	key func(*bufio.Reader, int, int) (int, error),
	val func(*bufio.Reader, int, int) (int, error),
) (offset int64, timestamp int64, headers []Header, err error) {
	if r.empty {
		return 0, 0, nil, RequestTimedOut
	}
	switch r.version {
	case 1:
		return r.v1.readMessage(min, key, val)
	case 2:
		return r.v2.readMessage(min, key, val)
	default:
		panic("Invalid messageSetReader - unknown message reader version")
	}
}

func (r *messageSetReader) remaining() (remain int) {
	if r.empty {
		return 0
	}
	switch r.version {
	case 1:
		return r.v1.remaining()
	case 2:
		return r.v2.remaining()
	default:
		panic("Invalid messageSetReader - unknown message reader version")
	}
}

func (r *messageSetReader) discard() (err error) {
	if r.empty {
		return nil
	}
	switch r.version {
	case 1:
		return r.v1.discard()
	case 2:
		return r.v2.discard()
	default:
		panic("Invalid messageSetReader - unknown message reader version")
	}
}

type messageSetReaderV1 struct {
	*readerStack
}

type readerStack struct {
	reader *bufio.Reader
	remain int
	base   int64
	parent *readerStack
}

func newMessageSetReader(reader *bufio.Reader, remain int) (*messageSetReader, error) {
	headerLength := 8 + 4 + 4 + 1 // offset + messageSize + crc + magicByte

	if headerLength > remain {
		return nil, errShortRead
	}

	b, err := reader.Peek(headerLength)
	if err != nil {
		return nil, err
	}
	var version int8 = int8(b[headerLength-1])

	switch version {
	case 0, 1:
		return &messageSetReader{
			version: 1,
			v1: messageSetReaderV1{&readerStack{
				reader: reader,
				remain: remain,
			}}}, nil
	case 2:
		mr := &messageSetReader{
			version: 2,
			v2: messageSetReaderV2{
				readerStack: &readerStack{
					reader: reader,
					remain: remain,
				},
				messageCount: 0,
			}}
		return mr, nil
	default:
		return nil, fmt.Errorf("unsupported message version %d found in fetch response", version)
	}
}

func (r *messageSetReaderV1) readMessage(min int64,
	key func(*bufio.Reader, int, int) (int, error),
	val func(*bufio.Reader, int, int) (int, error),
) (offset int64, timestamp int64, headers []Header, err error) {
	for r.readerStack != nil {
		if r.remain == 0 {
			r.readerStack = r.parent
			continue
		}

		var attributes int8
		if offset, attributes, timestamp, r.remain, err = readMessageHeader(r.reader, r.remain); err != nil {
			return
		}

		// if the message is compressed, decompress it and push a new reader
		// onto the stack.
		code := attributes & compressionCodecMask
		if code != 0 {
			var codec CompressionCodec
			if codec, err = resolveCodec(code); err != nil {
				return
			}

			// discard next four bytes...will be -1 to indicate null key
			if r.remain, err = discardN(r.reader, r.remain, 4); err != nil {
				return
			}

			// read and decompress the contained message set.
			var decompressed bytes.Buffer

			if r.remain, err = readBytesWith(r.reader, r.remain, func(r *bufio.Reader, sz, n int) (remain int, err error) {
				// x4 as a guess that the average compression ratio is near 75%
				decompressed.Grow(4 * n)

				l := io.LimitedReader{R: r, N: int64(n)}
				d := codec.NewReader(&l)

				_, err = decompressed.ReadFrom(d)
				remain = sz - (n - int(l.N))

				d.Close()
				return
			}); err != nil {
				return
			}

			// the compressed message's offset will be equal to the offset of
			// the last message in the set.  within the compressed set, the
			// offsets will be relative, so we have to scan through them to
			// get the base offset.  for example, if there are four compressed
			// messages at offsets 10-13, then the container message will have
			// offset 13 and the contained messages will be 0,1,2,3.  the base
			// offset for the container, then is 13-3=10.
			if offset, err = extractOffset(offset, decompressed.Bytes()); err != nil {
				return
			}

			r.readerStack = &readerStack{
				// Allocate a buffer of size 0, which gets capped at 16 bytes
				// by the bufio package. We are already reading buffered data
				// here, no need to reserve another 4KB buffer.
				reader: bufio.NewReaderSize(&decompressed, 0),
				remain: decompressed.Len(),
				base:   offset,
				parent: r.readerStack,
			}
			continue
		}

		// adjust the offset in case we're reading compressed messages.  the
		// base will be zero otherwise.
		offset += r.base

		// When the messages are compressed kafka may return messages at an
		// earlier offset than the one that was requested, it's the client's
		// responsibility to ignore those.
		if offset < min {
			if r.remain, err = discardBytes(r.reader, r.remain); err != nil {
				return
			}
			if r.remain, err = discardBytes(r.reader, r.remain); err != nil {
				return
			}
			continue
		}

		if r.remain, err = readBytesWith(r.reader, r.remain, key); err != nil {
			return
		}
		r.remain, err = readBytesWith(r.reader, r.remain, val)
		return
	}

	err = errShortRead
	return
}

func (r *messageSetReaderV1) remaining() (remain int) {
	for s := r.readerStack; s != nil; s = s.parent {
		remain += s.remain
	}
	return
}

func (r *messageSetReaderV1) discard() (err error) {
	if r.readerStack == nil {
		return
	}
	// rewind up to the top-most reader b/c it's the only one that's doing
	// actual i/o.  the rest are byte buffers that have been pushed on the stack
	// while reading compressed message sets.
	for r.parent != nil {
		r.readerStack = r.parent
	}
	r.remain, err = discardN(r.reader, r.remain, r.remain)
	return
}

func extractOffset(base int64, msgSet []byte) (offset int64, err error) {
	r, remain := bufio.NewReader(bytes.NewReader(msgSet)), len(msgSet)
	for remain > 0 {
		if remain, err = readInt64(r, remain, &offset); err != nil {
			return
		}
		var sz int32
		if remain, err = readInt32(r, remain, &sz); err != nil {
			return
		}
		if remain, err = discardN(r, remain, int(sz)); err != nil {
			return
		}
	}
	offset = base - offset
	return
}

type messageSetHeaderV2 struct {
	firstOffset          int64
	length               int32
	partitionLeaderEpoch int32
	magic                int8
	crc                  int32
	batchAttributes      int16
	lastOffsetDelta      int32
	firstTimestamp       int64
	maxTimestamp         int64
	producerId           int64
	producerEpoch        int16
	firstSequence        int32
}

type timestampType int8

const (
	createTime    timestampType = 0
	logAppendTime timestampType = 1
)

type transactionType int8

const (
	nonTransactional transactionType = 0
	transactional    transactionType = 1
)

type controlType int8

const (
	nonControlMessage controlType = 0
	controlMessage    controlType = 1
)

func (h *messageSetHeaderV2) compression() int8 {
	return int8(h.batchAttributes & 7)
}

func (h *messageSetHeaderV2) timestampType() timestampType {
	return timestampType((h.batchAttributes & (1 << 3)) >> 3)
}

func (h *messageSetHeaderV2) transactionType() transactionType {
	return transactionType((h.batchAttributes & (1 << 4)) >> 4)
}

func (h *messageSetHeaderV2) controlType() controlType {
	return controlType((h.batchAttributes & (1 << 5)) >> 5)
}

type messageSetReaderV2 struct {
	*readerStack
	messageCount int

	header messageSetHeaderV2
}

func (r *messageSetReaderV2) readHeader() (err error) {
	h := &r.header
	if r.remain, err = readInt64(r.reader, r.remain, &h.firstOffset); err != nil {
		return
	}
	if r.remain, err = readInt32(r.reader, r.remain, &h.length); err != nil {
		return
	}
	if r.remain, err = readInt32(r.reader, r.remain, &h.partitionLeaderEpoch); err != nil {
		return
	}
	if r.remain, err = readInt8(r.reader, r.remain, &h.magic); err != nil {
		return
	}
	if r.remain, err = readInt32(r.reader, r.remain, &h.crc); err != nil {
		return
	}
	if r.remain, err = readInt16(r.reader, r.remain, &h.batchAttributes); err != nil {
		return
	}
	if r.remain, err = readInt32(r.reader, r.remain, &h.lastOffsetDelta); err != nil {
		return
	}
	if r.remain, err = readInt64(r.reader, r.remain, &h.firstTimestamp); err != nil {
		return
	}
	if r.remain, err = readInt64(r.reader, r.remain, &h.maxTimestamp); err != nil {
		return
	}
	if r.remain, err = readInt64(r.reader, r.remain, &h.producerId); err != nil {
		return
	}
	if r.remain, err = readInt16(r.reader, r.remain, &h.producerEpoch); err != nil {
		return
	}
	if r.remain, err = readInt32(r.reader, r.remain, &h.firstSequence); err != nil {
		return
	}
	var messageCount int32
	if r.remain, err = readInt32(r.reader, r.remain, &messageCount); err != nil {
		return
	}
	r.messageCount = int(messageCount)

	return nil
}

func (r *messageSetReaderV2) readMessage(min int64,
	key func(*bufio.Reader, int, int) (int, error),
	val func(*bufio.Reader, int, int) (int, error),
) (offset int64, timestamp int64, headers []Header, err error) {

	if r.messageCount == 0 {
		if r.remain == 0 {
			if r.parent != nil {
				r.readerStack = r.parent
			}
		}

		if err = r.readHeader(); err != nil {
			return
		}

		if code := r.header.compression(); code != 0 {
			var codec CompressionCodec
			if codec, err = resolveCodec(code); err != nil {
				return
			}

			var batchRemain = int(r.header.length - 49)
			if batchRemain > r.remain {
				err = errShortRead
				return
			}

			var decompressed bytes.Buffer
			decompressed.Grow(4 * batchRemain)

			l := io.LimitedReader{R: r.reader, N: int64(batchRemain)}
			d := codec.NewReader(&l)

			_, err = decompressed.ReadFrom(d)
			r.remain = r.remain - (batchRemain - int(l.N))
			d.Close()

			if err != nil {
				return
			}

			r.readerStack = &readerStack{
				reader: bufio.NewReaderSize(&decompressed, 0),
				remain: decompressed.Len(),
				base:   -1, // base is unused here
				parent: r.readerStack,
			}
		}
	}

	var length int64
	if r.remain, err = readVarInt(r.reader, r.remain, &length); err != nil {
		return
	}

	var attrs int8
	if r.remain, err = readInt8(r.reader, r.remain, &attrs); err != nil {
		return
	}
	var timestampDelta int64
	if r.remain, err = readVarInt(r.reader, r.remain, &timestampDelta); err != nil {
		return
	}
	var offsetDelta int64
	if r.remain, err = readVarInt(r.reader, r.remain, &offsetDelta); err != nil {
		return
	}
	var keyLen int64
	if r.remain, err = readVarInt(r.reader, r.remain, &keyLen); err != nil {
		return
	}

	if r.remain, err = key(r.reader, r.remain, int(keyLen)); err != nil {
		return
	}
	var valueLen int64
	if r.remain, err = readVarInt(r.reader, r.remain, &valueLen); err != nil {
		return
	}

	if r.remain, err = val(r.reader, r.remain, int(valueLen)); err != nil {
		return
	}

	var headerCount int64
	if r.remain, err = readVarInt(r.reader, r.remain, &headerCount); err != nil {
		return
	}

	headers = make([]Header, headerCount)

	for i := 0; i < int(headerCount); i++ {
		if err = r.readMessageHeader(&headers[i]); err != nil {
			return
		}
	}
	r.messageCount--
	return r.header.firstOffset + offsetDelta, r.header.firstTimestamp + timestampDelta, headers, nil
}

func (r *messageSetReaderV2) readMessageHeader(header *Header) (err error) {
	var keyLen int64
	if r.remain, err = readVarInt(r.reader, r.remain, &keyLen); err != nil {
		return
	}
	if header.Key, r.remain, err = readNewString(r.reader, r.remain, int(keyLen)); err != nil {
		return
	}
	var valLen int64
	if r.remain, err = readVarInt(r.reader, r.remain, &valLen); err != nil {
		return
	}
	if header.Value, r.remain, err = readNewBytes(r.reader, r.remain, int(valLen)); err != nil {
		return
	}
	return nil
}

func (r *messageSetReaderV2) remaining() (remain int) {
	return r.remain
}

func (r *messageSetReaderV2) discard() (err error) {
	r.remain, err = discardN(r.reader, r.remain, r.remain)
	return
}
<commit_msg>Remove existing message set reader.<commit_after>package kafka

import (
	"time"
)

// Message is a data structure representing kafka messages.
type Message struct {
	// Topic indicates which topic this message was consumed from via Reader.
	//
	// When being used with Writer, this can be used to configured the topic if
	// not already specified on the writer itself.
	Topic string

	// Partition is read-only and MUST NOT be set when writing messages
	Partition     int
	Offset        int64
	HighWaterMark int64
	Key           []byte
	Value         []byte
	Headers       []Header

	// If not set at the creation, Time will be automatically set when
	// writing the message.
	Time time.Time
}

func (msg Message) message(cw *crc32Writer) message {
	m := message{
		MagicByte: 1,
		Key:       msg.Key,
		Value:     msg.Value,
		Timestamp: timestamp(msg.Time),
	}
	if cw != nil {
		m.CRC = m.crc32(cw)
	}
	return m
}

const timestampSize = 8

func (msg *Message) size() int32 {
	return 4 + 1 + 1 + sizeofBytes(msg.Key) + sizeofBytes(msg.Value) + timestampSize
}

type message struct {
	CRC        int32
	MagicByte  int8
	Attributes int8
	Timestamp  int64
	Key        []byte
	Value      []byte
}

func (m message) crc32(cw *crc32Writer) int32 {
	cw.crc32 = 0
	cw.writeInt8(m.MagicByte)
	cw.writeInt8(m.Attributes)
	if m.MagicByte != 0 {
		cw.writeInt64(m.Timestamp)
	}
	cw.writeBytes(m.Key)
	cw.writeBytes(m.Value)
	return int32(cw.crc32)
}

func (m message) size() int32 {
	size := 4 + 1 + 1 + sizeofBytes(m.Key) + sizeofBytes(m.Value)
	if m.MagicByte != 0 {
		size += timestampSize
	}
	return size
}

func (m message) writeTo(wb *writeBuffer) {
	wb.writeInt32(m.CRC)
	wb.writeInt8(m.MagicByte)
	wb.writeInt8(m.Attributes)
	if m.MagicByte != 0 {
		wb.writeInt64(m.Timestamp)
	}
	wb.writeBytes(m.Key)
	wb.writeBytes(m.Value)
}

type messageSetItem struct {
	Offset      int64
	MessageSize int32
	Message     message
}

func (m messageSetItem) size() int32 {
	return 8 + 4 + m.Message.size()
}

func (m messageSetItem) writeTo(wb *writeBuffer) {
	wb.writeInt64(m.Offset)
	wb.writeInt32(m.MessageSize)
	m.Message.writeTo(wb)
}

type messageSet []messageSetItem

func (s messageSet) size() (size int32) {
	for _, m := range s {
		size += m.size()
	}
	return
}

func (s messageSet) writeTo(wb *writeBuffer) {
	for _, m := range s {
		m.writeTo(wb)
	}
}

type timestampType int8

const (
	createTime    timestampType = 0
	logAppendTime timestampType = 1
)

type transactionType int8

const (
	nonTransactional transactionType = 0
	transactional    transactionType = 1
)

type controlType int8

const (
	nonControlMessage controlType = 0
	controlMessage    controlType = 1
)
<|endoftext|>


<commit_before>package data

import (
	"crypto/rand"
	"encoding/base64"
	"log"
	"mime"
	"strings"
	"time"
)

// MessageID represents the ID of an SMTP message including the hostname part
type MessageID string

// NewMessageID generates a new message ID
func NewMessageID(hostname string) (MessageID, error) {
	size := 32

	rb := make([]byte, size)
	_, err := rand.Read(rb)

	if err != nil {
		return MessageID(""), err
	}

	rs := base64.URLEncoding.EncodeToString(rb)

	return MessageID(rs + "@" + hostname), nil
}

// Messages represents an array of Messages
// - TODO is this even required?
type Messages []Message

// Message represents a parsed SMTP message
type Message struct {
	ID      MessageID
	From    *Path
	To      []*Path
	Content *Content
	Created time.Time
	MIME    *MIMEBody // FIXME refactor to use Content.MIME
	Raw     *SMTPMessage
}

// Path represents an SMTP forward-path or return-path
type Path struct {
	Relays  []string
	Mailbox string
	Domain  string
	Params  string
}

// Content represents the body content of an SMTP message
type Content struct {
	Headers map[string][]string
	Body    string
	Size    int
	MIME    *MIMEBody
}

// SMTPMessage represents a raw SMTP message
type SMTPMessage struct {
	From string
	To   []string
	Data string
	Helo string
}

// MIMEBody represents a collection of MIME parts
type MIMEBody struct {
	Parts []*Content
}

// Parse converts a raw SMTP message to a parsed MIME message
func (m *SMTPMessage) Parse(hostname string) *Message {
	var arr []*Path
	for _, path := range m.To {
		arr = append(arr, PathFromString(path))
	}

	id, _ := NewMessageID(hostname)
	msg := &Message{
		ID:      id,
		From:    PathFromString(m.From),
		To:      arr,
		Content: ContentFromString(m.Data),
		Created: time.Now(),
		Raw:     m,
	}

	if msg.Content.IsMIME() {
		log.Printf("Parsing MIME body")
		msg.MIME = msg.Content.ParseMIMEBody()
	}

	// FIXME shouldn't be setting Message-ID, its a client thing
	msg.Content.Headers["Message-ID"] = []string{string(id)}
	msg.Content.Headers["Received"] = []string{"from " + m.Helo + " by " + hostname + " (Go-MailHog)\r\n          id " + string(id) + "; " + time.Now().Format(time.RFC1123Z)}
	msg.Content.Headers["Return-Path"] = []string{"<" + m.From + ">"}
	return msg
}

// IsMIME detects a valid MIME header
func (content *Content) IsMIME() bool {
	header, ok := content.Headers["Content-Type"]
	if !ok {
		return false
	}
	return strings.HasPrefix(header[0], "multipart/")
}

// ParseMIMEBody parses SMTP message content into multiple MIME parts
func (content *Content) ParseMIMEBody() *MIMEBody {
	var parts []*Content

	if hdr, ok := content.Headers["Content-Type"]; ok {
		if len(hdr) > 0 {
			boundary := extractBoundary(hdr[0])
			var p []string
			if len(boundary) > 0 {
				p = strings.Split(content.Body, "--"+boundary)
				log.Printf("Got boundary: %s", boundary)
			} else {
				log.Printf("Boundary not found: %s", hdr[0])
			}

			for _, s := range p {
				if len(s) > 0 {
					part := ContentFromString(strings.Trim(s, "\r\n"))
					if part.IsMIME() {
						log.Printf("Parsing inner MIME body")
						part.MIME = part.ParseMIMEBody()
					}
					parts = append(parts, part)
				}
			}
		}
	}

	return &MIMEBody{
		Parts: parts,
	}
}

// PathFromString parses a forward-path or reverse-path into its parts
func PathFromString(path string) *Path {
	var relays []string
	email := path
	if strings.Contains(path, ":") {
		x := strings.SplitN(path, ":", 2)
		r, e := x[0], x[1]
		email = e
		relays = strings.Split(r, ",")
	}
	mailbox, domain := "", ""
	if strings.Contains(email, "@") {
		x := strings.SplitN(email, "@", 2)
		mailbox, domain = x[0], x[1]
	} else {
		mailbox = email
	}

	return &Path{
		Relays:  relays,
		Mailbox: mailbox,
		Domain:  domain,
		Params:  "", // FIXME?
	}
}

// ContentFromString parses SMTP content into separate headers and body
func ContentFromString(data string) *Content {
	log.Printf("Parsing Content from string: '%s'", data)
	x := strings.SplitN(data, "\r\n\r\n", 2)
	h := make(map[string][]string, 0)

	if len(x) == 2 {
		headers, body := x[0], x[1]
		hdrs := strings.Split(headers, "\r\n")
		var lastHdr = ""
		for _, hdr := range hdrs {
			if lastHdr != "" && (strings.HasPrefix(hdr, " ") || strings.HasPrefix(hdr, "\t")) {
				h[lastHdr][len(h[lastHdr])-1] = h[lastHdr][len(h[lastHdr])-1] + hdr
			} else if strings.Contains(hdr, ": ") {
				y := strings.SplitN(hdr, ": ", 2)
				key, value := y[0], y[1]
				// TODO multiple header fields
				h[key] = []string{value}
				lastHdr = key
			} else {
				log.Printf("Found invalid header: '%s'", hdr)
			}
		}
		return &Content{
			Size:    len(data),
			Headers: h,
			Body:    body,
		}
	}
	return &Content{
		Size:    len(data),
		Headers: h,
		Body:    x[0],
	}
}

// extractBoundary extract boundary string in contentType.
// It returns empty string if no valid boundary found
func extractBoundary(contentType string) string {
	_, params, err := mime.ParseMediaType(contentType)
	if err == nil {
		return params["boundary"]
	}
	return ""
}
<commit_msg>Factor out use of log package to enable custom logging handlers. Fix #3<commit_after>package data

import (
	"crypto/rand"
	"encoding/base64"
	"log"
	"mime"
	"strings"
	"time"
)

// LogHandler is called for each log message. If nil, log messages will
// be output using log.Printf instead.
var LogHandler func(message string, args ...interface{})

func logf(message string, args ...interface{}) {
	if LogHandler != nil {
		LogHandler(message, args...)
	} else {
		log.Printf(message, args...)
	}
}

// MessageID represents the ID of an SMTP message including the hostname part
type MessageID string

// NewMessageID generates a new message ID
func NewMessageID(hostname string) (MessageID, error) {
	size := 32

	rb := make([]byte, size)
	_, err := rand.Read(rb)

	if err != nil {
		return MessageID(""), err
	}

	rs := base64.URLEncoding.EncodeToString(rb)

	return MessageID(rs + "@" + hostname), nil
}

// Messages represents an array of Messages
// - TODO is this even required?
type Messages []Message

// Message represents a parsed SMTP message
type Message struct {
	ID      MessageID
	From    *Path
	To      []*Path
	Content *Content
	Created time.Time
	MIME    *MIMEBody // FIXME refactor to use Content.MIME
	Raw     *SMTPMessage
}

// Path represents an SMTP forward-path or return-path
type Path struct {
	Relays  []string
	Mailbox string
	Domain  string
	Params  string
}

// Content represents the body content of an SMTP message
type Content struct {
	Headers map[string][]string
	Body    string
	Size    int
	MIME    *MIMEBody
}

// SMTPMessage represents a raw SMTP message
type SMTPMessage struct {
	From string
	To   []string
	Data string
	Helo string
}

// MIMEBody represents a collection of MIME parts
type MIMEBody struct {
	Parts []*Content
}

// Parse converts a raw SMTP message to a parsed MIME message
func (m *SMTPMessage) Parse(hostname string) *Message {
	var arr []*Path
	for _, path := range m.To {
		arr = append(arr, PathFromString(path))
	}

	id, _ := NewMessageID(hostname)
	msg := &Message{
		ID:      id,
		From:    PathFromString(m.From),
		To:      arr,
		Content: ContentFromString(m.Data),
		Created: time.Now(),
		Raw:     m,
	}

	if msg.Content.IsMIME() {
		logf("Parsing MIME body")
		msg.MIME = msg.Content.ParseMIMEBody()
	}

	// FIXME shouldn't be setting Message-ID, its a client thing
	msg.Content.Headers["Message-ID"] = []string{string(id)}
	msg.Content.Headers["Received"] = []string{"from " + m.Helo + " by " + hostname + " (Go-MailHog)\r\n          id " + string(id) + "; " + time.Now().Format(time.RFC1123Z)}
	msg.Content.Headers["Return-Path"] = []string{"<" + m.From + ">"}
	return msg
}

// IsMIME detects a valid MIME header
func (content *Content) IsMIME() bool {
	header, ok := content.Headers["Content-Type"]
	if !ok {
		return false
	}
	return strings.HasPrefix(header[0], "multipart/")
}

// ParseMIMEBody parses SMTP message content into multiple MIME parts
func (content *Content) ParseMIMEBody() *MIMEBody {
	var parts []*Content

	if hdr, ok := content.Headers["Content-Type"]; ok {
		if len(hdr) > 0 {
			boundary := extractBoundary(hdr[0])
			var p []string
			if len(boundary) > 0 {
				p = strings.Split(content.Body, "--"+boundary)
				logf("Got boundary: %s", boundary)
			} else {
				logf("Boundary not found: %s", hdr[0])
			}

			for _, s := range p {
				if len(s) > 0 {
					part := ContentFromString(strings.Trim(s, "\r\n"))
					if part.IsMIME() {
						logf("Parsing inner MIME body")
						part.MIME = part.ParseMIMEBody()
					}
					parts = append(parts, part)
				}
			}
		}
	}

	return &MIMEBody{
		Parts: parts,
	}
}

// PathFromString parses a forward-path or reverse-path into its parts
func PathFromString(path string) *Path {
	var relays []string
	email := path
	if strings.Contains(path, ":") {
		x := strings.SplitN(path, ":", 2)
		r, e := x[0], x[1]
		email = e
		relays = strings.Split(r, ",")
	}
	mailbox, domain := "", ""
	if strings.Contains(email, "@") {
		x := strings.SplitN(email, "@", 2)
		mailbox, domain = x[0], x[1]
	} else {
		mailbox = email
	}

	return &Path{
		Relays:  relays,
		Mailbox: mailbox,
		Domain:  domain,
		Params:  "", // FIXME?
	}
}

// ContentFromString parses SMTP content into separate headers and body
func ContentFromString(data string) *Content {
	logf("Parsing Content from string: '%s'", data)
	x := strings.SplitN(data, "\r\n\r\n", 2)
	h := make(map[string][]string, 0)

	if len(x) == 2 {
		headers, body := x[0], x[1]
		hdrs := strings.Split(headers, "\r\n")
		var lastHdr = ""
		for _, hdr := range hdrs {
			if lastHdr != "" && (strings.HasPrefix(hdr, " ") || strings.HasPrefix(hdr, "\t")) {
				h[lastHdr][len(h[lastHdr])-1] = h[lastHdr][len(h[lastHdr])-1] + hdr
			} else if strings.Contains(hdr, ": ") {
				y := strings.SplitN(hdr, ": ", 2)
				key, value := y[0], y[1]
				// TODO multiple header fields
				h[key] = []string{value}
				lastHdr = key
			} else {
				logf("Found invalid header: '%s'", hdr)
			}
		}
		return &Content{
			Size:    len(data),
			Headers: h,
			Body:    body,
		}
	}
	return &Content{
		Size:    len(data),
		Headers: h,
		Body:    x[0],
	}
}

// extractBoundary extract boundary string in contentType.
// It returns empty string if no valid boundary found
func extractBoundary(contentType string) string {
	_, params, err := mime.ParseMediaType(contentType)
	if err == nil {
		return params["boundary"]
	}
	return ""
}
<|endoftext|>


<commit_before>package twitch

import (
	"strconv"
	"strings"
	"time"
)

// MessageType different message types possible to receive via IRC
type MessageType int

const (
	// UNSET is the default message type, for whenever a new message type is added by twitch that we don't parse yet
	UNSET MessageType = -1
	// WHISPER private messages
	WHISPER MessageType = 0
	// PRIVMSG standard chat message
	PRIVMSG MessageType = 1
	// CLEARCHAT timeout messages
	CLEARCHAT MessageType = 2
	// ROOMSTATE changes like sub mode
	ROOMSTATE MessageType = 3
	// USERNOTICE messages like subs, resubs, raids, etc
	USERNOTICE MessageType = 4
	// USERSTATE messages
	USERSTATE MessageType = 5
	// NOTICE messages like sub mode, host on
	NOTICE MessageType = 6
)

type channelMessage struct {
	RawMessage
	Channel string
}

type roomMessage struct {
	channelMessage
	RoomID string
}

type chatMessage struct {
	roomMessage
	ID   string
	Time time.Time
}

// Emote twitch emotes
type Emote struct {
	Name  string
	ID    string
	Count int
}

// message is purely for internal use
type message struct {
	RawMessage RawMessage
	Channel    string
	Username   string
}

func parseMessage(line string) *message {
	if !strings.HasPrefix(line, "@") {
		return &message{
			RawMessage: RawMessage{
				Type:    UNSET,
				Raw:     line,
				Message: line,
			},
		}
	}

	split := strings.SplitN(line, " :", 3)
	if len(split) < 3 {
		for i := 0; i < 3-len(split); i++ {
			split = append(split, "")
		}
	}

	rawType, channel, username := parseMiddle(split[1])

	rawMessage := RawMessage{
		Type:    parseMessageType(rawType),
		RawType: rawType,
		Raw:     line,
		Tags:    parseTags(split[0]),
		Message: split[2],
	}

	return &message{
		Channel:    channel,
		Username:   username,
		RawMessage: rawMessage,
	}
}

func parseMiddle(middle string) (string, string, string) {
	var rawType, channel, username string

	for _, v := range strings.SplitN(middle, " ", 3) {
		if strings.Contains(v, "!") {
			username = strings.SplitN(v, "!", 2)[0]
		} else if strings.Contains(v, "#") {
			channel = strings.TrimPrefix(v, "#")
		} else {
			rawType = v
		}
	}

	return rawType, channel, username
}

func parseMessageType(messageType string) MessageType {
	switch messageType {
	case "PRIVMSG":
		return PRIVMSG
	case "WHISPER":
		return WHISPER
	case "CLEARCHAT":
		return CLEARCHAT
	case "NOTICE":
		return NOTICE
	case "ROOMSTATE":
		return ROOMSTATE
	case "USERSTATE":
		return USERSTATE
	case "USERNOTICE":
		return USERNOTICE
	default:
		return UNSET
	}
}

func parseTags(tagsRaw string) map[string]string {
	tags := make(map[string]string)

	tagsRaw = strings.TrimPrefix(tagsRaw, "@")
	for _, v := range strings.Split(tagsRaw, ";") {
		tag := strings.SplitN(v, "=", 2)

		var value string
		if len(tag) > 1 {
			value = tag[1]
		}

		tags[tag[0]] = value
	}
	return tags
}

func (m *message) parsePRIVMSGMessage() (*User, *PRIVMSGMessage) {
	privateMessage := PRIVMSGMessage{
		chatMessage: *m.parseChatMessage(),
		Emotes:      m.parseEmotes(),
	}

	text := privateMessage.Message
	if strings.HasPrefix(text, "\u0001ACTION") && strings.HasSuffix(text, "\u0001") {
		privateMessage.Action = true
		privateMessage.Message = text[8 : len(text)-1]
	}

	rawBits, ok := m.RawMessage.Tags["bits"]
	if !ok {
		return m.parseUser(), &privateMessage
	}

	bits, _ := strconv.Atoi(rawBits)
	privateMessage.Bits = bits
	return m.parseUser(), &privateMessage
}

func (m *message) parseCLEARCHATMessage() *CLEARCHATMessage {
	clearchatMessage := CLEARCHATMessage{
		chatMessage:  *m.parseChatMessage(),
		TargetUserID: m.RawMessage.Tags["target-user-id"],
	}

	clearchatMessage.TargetUsername = clearchatMessage.Message
	clearchatMessage.Message = ""

	rawBanDuration, ok := m.RawMessage.Tags["ban-duration"]
	if !ok {
		return &clearchatMessage
	}

	banDuration, _ := strconv.Atoi(rawBanDuration)
	clearchatMessage.BanDuration = banDuration
	return &clearchatMessage
}

func (m *message) parseUser() *User {
	user := User{
		ID:          m.RawMessage.Tags["user-id"],
		Name:        m.Username,
		DisplayName: m.RawMessage.Tags["display-name"],
		Color:       m.RawMessage.Tags["color"],
		Badges:      m.parseBadges(),
	}

	// USERSTATE doesn't contain a Username, but it does have a display-name tag.
	if user.Name == "" {
		user.Name = strings.ToLower(user.DisplayName)
	}

	return &user
}
func (m *message) parseBadges() map[string]int {
	badges := make(map[string]int)

	rawBadges, ok := m.RawMessage.Tags["badges"]
	if !ok {
		return badges
	}

	for _, v := range strings.Split(rawBadges, ",") {
		badge := strings.SplitN(v, "/", 2)
		if len(badge) < 2 {
			continue
		}

		badges[badge[0]], _ = strconv.Atoi(badge[1])
	}

	return badges
}

func (m *message) parseChatMessage() *chatMessage {
	chatMessage := chatMessage{
		roomMessage: *m.parseRoomMessage(),
		ID:          m.RawMessage.Tags["id"],
	}

	i, err := strconv.ParseInt(m.RawMessage.Tags["tmi-sent-ts"], 10, 64)
	if err != nil {
		return &chatMessage
	}

	chatMessage.Time = time.Unix(0, int64(i*1e6))
	return &chatMessage
}

func (m *message) parseRoomMessage() *roomMessage {
	return &roomMessage{
		channelMessage: *m.parseChannelMessage(),
		RoomID:         m.RawMessage.Tags["room-id"],
	}
}

func (m *message) parseChannelMessage() *channelMessage {
	return &channelMessage{
		RawMessage: m.RawMessage,
		Channel:    m.Channel,
	}
}

func (m *message) parseEmotes() []*Emote {
	var emotes []*Emote

	rawEmotes := m.RawMessage.Tags["emotes"]
	if rawEmotes == "" {
		return emotes
	}

	runes := []rune(m.RawMessage.Message)

	for _, v := range strings.Split(rawEmotes, "/") {
		split := strings.SplitN(v, ":", 2)
		pos := strings.SplitN(split[1], ",", 2)
		indexPair := strings.SplitN(pos[0], "-", 2)
		firstIndex, _ := strconv.Atoi(indexPair[0])
		lastIndex, _ := strconv.Atoi(indexPair[1])

		e := &Emote{
			Name:  string(runes[firstIndex:lastIndex]),
			ID:    split[0],
			Count: strings.Count(split[1], ",") + 1,
		}

		emotes = append(emotes, e)
	}

	return emotes
}

func parseJoinPart(text string) (string, string) {
	username := strings.Split(text, "!")
	channel := strings.Split(username[1], "#")
	return strings.Trim(channel[1], " "), strings.Trim(username[0], " :")
}

func parseNames(text string) (string, []string) {
	lines := strings.Split(text, ":")
	channelDirty := strings.Split(lines[1], "#")
	channel := strings.Trim(channelDirty[1], " ")
	users := strings.Split(lines[2], " ")

	return channel, users
}
<commit_msg>Fix parseMiddle incorrectly assigning rawType<commit_after>package twitch

import (
	"strconv"
	"strings"
	"time"
)

// MessageType different message types possible to receive via IRC
type MessageType int

const (
	// UNSET is the default message type, for whenever a new message type is added by twitch that we don't parse yet
	UNSET MessageType = -1
	// WHISPER private messages
	WHISPER MessageType = 0
	// PRIVMSG standard chat message
	PRIVMSG MessageType = 1
	// CLEARCHAT timeout messages
	CLEARCHAT MessageType = 2
	// ROOMSTATE changes like sub mode
	ROOMSTATE MessageType = 3
	// USERNOTICE messages like subs, resubs, raids, etc
	USERNOTICE MessageType = 4
	// USERSTATE messages
	USERSTATE MessageType = 5
	// NOTICE messages like sub mode, host on
	NOTICE MessageType = 6
)

type channelMessage struct {
	RawMessage
	Channel string
}

type roomMessage struct {
	channelMessage
	RoomID string
}

type chatMessage struct {
	roomMessage
	ID   string
	Time time.Time
}

// Emote twitch emotes
type Emote struct {
	Name  string
	ID    string
	Count int
}

// message is purely for internal use
type message struct {
	RawMessage RawMessage
	Channel    string
	Username   string
}

func parseMessage(line string) *message {
	if !strings.HasPrefix(line, "@") {
		return &message{
			RawMessage: RawMessage{
				Type:    UNSET,
				Raw:     line,
				Message: line,
			},
		}
	}

	split := strings.SplitN(line, " :", 3)
	if len(split) < 3 {
		for i := 0; i < 3-len(split); i++ {
			split = append(split, "")
		}
	}

	rawType, channel, username := parseMiddle(split[1])

	rawMessage := RawMessage{
		Type:    parseMessageType(rawType),
		RawType: rawType,
		Raw:     line,
		Tags:    parseTags(split[0]),
		Message: split[2],
	}

	return &message{
		Channel:    channel,
		Username:   username,
		RawMessage: rawMessage,
	}
}

func parseMiddle(middle string) (string, string, string) {
	var rawType, channel, username string

	for i, v := range strings.SplitN(middle, " ", 3) {
		switch {
		case i == 1:
			rawType = v
		case strings.Contains(v, "!"):
			username = strings.SplitN(v, "!", 2)[0]
		case strings.Contains(v, "#"):
			channel = strings.TrimPrefix(v, "#")
		}
	}

	return rawType, channel, username
}

func parseMessageType(messageType string) MessageType {
	switch messageType {
	case "PRIVMSG":
		return PRIVMSG
	case "WHISPER":
		return WHISPER
	case "CLEARCHAT":
		return CLEARCHAT
	case "NOTICE":
		return NOTICE
	case "ROOMSTATE":
		return ROOMSTATE
	case "USERSTATE":
		return USERSTATE
	case "USERNOTICE":
		return USERNOTICE
	default:
		return UNSET
	}
}

func parseTags(tagsRaw string) map[string]string {
	tags := make(map[string]string)

	tagsRaw = strings.TrimPrefix(tagsRaw, "@")
	for _, v := range strings.Split(tagsRaw, ";") {
		tag := strings.SplitN(v, "=", 2)

		var value string
		if len(tag) > 1 {
			value = tag[1]
		}

		tags[tag[0]] = value
	}
	return tags
}

func (m *message) parsePRIVMSGMessage() (*User, *PRIVMSGMessage) {
	privateMessage := PRIVMSGMessage{
		chatMessage: *m.parseChatMessage(),
		Emotes:      m.parseEmotes(),
	}

	text := privateMessage.Message
	if strings.HasPrefix(text, "\u0001ACTION") && strings.HasSuffix(text, "\u0001") {
		privateMessage.Action = true
		privateMessage.Message = text[8 : len(text)-1]
	}

	rawBits, ok := m.RawMessage.Tags["bits"]
	if !ok {
		return m.parseUser(), &privateMessage
	}

	bits, _ := strconv.Atoi(rawBits)
	privateMessage.Bits = bits
	return m.parseUser(), &privateMessage
}

func (m *message) parseCLEARCHATMessage() *CLEARCHATMessage {
	clearchatMessage := CLEARCHATMessage{
		chatMessage:  *m.parseChatMessage(),
		TargetUserID: m.RawMessage.Tags["target-user-id"],
	}

	clearchatMessage.TargetUsername = clearchatMessage.Message
	clearchatMessage.Message = ""

	rawBanDuration, ok := m.RawMessage.Tags["ban-duration"]
	if !ok {
		return &clearchatMessage
	}

	banDuration, _ := strconv.Atoi(rawBanDuration)
	clearchatMessage.BanDuration = banDuration
	return &clearchatMessage
}

func (m *message) parseUser() *User {
	user := User{
		ID:          m.RawMessage.Tags["user-id"],
		Name:        m.Username,
		DisplayName: m.RawMessage.Tags["display-name"],
		Color:       m.RawMessage.Tags["color"],
		Badges:      m.parseBadges(),
	}

	// USERSTATE doesn't contain a Username, but it does have a display-name tag.
	if user.Name == "" {
		user.Name = strings.ToLower(user.DisplayName)
	}

	return &user
}
func (m *message) parseBadges() map[string]int {
	badges := make(map[string]int)

	rawBadges, ok := m.RawMessage.Tags["badges"]
	if !ok {
		return badges
	}

	for _, v := range strings.Split(rawBadges, ",") {
		badge := strings.SplitN(v, "/", 2)
		if len(badge) < 2 {
			continue
		}

		badges[badge[0]], _ = strconv.Atoi(badge[1])
	}

	return badges
}

func (m *message) parseChatMessage() *chatMessage {
	chatMessage := chatMessage{
		roomMessage: *m.parseRoomMessage(),
		ID:          m.RawMessage.Tags["id"],
	}

	i, err := strconv.ParseInt(m.RawMessage.Tags["tmi-sent-ts"], 10, 64)
	if err != nil {
		return &chatMessage
	}

	chatMessage.Time = time.Unix(0, int64(i*1e6))
	return &chatMessage
}

func (m *message) parseRoomMessage() *roomMessage {
	return &roomMessage{
		channelMessage: *m.parseChannelMessage(),
		RoomID:         m.RawMessage.Tags["room-id"],
	}
}

func (m *message) parseChannelMessage() *channelMessage {
	return &channelMessage{
		RawMessage: m.RawMessage,
		Channel:    m.Channel,
	}
}

func (m *message) parseEmotes() []*Emote {
	var emotes []*Emote

	rawEmotes := m.RawMessage.Tags["emotes"]
	if rawEmotes == "" {
		return emotes
	}

	runes := []rune(m.RawMessage.Message)

	for _, v := range strings.Split(rawEmotes, "/") {
		split := strings.SplitN(v, ":", 2)
		pos := strings.SplitN(split[1], ",", 2)
		indexPair := strings.SplitN(pos[0], "-", 2)
		firstIndex, _ := strconv.Atoi(indexPair[0])
		lastIndex, _ := strconv.Atoi(indexPair[1])

		e := &Emote{
			Name:  string(runes[firstIndex:lastIndex]),
			ID:    split[0],
			Count: strings.Count(split[1], ",") + 1,
		}

		emotes = append(emotes, e)
	}

	return emotes
}

func parseJoinPart(text string) (string, string) {
	username := strings.Split(text, "!")
	channel := strings.Split(username[1], "#")
	return strings.Trim(channel[1], " "), strings.Trim(username[0], " :")
}

func parseNames(text string) (string, []string) {
	lines := strings.Split(text, ":")
	channelDirty := strings.Split(lines[1], "#")
	channel := strings.Trim(channelDirty[1], " ")
	users := strings.Split(lines[2], " ")

	return channel, users
}
<|endoftext|>