package main

// a word-count application "plugin" for MapReduce.
// go build -buildmode=plugin wc.go

import "../mr"
import "unicode"
import "strings"
import "strconv"

// The map function is called once for each file of input. The first
// argument is the name of the input file, and the second is the
// file's complete contents. You should ignore the input file name,
// and look only at the contents argument. The return value is a slice
// of key/value pairs.
func Map(filename string, contents string) []mr.KeyValue {
	// function to detect word separators.
	ff := func(r rune) bool { return !unicode.IsLetter(r) }

	// split contents into an array of words.
	words := strings.FieldsFunc(contents, ff)

	kva := []mr.KeyValue{}
	for _, w := range words {
		kv := mr.KeyValue{w, "1"}
		kva = append(kva, kv)
	}
	return kva
}

// The reduce function is called once for each key generated by the
// map tasks, with a list of all the values created for that key by
// any map task.
func Reduce(key string, values []string) string {
	// return the number of occurrences of this word.
	return strconv.Itoa(len(values))
}

package main

// simple sequential MapReduce.
// go run mrsequential.go wc.so pg*.txt

import "fmt"
import "../mr"
import "plugin"
import "os"
import "log"
import "io/ioutil"
import "sort"

// for sorting by key.
type ByKey []mr.KeyValue

// for sorting by key.
func (a ByKey) Len() int           { return len(a) }
func (a ByKey) Swap(i, j int)      { a[i], a[j] = a[j], a[i] }
func (a ByKey) Less(i, j int) bool { return a[i].Key < a[j].Key }

func main() {
	if len(os.Args) < 3 {
		fmt.Fprintf(os.Stderr, "Usage: mrsequential xxx.so inputfiles...\n")
		os.Exit(1)
	}

	mapf, reducef := loadPlugin(os.Args[1])

	// read each input file,
	// pass it to Map,
	// accumulate the intermediate Map output.
	intermediate := []mr.KeyValue{}
	for _, filename := range os.Args[2:] {
		file, err := os.Open(filename)
		if err != nil {
			log.Fatalf("cannot open %v", filename)
		}
		content, err := ioutil.ReadAll(file)
		if err != nil {
			log.Fatalf("cannot read %v", filename)
		}
		file.Close()
		kva := mapf(filename, string(content))
		intermediate = append(intermediate, kva...)
	}

	// a big difference from real MapReduce is that all the
	// intermediate data is in one place, intermediate[],
	// rather than being partitioned into NxM buckets.

	sort.Sort(ByKey(intermediate))

	oname := "mr-out-0"
	ofile, _ := os.Create(oname)

	// call Reduce on each distinct key in intermediate[],
	// and print the result to mr-out-0.
	i := 0
	for i < len(intermediate) {
		j := i + 1
		for j < len(intermediate) && intermediate[j].Key == intermediate[i].Key {
			j++
		}
		values := []string{}
		for k := i; k < j; k++ {
			values = append(values, intermediate[k].Value)
		}
		output := reducef(intermediate[i].Key, values)

		// this is the correct format for each line of Reduce output.
		fmt.Fprintf(ofile, "%v %v\n", intermediate[i].Key, output)

		i = j
	}

	ofile.Close()
}

// load the application Map and Reduce functions
// from a plugin file, e.g. ../mrapps/wc.so
func loadPlugin(filename string) (func(string, string) []mr.KeyValue, func(string, []string) string) {
	p, err := plugin.Open(filename)
	if err != nil {
		log.Fatalf("cannot load plugin %v", filename)
	}
	xmapf, err := p.Lookup("Map")
	if err != nil {
		log.Fatalf("cannot find Map in %v", filename)
	}
	mapf := xmapf.(func(string, string) []mr.KeyValue)
	xreducef, err := p.Lookup("Reduce")
	if err != nil {
		log.Fatalf("cannot find Reduce in %v", filename)
	}
	reducef := xreducef.(func(string, []string) string)

	return mapf, reducef
}

// TaskTypeOpt The MapReduce task type
type TaskTypeOpt string

const (
	TaskTypeMap      TaskTypeOpt = "Map"
	TaskTypeReduce   TaskTypeOpt = "Reduce"
	TaskTypeFinished TaskTypeOpt = "Finished"
)

