// Copyright 2017 Google Inc. All rights reserved.
//
// Licensed under the Apache License, Version 2.0 (the "License");
// you may not use this file except in compliance with the License.
// You may obtain a copy of the License at
//
//     http://www.apache.org/licenses/LICENSE-2.0
//
// Unless required by applicable law or agreed to in writing, software
// distributed under the License is distributed on an "AS IS" BASIS,
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
// See the License for the specific language governing permissions and
// limitations under the License.

// This package implements a trace file writer, whose files can be opened in
// chrome://tracing.
//
// It implements the JSON Array Format defined here:
// https://docs.google.com/document/d/1CvAClvFfyA5R-PhYUmn5OOQtYMH4h6I0nSsKchNAySU/edit
package tracer

import (
 "bytes"
 "compress/gzip"
 "encoding/json"
 "fmt"
 "io"
 "os"
 "strings"
 "sync"
 "time"

 "android/soong/ui/logger"
 "android/soong/ui/status"
)

type Thread uint64

const (
 MainThread     = Thread(iota)
 MaxInitThreads = Thread(iota)
)

type Tracer interface {
 Begin(name string, thread Thread)
 End(thread Thread)
 Complete(name string, thread Thread, begin, end uint64)

 CountersAtTime(name string, thread Thread, time uint64, counters []Counter)

 ImportMicrofactoryLog(filename string)

 StatusTracer() status.StatusOutput

 NewThread(name string) Thread
}

type tracerImpl struct {
 lock sync.Mutex
 log  logger.Logger

 buf  bytes.Buffer
 file *os.File
 w    io.WriteCloser

 firstEvent bool
 nextTid    uint64
}

var _ Tracer = &tracerImpl{}

type viewerEvent struct {
 Name  string      `json:"name,omitempty"`
 Phase string      `json:"ph"`
 Scope string      `json:"s,omitempty"`
 Time  uint64      `json:"ts"`
 Dur   uint64      `json:"dur,omitempty"`
 Pid   uint64      `json:"pid"`
 Tid   uint64      `json:"tid"`
 ID    uint64      `json:"id,omitempty"`
 Arg   interface{} `json:"args,omitempty"`
}

type nameArg struct {
 Name string `json:"name"`
}

type nopCloser struct{ io.Writer }

func (nopCloser) Close() error { return nil }

// New creates a new Tracer, storing log in order to log errors later.
// Events are buffered in memory until SetOutput is called.
func New(log logger.Logger) *tracerImpl {
 ret := &tracerImpl{
  log: log,

  firstEvent: true,
  nextTid:    uint64(MaxInitThreads),
 }
 ret.startBuffer()

 return ret
}

func (t *tracerImpl) startBuffer() {
 t.w = nopCloser{&t.buf}
 fmt.Fprintln(t.w, "[")

 t.defineThread(MainThread, "main")
}

func (t *tracerImpl) close() {
 if t.file != nil {
  fmt.Fprintln(t.w, "]")

  if err := t.w.Close(); err != nil {
   t.log.Println("Error closing trace writer:", err)
  }

  if err := t.file.Close(); err != nil {
   t.log.Println("Error closing trace file:", err)
  }
  t.file = nil
  t.startBuffer()
 }
}

// SetOutput creates the output file (rotating old files).
func (t *tracerImpl) SetOutput(filename string) {
 t.lock.Lock()
 defer t.lock.Unlock()

 t.close()

 // chrome://tracing requires that compressed trace files end in .gz
 if !strings.HasSuffix(filename, ".gz") {
  filename += ".gz"
 }

 f, err := logger.CreateFileWithRotation(filename, 5)
 if err != nil {
  t.log.Println("Failed to create trace file:", err)
  return
 }
 // Save the file, since closing the gzip Writer doesn't close the
 // underlying file.
 t.file = f
 t.w = gzip.NewWriter(f)

 // Write out everything that happened since the start
 if _, err := io.Copy(t.w, &t.buf); err != nil {
  t.log.Println("Failed to write trace buffer to file:", err)
 }
 t.buf = bytes.Buffer{}
}

