type WorkerArgs struct {
	MapTaskNumber    int // finished map task number
	ReduceTaskNumber int // finished reduce task number
}

type WorkerReply struct {
	TaskType         int    // 0:map task 1:reduce task 2:waiting 3:finished
	NMap             int    // total num of map task
	NReduce          int    // total num of reduce task
	MapTaskNumber    int    // number of map task
	ReduceTaskNumber int  	// number of reduce task 
	Filename         string // filename for worker to map
}

const (
	TaskMap int = iota
	TaskReduce
	TaskWaiting
	TaskFinished
)

type Coordinator struct {
	// Your definitions here.
	nMap             int      // total num of map task
	nReduce          int      // total num of reduce task
	mapFinished      int      // number of finished map task
	reduceFinished   int      // number of finished reduce task
	mapTaskStatus    []int    // status array for map task
	reduceTaskStatus []int    // status array for reduce task
	files            []string // input files
	mu               sync.Mutex
}

const (
	StatusNotAllocated int = iota
	StatusWaiting
	StatusFinished
)

func (c *Coordinator) ReceiveFinishedMap(args *WorkerArgs, reply *WorkerReply) error {
	c.mu.Lock()
	defer c.mu.Unlock()
	c.mapFinished++
	c.mapTaskStatus[args.MapTaskNumber] = StatusFinished
	return nil
}

func (c *Coordinator) ReceiveFinishedReduce(args *WorkerArgs, reply *WorkerReply) error {
	c.mu.Lock()
	defer c.mu.Unlock()
	c.reduceFinished++
	c.reduceTaskStatus[args.ReduceTaskNumber] = StatusFinished
	return nil
}

func (c *Coordinator) AllocateTask(args *WorkerArgs, reply *WorkerReply) error {
	c.mu.Lock()
	defer c.mu.Unlock()
	if c.mapFinished < c.nMap {
		// allocate new map task
		allocate := -1
		for i := 0; i < c.nMap; i++ {
			if c.mapTaskStatus[i] == StatusNotAllocated {
				allocate = i
				break
			}
		}
		if allocate == -1 {
			reply.TaskType = TaskWaiting
		} else {
			reply.NReduce = c.nReduce
			reply.TaskType = TaskMap
			reply.MapTaskNumber = allocate
			reply.Filename = c.files[allocate]
			c.mapTaskStatus[allocate] = StatusWaiting
			// 10s not finished
			go func() {
				select {
				case <-time.After(10 * time.Second):
					c.mu.Lock()
					if c.mapTaskStatus[allocate] == StatusWaiting {
						c.mapTaskStatus[allocate] = StatusNotAllocated
					}
					c.mu.Unlock()
				}
			}()
		}
	} else if c.mapFinished == c.nMap && c.reduceFinished < c.nReduce {
		// all map task finished, allocate new reduce task
		allocate := -1
		for i := 0; i < c.nReduce; i++ {
			if c.reduceTaskStatus[i] == StatusNotAllocated {
				allocate = i
				break
			}
		}
		if allocate == -1 {
			reply.TaskType = TaskWaiting
		} else {
			reply.NMap = c.nMap
			reply.TaskType = TaskReduce
			reply.ReduceTaskNumber = allocate
			c.reduceTaskStatus[allocate] = StatusWaiting
			// 10s not finished
			go func() {
				select {
				case <-time.After(10 * time.Second):
					c.mu.Lock()
					if c.reduceTaskStatus[allocate] == StatusWaiting {
						c.reduceTaskStatus[allocate] = StatusNotAllocated
					}
					c.mu.Unlock()
				}
			}()
		}
	} else {
		// all finished
		reply.TaskType = TaskFinished
	}
	return nil
}

func MakeCoordinator(files []string, nReduce int) *Coordinator {
	c := Coordinator{}

	// Your code here.
	c.files = files
	c.nMap = len(files)
	c.nReduce = nReduce
	c.mapTaskStatus = make([]int, c.nMap)
	c.reduceTaskStatus = make([]int, c.nReduce)

	c.server()
	return &c
}