// TaskInfo The MapReduce task
type TaskInfo struct {
	Id       string
	Type     TaskTypeOpt
	Index    int
	File     string
	WorkerId string
	Deadline time.Time
}

// AckAndQueryNewTaskRequest RPC request for Workers to query a task after finished previously task
type AckAndQueryNewTaskRequest struct {
	// The finished previous task index (if it has finished task)
	PreviousTaskIndex int
	TaskType          TaskTypeOpt
	WorkerId          string
}

// AckAndQueryNewTaskResponse RPC response for Workers to query a task
type AckAndQueryNewTaskResponse struct {
	// The task id(filename) for Map or Reduce to yield results(if there has)
	Task            *TaskInfo
	MapWorkerCnt    int
	ReduceWorkerCnt int
}

type Master struct {
	// Use lock to avoid data race
	lock sync.RWMutex

	// The phase of all tasks
	status TaskTypeOpt

	// The count of the workers
	mapCnt    int
	reduceCnt int

	// All tasks
	tasks map[string]*TaskInfo
	// All ongoing tasks
	availableTasks chan *TaskInfo
}

// MakeMaster create a Master.
// main/mrmaster.go calls this function.
// nReduce is the number of reduce tasks to use.
func MakeMaster(files []string, nReduce int) *Master {

	// Step 1: Create master
	m := &Master{
		status:         TaskTypeMap,
		mapCnt:         len(files),
		reduceCnt:      nReduce,
		tasks:          make(map[string]*TaskInfo),
		availableTasks: make(chan *TaskInfo, max(len(files), nReduce)),
	}

	// Step 2: Store the data
	for idx, file := range files {
		task := &TaskInfo{
			Id:    generateTaskId(TaskTypeMap, idx),
			Type:  TaskTypeMap,
			Index: idx,
			File:  file,
		}
		m.tasks[task.Id] = task  // Store the tasks
		m.availableTasks <- task // Send the task to the channel
	}

	// Step 3: Start master server
	m.server()
	infof("master server started: %v", m)

	// Step 4: Start workers heartbeats checker
	go func() {
		for {
			time.Sleep(500 * time.Millisecond)
			m.checkWorkers()
		}
	}()

	return m
}

// generateTaskId Generate TaskId for the given task
//  The id follows the format: "taskType-taskIndex"
func generateTaskId(taskType TaskTypeOpt, index int) string {
	return fmt.Sprintf("%s-%d", taskType, index)
}

// start a thread that listens for RPCs from worker.go
func (m *Master) server() {
	_ = rpc.Register(m)
	rpc.HandleHTTP()
	//l, e := net.Listen("tcp", ":1234")
	sockname := masterSock()
	_ = os.Remove(sockname)
	l, e := net.Listen("unix", sockname)
	if e != nil {
		log.Fatal("listen error:", e)
	}
	go http.Serve(l, nil)
}

m := mr.MakeMaster(os.Args[1:], 10)
for m.Done() == false {
    time.Sleep(time.Second)
}

func (m *Master) AckAndQueryNewTask(req *AckAndQueryNewTaskRequest,
	resp *AckAndQueryNewTaskResponse) error {

	// Step 1: Mark previous task finished if necessary
	if req.TaskType != "" {
		err := m.handlePreviousTask(req)
		if err != nil {
			errorf("handlePreviousTask err: %v", err)
			return err
		}
	}

	// Step 2: Get the next task
	task, ok := <-m.availableTasks
	if !ok { // Channel closed: no available tasks
		m.lock.RLock()
		defer m.lock.RUnlock()
		retTask := &TaskInfo{}
		if m.status == TaskTypeFinished {
			retTask.Type = TaskTypeFinished
		}
		resp.Task = retTask
		return nil
	}

	// Step 3: Assign the task to the worker
	m.lock.Lock()
	defer m.lock.Unlock()
	infof("Assign task %v to worker %s", task, req.WorkerId)
	task.WorkerId = req.WorkerId
	task.Deadline = time.Now().Add(10 * time.Second)
	m.tasks[generateTaskId(task.Type, task.Index)] = task

	// Step 4: Handle response
	resp.Task = &TaskInfo{
		Id:       task.Id,
		Type:     task.Type,
		Index:    task.Index,
		File:     task.File,
		WorkerId: task.WorkerId,
		Deadline: task.Deadline,
	}
	resp.MapWorkerCnt = m.mapCnt
	resp.ReduceWorkerCnt = m.reduceCnt

	return nil
}