// Close closes the output file. Any future events will be buffered until the
// next call to SetOutput.
func (t *tracerImpl) Close() {
 t.lock.Lock()
 defer t.lock.Unlock()

 t.close()
}

func (t *tracerImpl) writeEvent(event *viewerEvent) {
 t.lock.Lock()
 defer t.lock.Unlock()

 t.writeEventLocked(event)
}

func (t *tracerImpl) writeEventLocked(event *viewerEvent) {
 bytes, err := json.Marshal(event)
 if err != nil {
  t.log.Println("Failed to marshal event:", err)
  t.log.Verbosef("Event: %#v", event)
  return
 }

 if !t.firstEvent {
  fmt.Fprintln(t.w, ",")
 } else {
  t.firstEvent = false
 }

 if _, err = t.w.Write(bytes); err != nil {
  t.log.Println("Trace write error:", err)
 }
}

func (t *tracerImpl) defineThread(thread Thread, name string) {
 t.writeEventLocked(&viewerEvent{
  Name:  "thread_name",
  Phase: "M",
  Pid:   0,
  Tid:   uint64(thread),
  Arg: &nameArg{
   Name: name,
  },
 })
}

// NewThread returns a new Thread with an unused tid, writing the name out to
// the trace file.
func (t *tracerImpl) NewThread(name string) Thread {
 t.lock.Lock()
 defer t.lock.Unlock()

 ret := Thread(t.nextTid)
 t.nextTid += 1

 t.defineThread(ret, name)
 return ret
}

// Begin starts a new Duration Event. More than one Duration Event may be active
// at a time on each Thread, but they're nested.
func (t *tracerImpl) Begin(name string, thread Thread) {
 t.writeEvent(&viewerEvent{
  Name:  name,
  Phase: "B",
  Time:  uint64(time.Now().UnixNano()) / 1000,
  Pid:   0,
  Tid:   uint64(thread),
 })
}

// End finishes the most recent active Duration Event on the thread.
func (t *tracerImpl) End(thread Thread) {
 t.writeEvent(&viewerEvent{
  Phase: "E",
  Time:  uint64(time.Now().UnixNano()) / 1000,
  Pid:   0,
  Tid:   uint64(thread),
 })
}

// Complete writes a Complete Event, which are like Duration Events, but include
// a begin and end timestamp in the same event.
func (t *tracerImpl) Complete(name string, thread Thread, begin, end uint64) {
 t.writeEvent(&viewerEvent{
  Name:  name,
  Phase: "X",
  Time:  begin / 1000,
  Dur:   (end - begin) / 1000,
  Pid:   0,
  Tid:   uint64(thread),
 })
}

type Counter struct {
 Name  string
 Value int64
}

type countersMarshaller []Counter

var _ json.Marshaler = countersMarshaller(nil)

func (counters countersMarshaller) MarshalJSON() ([]byte, error) {
 // This produces similar output to a map[string]int64, but maintains the order of the slice.
 buf := bytes.Buffer{}
 buf.WriteRune('{')
 for i, counter := range counters {
  name, err := json.Marshal(counter.Name)
  if err != nil {
   return nil, err
  }
  buf.Write(name)
  buf.WriteByte(':')
  value, err := json.Marshal(counter.Value)
  if err != nil {
   return nil, err
  }
  buf.Write(value)
  if i != len(counters)-1 {
   buf.WriteRune(',')
  }
 }
 buf.WriteRune('}')
 return buf.Bytes(), nil
}

// CountersAtTime writes a Counter event at the given timestamp in nanoseconds.
func (t *tracerImpl) CountersAtTime(name string, thread Thread, time uint64, counters []Counter) {
 t.writeEvent(&viewerEvent{
  Name:  name,
  Phase: "C",
  Time:  time / 1000,
  Pid:   0,
  Tid:   uint64(thread),
  Arg:   countersMarshaller(counters),
 })
}