package main

import (
	"code.justin.tv/release/trace/analysis/routes"
	"code.justin.tv/release/trace/api"
	"code.justin.tv/release/trace/pbmsg"
)

// Server for code.justin.tv/release/trace/api

type apiServer struct {
	manager *queryManager
}

func newApiServer(manager *queryManager) *apiServer {
	return &apiServer{
		manager: manager,
	}
}

func (s *apiServer) Firehose(req *api.FirehoseRequest, resp api.Trace_FirehoseServer) error {
	ch, _ := s.manager.Add(req)
	defer s.manager.Remove(req)

	for {
		select {
		case tx, ok := <-ch:
			if !ok {
				return nil
			}
			if err := resp.Send(tx); err != nil {
				return err
			}
		// Stop sending data when the request is canceled
		case <-resp.Context().Done():
			return nil
		}
	}
}

func cleanAPITransaction(tx *api.Transaction) {
	cleanAPICall(tx.GetRoot())
}

func cleanAPICall(call *api.Call) {
	for _, c := range call.GetSubcalls() {
		cleanAPICall(c)
	}

	httpParams := call.GetParams().GetHttp()
	if httpParams != nil {
		path := httpParams.UriPath
		httpParams.UriPath = ""

		svc := call.GetSvc()
		if svc != nil {
			rm := routes.Routes[svc.Name]
			httpParams.Route = rm.LookupTrimSlash(path, httpParams.Method.String())
		}

	}
}

// Server for legacy code.justin.tv/release/trace/pbmsg

type traceServer struct {
	manager *queryManager
}

func newServer(manager *queryManager) *traceServer {
	return &traceServer{
		manager: manager,
	}
}

func (s *traceServer) Firehose(req *pbmsg.FirehoseRequest, resp pbmsg.Trace_FirehoseServer) error {
	apiReq := &api.FirehoseRequest{
		Query:    pbmsg.UpgradeQuery(req.Query),
		Sampling: req.Sampling,
	}
	_, ch := s.manager.Add(apiReq)
	defer s.manager.Remove(apiReq)

	for {
		select {
		case tx, ok := <-ch:
			if !ok {
				return nil
			}
			if err := resp.Send(tx); err != nil {
				return err
			}
		// Stop sending data when the request is canceled
		case <-resp.Context().Done():
			return nil
		}
	}
}