var (
	taskFinishHandlerMap = map[TaskTypeOpt]func(workerId string, taskIdx, reduceCnt int) error{
		TaskTypeMap:    handleFinishedMapTask,
		TaskTypeReduce: handleFinishedReduceTask,
	}
)

// handle the previous finished task
func (m *Master) handlePreviousTask(req *AckAndQueryNewTaskRequest) error {

	previousTaskId := generateTaskId(req.TaskType, req.PreviousTaskIndex)

	m.lock.Lock()
	defer m.lock.Unlock()
	taskInfo, exists := m.tasks[previousTaskId]
	if exists {
		if taskInfo.WorkerId == req.WorkerId { // This task belongs to the worker
			infof("Mark task [%v] finished on worker %s", taskInfo, req.WorkerId)

			// Step 1: Handle the previous finished task
			handler, handlerExists := taskFinishHandlerMap[taskInfo.Type]
			if !handlerExists || handler == nil {
				return fmt.Errorf("handler not found for task: %v", taskInfo)
			}
			err := handler(req.WorkerId, req.PreviousTaskIndex, m.reduceCnt)
			if err != nil {
				errorf("Failed to handle previous task: %v", err)
				return err
			}
			delete(m.tasks, previousTaskId)

			// Step 2: Transit job phase if necessary
			if len(m.tasks) <= 0 {
				m.transit()
			}
			return nil
		} else { // The task is no longer belongs to this worker
			infof("Task %v is no longer belongs to this worker", taskInfo)
			return nil
		}
	} else { // Previous task not found in task map(worker retry ack maybe)!
		warnf("[Warn] Previous task: %v not found in map", taskInfo)
		return nil
	}
}

var (
	taskFinishHandlerMap = map[TaskTypeOpt]func(workerId string, taskIdx, reduceCnt int) error{
		TaskTypeMap:    handleFinishedMapTask,
		TaskTypeReduce: handleFinishedReduceTask,
	}
)

// The map-type task finished handler
func handleFinishedMapTask(workerId string, taskIdx, reduceCnt int) error {
	// Mark the task's temporary file to final file(Rename)
	for reduceIdx := 0; reduceIdx < reduceCnt; reduceIdx++ {
		tmpMapFileName := tmpMapOutFile(workerId, taskIdx, reduceIdx)
		finalMapOutFileName := finalMapOutFile(taskIdx, reduceIdx)
		err := os.Rename(tmpMapFileName, finalMapOutFileName)
		if err != nil {
			errorf("Failed to mark map output file `%s` as final: %e", tmpMapFileName, err)
			return err
		}
	}

	infof("handleFinishedMapTask success: workerId: %s, taskIdx: %d", workerId, taskIdx)

	return nil
}

// The reduce-type task finished handler
func handleFinishedReduceTask(workerId string, taskIdx, _ int) error {
	// Mark the task's temporary file to final file(Rename)
	tmpReduceFileName := tmpReduceOutFile(workerId, taskIdx)
	finalReduceOutFileName := finalReduceOutFile(taskIdx)
	err := os.Rename(tmpReduceFileName, finalReduceOutFileName)
	if err != nil {
		errorf("Failed to mark reduce output file `%s` as final: %v", tmpReduceFileName, err)
		return err
	}

	infof("handleFinishedReduceTask success: workerId: %s, taskIdx: %d, finalReduceOutFileName: %s",
		workerId, taskIdx, finalReduceOutFileName)

	return nil
}

// The temporary file that map-type task yield
func tmpMapOutFile(workerId string, taskIdx, reduceIdx int) string {
	return fmt.Sprintf("mr-map-%s-%d-%d", workerId, taskIdx, reduceIdx)
}

// The final file that map-type task yield(for reduce)
func finalMapOutFile(taskIdx, reduceIdx int) string {
	return fmt.Sprintf("mr-map-%d-%d", taskIdx, reduceIdx)
}

// The temporary file that reduce-type task yield
func tmpReduceOutFile(workerId string, reduceIdx int) string {
	return fmt.Sprintf("mr-reduce-%s-%d", workerId, reduceIdx)
}