func (c *Coordinator) Done() bool {
    //ret := false
    
    // Your code here.
    c.mu.Lock()
    defer c.mu.Unlock()
    ret := c.nReduce == c.reduceFinished
    return ret
}

// for sorting by key.
type ByKey []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 Worker(mapf func(string, string) []KeyValue,
	reducef func(string, []string) string) {

	// Your worker implementation here.

	// uncomment to send the Example RPC to the coordinator.
	// CallExample()
	for {
		args := WorkerArgs{}
		reply := WorkerReply{}
		ok := call("Coordinator.AllocateTask", &args, &reply)
		if !ok || reply.TaskType == TaskFinished {
			// coordinator died || job finished
			break
		}
		if reply.TaskType == TaskMap {
			// map task
			intermediate := []KeyValue{}
			file, err := os.Open(reply.Filename)
			if err != nil {
				log.Fatalf("cannot open %v: %v", reply.Filename, err)
			}
			content, err := ioutil.ReadAll(file)
			if err != nil {
				log.Fatalf("cannot read %v: %v", reply.Filename, err)
			}
			file.Close()
			kva := mapf(reply.Filename, string(content))
			intermediate = append(intermediate, kva...)

			// hash into buckets
			buckets := make([][]KeyValue, reply.NReduce)
			for i := range buckets {
				buckets[i] = []KeyValue{}
			}
			for _, kva := range intermediate {
				buckets[ihash(kva.Key)%reply.NReduce] = append(buckets[ihash(kva.Key)%reply.NReduce], kva)
			}

			// write into intermediate files
			for i := range buckets {
				oname := fmt.Sprintf("mr-%d-%d", reply.MapTaskNumber, i)
				ofile, _ := ioutil.TempFile("", oname+"*")
				enc := json.NewEncoder(ofile)
				for _, kva := range buckets[i] {
					err := enc.Encode(&kva)
					if err != nil {
						log.Fatalf("cannot write into %v: %v", oname, err)
					}
				}
				os.Rename(ofile.Name(), oname)
				ofile.Close()
			}

			// call to send finish msg
			finishedArgs := WorkerArgs{reply.MapTaskNumber, -1}
			finishedReply := WorkerReply{}
			call("Coordinator.ReceiveFinishedMap", &finishedArgs, &finishedReply)
		} else if reply.TaskType == TaskReduce {
			// reduce task
			// collect Key-Value from mr-X-Y
			intermediate := []KeyValue{}
			for i := 0; i < reply.NMap; i++ {
				iname := fmt.Sprintf("mr-%d-%d", i, reply.ReduceTaskNumber)
				file, err := os.Open(iname)
				if err != nil {
					log.Fatalf("cannot open %v: %v", iname, err)
				}
				dec := json.NewDecoder(file)
				for {
					var kv KeyValue
					if err := dec.Decode(&kv); err != nil {
						break
					}
					intermediate = append(intermediate, kv)
				}
				file.Close()
			}

			// sort by key
			sort.Sort(ByKey(intermediate))

			// output file
			oname := fmt.Sprintf("mr-out-%d", reply.ReduceTaskNumber)
			ofile, _ := ioutil.TempFile("", 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
			}
			os.Rename(ofile.Name(), oname)
			ofile.Close()

			// delete intermediate files
			for i := 0; i < reply.NMap; i++ {
				iname := fmt.Sprintf("mr-%d-%d", i, reply.ReduceTaskNumber)
				err := os.Remove(iname)
				if err != nil {
					log.Fatalf("cannot delete %v: %v", iname, err)
				}
			}

			// call to send finish msg
			finishedArgs := WorkerArgs{-1, reply.ReduceTaskNumber}
			finishedReply := WorkerReply{}
			call("Coordinator.ReceiveFinishedReduce", &finishedArgs, &finishedReply)
		}
		// else: reply.TaskType == TaskWaiting
		time.Sleep(time.Second)
	}
	return
}