// The final file that reduce-type task yield(the MapReduce task yield)
func finalReduceOutFile(taskIndex int) string {
	return fmt.Sprintf("mr-out-%d", taskIndex)
}

// Transit the job phase (from Map to Reduce)
func (m *Master) transit() {

	if m.status == TaskTypeMap {
		// All map-type tasks finished, change to reduce phase
		infof("All map-type tasks finished. Transit to REDUCE stage!")
		m.status = TaskTypeReduce

		// Yield Reduce Tasks
		for reduceIdx := 0; reduceIdx < m.reduceCnt; reduceIdx++ {
			task := &TaskInfo{
				Type:  TaskTypeReduce,
				Index: reduceIdx,
			}
			m.tasks[generateTaskId(task.Type, task.Index)] = task
			m.availableTasks <- task
		}
	} else if m.status == TaskTypeReduce {
		// All reduce-type tasks finished, ready to exit
		infof("All reduce-type tasks finished. Prepare to exit!")
		// Close channel
		close(m.availableTasks)
		// Mark status to TaskTypeFinished for job completion
		m.status = TaskTypeFinished
	}
}

// Check all workers heartbeats
func (m *Master) checkWorkers() {

	m.lock.Lock()
	defer m.lock.Unlock()
	for _, task := range m.tasks {
		if task.WorkerId != "" && time.Now().After(task.Deadline) {
			infof(
				"Found timed-out task：%v, previously running on worker %s. Prepare to re-assign...",
				task, task.WorkerId)
			task.WorkerId = ""
			m.availableTasks <- task
		}
	}
}

// main/mrmaster.go calls Done() periodically to find out
// if the entire job has finished.
func (m *Master) Done() bool {
	m.lock.RLock()
	defer m.lock.RUnlock()
	
    taskPhase := m.status
	infof("current task phase: %s, res task count: %d", taskPhase, len(m.availableTasks))

	// All tasks have finished
	return taskPhase == TaskTypeFinished
}

// main/mrworker.go calls this function.
func Worker(mapFunc func(string, string) []KeyValue,
	reduceFunc func(string, []string) string) {

	// Done Your worker implementation here.
	workerId := generateWorkerId()
	infof("Worker %v started!\n", workerId)

	// Fire the worker to receive tasks
	var previousTaskType TaskTypeOpt
	var previousTaskIndex int
	var taskErr error
	for {
		// Step 1: Query Task & Ack the last task
		req := AckAndQueryNewTaskRequest{
			WorkerId:          workerId,
			PreviousTaskIndex: previousTaskIndex,
			TaskType:          previousTaskType,
		}
		resp := AckAndQueryNewTaskResponse{}
		succeed := call("Master.AckAndQueryNewTask", &req, &resp)
		if !succeed {
			errorf("Failed to call AckAndQueryNewTask, retry 1 second later")
			time.Sleep(time.Second)
			continue
		}
		infof("Call AckAndQueryNewTaskResponse success! req: %v, resp: %v", &req, &resp)

		// Extra Step: Job finished, exit
		if resp.Task.Type == TaskTypeFinished {
			// Job finished, exit
			infof("Received job finish signal from master, exit")
			break
		}

		// Step 2: handle the queried task
		if resp.Task.Type == TaskTypeMap {
			taskErr = handleMapTask(&resp, workerId, mapFunc)
			if taskErr != nil {
				errorf("Failed to handleMapTask: %v, err: %v", resp.Task, taskErr)
				continue
			}
		} else if resp.Task.Type == TaskTypeReduce {
			taskErr = handleReduceTask(&resp, workerId, reduceFunc)
			if taskErr != nil {
				errorf("Failed to handleReduceTask: %v, err: %v", resp.Task, taskErr)
				continue
			}
		} else {
			errorf("No handler to handle task: %v", resp.Task)
			continue
		}

		// Step 3: save finished task info, ack for the next iteration
		previousTaskType = resp.Task.Type
		previousTaskIndex = resp.Task.Index
		infof("Finished task: %v on worker: %s\n", resp.Task, workerId)
	}

	infof("Worker %v finished!\n", workerId)
}

// Use pid as the workerId(standalone type)
func generateWorkerId() string {
	return strconv.Itoa(os.Getpid())
}

// The map-type task handler for worker
func handleMapTask(resp *AckAndQueryNewTaskResponse,
	workerId string, mapFunc func(string, string) []KeyValue) error {

	// Step 1: Read task file
	taskFile, err := os.Open(resp.Task.File)
	if err != nil {
		errorf("Failed to open map input file %s: %v", resp.Task.File, err)
		return err
	}
	content, err := ioutil.ReadAll(taskFile)
	if err != nil {
		errorf("Failed to read map input file %s: %v", resp.Task.File, err)
		return err
	}

	// Step 2: Use MapFunc to yield intermediate results
	// Key: FileName, Value: FileContent
	// Then Hash the key by ihash func, and push result into different buckets
	mapResults := mapFunc(resp.Task.File, string(content))
	hashedKva := make(map[int][]KeyValue)
	for _, kv := range mapResults {
		hashed := ihash(kv.Key) % resp.ReduceWorkerCnt
		hashedKva[hashed] = append(hashedKva[hashed], kv)
	}

	// Step 3: Writes all intermediate results into intermediate files
	for idx := 0; idx < resp.ReduceWorkerCnt; idx++ {
		tmpFileName := tmpMapOutFile(workerId, resp.Task.Index, idx)
		intermediateFile, _ := os.Create(tmpFileName)
		for _, kv := range hashedKva[idx] {
			// The intermediate file format is: ${key}\t${value}\n
			_, err = fmt.Fprintf(intermediateFile, "%v\t%v\n", kv.Key, kv.Value)
			if err != nil {
				errorf("Write intermediate file: %s failed, err: %v", intermediateFile, err)
				intermediateFile.Close()
				return err
			}
		}
		intermediateFile.Close()
	}

	infof("Worker[%s] writes intermediate files success, task: %v", workerId, resp.Task)

	return nil
}

// use ihash(key) % NReduce to choose the reducer
// task number for each KeyValue emitted by Map.
func ihash(key string) int {
	h := fnv.New32a()
	_, _ = h.Write([]byte(key))
	return int(h.Sum32() & 0x7fffffff)
}

// The reduce-type task handler for worker
func handleReduceTask(resp *AckAndQueryNewTaskResponse,
	workerId string, reduceFunc func(string, []string) string) error {

	// Step 1: Read the corresponding file
	var lines []string
	for mi := 0; mi < resp.MapWorkerCnt; mi++ {
		inputFile := finalMapOutFile(mi, resp.Task.Index)
		file, err := os.Open(inputFile)
		if err != nil {
			errorf("Failed to open map output file %s: %v", inputFile, err)
			return err
		}
		content, err := ioutil.ReadAll(file)
		if err != nil {
			errorf("Failed to read map output file %s: %v", inputFile, err)
			return err
		}
		lines = append(lines, strings.Split(string(content), "\n")...)
	}

	// Step 2: Format the lines
	var mapResults []KeyValue
	for _, line := range lines {
		if strings.TrimSpace(line) == "" {
			continue
		}
		parts := strings.Split(line, "\t")
		mapResults = append(mapResults, KeyValue{
			Key:   parts[0],
			Value: parts[1],
		})
	}

	// Step 3: Sort the results
	sort.Sort(ByKey(mapResults))

	// Step 4: Write the results(Just as the mrsequential.go do!)
	tmpFileName := tmpReduceOutFile(workerId, resp.Task.Index)
	tmpFile, _ := os.Create(tmpFileName)
	defer tmpFile.Close()

	// Step 5: Call Reduce on each distinct key in mapResults[], and write the intermediate result.
	i := 0
	for i < len(mapResults) {
		// Step 5.1: Find the same key in mapResults
		j := i + 1
		for j < len(mapResults) && mapResults[j].Key == mapResults[i].Key {
			j++
		}
		var values []string
		for k := i; k < j; k++ {
			values = append(values, mapResults[k].Value)
		}

		// Step 5.2: Call reduceFunc
		output := reduceFunc(mapResults[i].Key, values)

		// Step 5.3: Save the yield reduce results to intermediate files.
		_, _ = fmt.Fprintf(tmpFile, "%v %v\n", mapResults[i].Key, output)

		i = j
	}

	return nil
}