// ---------- for RequestVote ----------
type RequestVoteArgs struct {
	// Your data here (2A, 2B).
	Term        int // candidate's curTerm
	CandidateId int // candidate requesting vote
}

type RequestVoteReply struct {
	// Your data here (2A).
	Term        int  // currentTerm, for candidate to update itself
	VoteGranted bool // true means candidate received vote
}

type voteParam struct {
	args   *RequestVoteArgs
	reply  *RequestVoteReply
	notify chan struct{}
}

func (rf *Raft) RequestVote(args *RequestVoteArgs, reply *RequestVoteReply) {
	// Your code here (2A, 2B).
	notify := make(chan struct{})
	rf.voteChan <- voteParam{args, reply, notify}
	<-notify
}

func (rf *Raft) sendRequestVote(server int, args *RequestVoteArgs, reply *RequestVoteReply) bool {
	ok := rf.peers[server].Call("Raft.RequestVote", args, reply)
	return ok
}

// ---------- for AppendEntries ----------
type AppendEntriesArgs struct {
	Term     int // leader's curTerm
	LeaderId int // so follower can redirect clients
}

type AppendEntriesReply struct {
	Term    int  // currentTerm, for leader to update itself
	Success bool // true if follower contained entry matching prevLogIndex and prevLogTerm
}

type appendEntryParam struct {
	args   *AppendEntriesArgs
	reply  *AppendEntriesReply
	notify chan struct{}
}

func (rf *Raft) AppendEntries(args *AppendEntriesArgs, reply *AppendEntriesReply) {
	notify := make(chan struct{})
	rf.entryChan <- appendEntryParam{args, reply, notify}
	<-notify
}

func (rf *Raft) sendAppendEntries(server int, args *AppendEntriesArgs, reply *AppendEntriesReply) bool {
	ok := rf.peers[server].Call("Raft.AppendEntries", args, reply)
	return ok
}

// Your data here (2A, 2B, 2C).
// Look at the paper's Figure 2 for a description of what
// state a Raft server must maintain.
curTerm  int    // current curTerm
state    RState // current state
votedFor int    // candidate id that received vote in current curTerm

voteChan  chan voteParam        // channel for vote request
entryChan chan appendEntryParam // channel for entry request

type RState interface {
	Run(tf *Raft)
}

type Follower struct{}
type Candidate struct{}
type Leader struct{}

func (rf *Raft) GetState() (int, bool) {
	rf.mu.Lock()
	defer rf.mu.Unlock()
	return rf.curTerm, reflect.TypeOf(rf.state) == reflect.TypeOf(&Leader{})
}

func (rf *Raft) ticker() {
	for !rf.killed() {
		rf.state.Run(rf)
	}
}

func Make(peers []*labrpc.ClientEnd, me int,
	persister *Persister, applyCh chan ApplyMsg) *Raft {
	rf := &Raft{}
	rf.peers = peers
	rf.persister = persister
	rf.me = me

	rf.curTerm = 0
	rf.state = &Follower{}
	rf.votedFor = -1

	rf.voteChan = make(chan voteParam)
	rf.entryChan = make(chan appendEntryParam)
	rf.readPersist(persister.ReadRaftState())

	// start ticker goroutine to start elections
	go rf.ticker()

	return rf
}

const (
	ReElectLower = 150
	ReElectUpper = 300

	HeartBeatTimeout = 100
)

func randBetween(lower, upper int) int {
	return rand.Intn(upper-lower) + lower
}

func electionTimeout() time.Duration {
	return time.Duration(randBetween(ReElectLower, ReElectUpper)) * time.Millisecond
}

func heartbeatTimeout() time.Duration {
	return time.Duration(HeartBeatTimeout) * time.Millisecond
}

func resetTimer(timer *time.Timer, timeout time.Duration) {
	if !timer.Stop() {
		<-timer.C
	}
	timer.Reset(timeout)
}

func (f *Follower) Run(rf *Raft) {
	timer := time.NewTimer(electionTimeout())
	for {
		select {
		case <-timer.C:
		// a 处理超时
		case vote := <-rf.voteChan:
		// b 处理 RequestVote 请求
		case entry := <-rf.entryChan:
		// c 处理 AppendEntries 请求
	}
}

rf.mu.Lock()
rf.state = &Candidate{}
rf.votedFor = rf.me
rf.mu.Unlock()
return

vote.reply.Term = vote.args.Term
vote.reply.VoteGranted = false

vote.args.CandidateId)
rf.mu.Lock()
if vote.args.Term < rf.curTerm {
	rf.mu.Unlock()
	vote.notify <- struct{}{}
	continue
}

// grant vote, update curTerm
if vote.args.Term > rf.curTerm ||
	(vote.args.Term == rf.curTerm && rf.votedFor == vote.args.CandidateId) {
	rf.curTerm = vote.args.Term
	rf.votedFor = vote.args.CandidateId
	vote.reply.Term = rf.curTerm
	vote.reply.VoteGranted = true
	// reset timer
	resetTimer(timer, electionTimeout())
}
rf.mu.Unlock()
vote.notify <- struct{}{}

entry.args.LeaderId)
rf.mu.Lock()
if entry.args.Term < rf.curTerm {
	// stale curTerm, do nothing
	entry.reply.Term = rf.curTerm
	entry.reply.Success = false
} else if entry.args.Term > rf.curTerm {
	// larger curTerm, may this server is out of date
	rf.curTerm = entry.args.Term
	rf.votedFor = entry.args.LeaderId
	entry.reply.Term = rf.curTerm
	entry.reply.Success = true
	resetTimer(timer, electionTimeout())
} else {
	// same curTerm, reset timer
	entry.reply.Term = rf.curTerm
	entry.reply.Success = true
	resetTimer(timer, electionTimeout())
}
rf.mu.Unlock()
entry.notify <- struct{}{}

func (c *Candidate) Run(rf *Raft) {
	rf.mu.Lock()
	rf.curTerm++

    // TODO: send RequestVote RPC
	rf.mu.Unlock()

	timer := time.NewTimer(electionTimeout())
	for {
		select {
		case <-timer.C:
		// d 超时
		case reply := <-replyChan:
		// g RequestVote 回答
		case vote := <-rf.voteChan:
		// e RequestVote 请求
		case entry := <-rf.entryChan:
		// f AppendEntries 请求
		}
	}
}

return

rf.mu.Lock()
if vote.args.Term > rf.curTerm {
	rf.curTerm = vote.args.Term
	rf.votedFor = vote.args.CandidateId
	rf.state = &Follower{}
	rf.mu.Unlock()
	vote.reply.VoteGranted = true
	vote.reply.Term = vote.args.Term
	vote.notify <- struct{}{}
	return
} else if vote.args.Term < rf.curTerm {
	vote.reply.VoteGranted = false
	vote.reply.Term = rf.curTerm
} else {
	vote.reply.VoteGranted = false
	vote.reply.Term = rf.curTerm
}
rf.mu.Unlock()
vote.notify <- struct{}{}

rf.mu.Lock()
// a leader send AppendEntries, if args.term >= curTerm
// acknowledge the leader and become follower
// or reject the request then stay as candidate
if entry.args.Term >= rf.curTerm {
	rf.curTerm = entry.args.Term
	rf.votedFor = entry.args.LeaderId
	rf.state = &Follower{}
	rf.mu.Unlock()
	entry.reply.Term = entry.args.Term
	entry.reply.Success = true
	entry.notify <- struct{}{}
	return
} else {
	entry.reply.Term = rf.curTerm
	entry.reply.Success = false
}
rf.mu.Unlock()
entry.notify <- struct{}{}

// send RequestVote RPC
replyChan := make(chan RequestVoteReply, len(rf.peers)-1)
args := RequestVoteArgs{
	Term:        rf.curTerm,
	CandidateId: rf.me,
}
for i := range rf.peers {
	if i == rf.me {
		continue
	}
	go func(pees *labrpc.ClientEnd) {
		var reply RequestVoteReply
		if pees.Call("Raft.RequestVote", &args, &reply) {
			replyChan <- reply
		}
	}(rf.peers[i])
}

grantedCnt := 1
minVote := len(rf.peers)/2 + 1

rf.mu.Lock()
if reply.Term > rf.curTerm {
	// received larger term, become follower
	rf.curTerm = reply.Term
	rf.votedFor = -1
	rf.state = &Follower{}
	rf.mu.Unlock()
	return
} else if reply.Term == rf.curTerm && reply.VoteGranted {
	// received a grantVote
	grantedCnt++
	if grantedCnt >= minVote {
		rf.state = &Leader{}
		rf.mu.Unlock()
		return
	}
}
rf.mu.Unlock()

func (l *Leader) Run(rf *Raft) {
    // TODO: send heartbeat
    // TODO: send AppendEntries RPCs to all other servers

	// make heartbeat timer
	timer := time.NewTimer(heartbeatTimeout())
	for {
		select {
		case <-timer.C:
		// h 超时
		case reply := <-replyChan:
		// k AppendEntries 回答
		case vote := <-rf.voteChan:
		// i RequestVote 请求
		case entry := <-rf.entryChan:
		// j AppendEntries 请求
		}
	}
}

return

rf.mu.Lock()
if vote.args.Term > rf.curTerm {
	rf.curTerm = vote.args.Term
	rf.votedFor = vote.args.CandidateId
	rf.state = &Follower{}
	vote.reply.VoteGranted = true
	vote.reply.Term = rf.curTerm
	rf.mu.Unlock()
	vote.notify <- struct{}{}
	return
} else if vote.args.Term <= rf.curTerm {
	vote.reply.VoteGranted = false
	vote.reply.Term = rf.curTerm
}
rf.mu.Unlock()
vote.notify <- struct{}{}

rf.mu.Lock()
if entry.args.Term > rf.curTerm {
	rf.curTerm = entry.args.Term
	rf.votedFor = entry.args.LeaderId
	rf.state = &Follower{}
	entry.reply.Success = true
	entry.reply.Term = rf.curTerm
	rf.mu.Unlock()
	entry.notify <- struct{}{}
	return
} else if entry.args.Term <= rf.curTerm {
	entry.reply.Success = false
	entry.reply.Term = rf.curTerm
}
rf.mu.Unlock()
entry.notify <- struct{}{}

// send heartbeat
rf.mu.Lock()
replyChan := make(chan AppendEntriesReply, len(rf.peers))
args := AppendEntriesArgs{
	Term:     rf.curTerm,
	LeaderId: rf.me,
}
// send AppendEntries RPCs to all other servers
for i := range rf.peers {
	if i == rf.me {
		continue
	}
	go func(peer *labrpc.ClientEnd) {
		var reply AppendEntriesReply
		if peer.Call("Raft.AppendEntries", &args, &reply) {
			replyChan <- reply
		}
	}(rf.peers[i])
}
rf.mu.Unlock()

rf.mu.Lock()
if vote.args.Term > rf.curTerm {
	rf.curTerm = vote.args.Term
	rf.votedFor = vote.args.CandidateId
	rf.state = &Follower{}
	vote.reply.VoteGranted = true
	vote.reply.Term = rf.curTerm
	rf.mu.Unlock()
	vote.notify <- struct{}{}
	return
} else if vote.args.Term <= rf.curTerm {
	vote.reply.VoteGranted = false
	vote.reply.Term = rf.curTerm
}
rf.mu.Unlock()
vote.notify <- struct{}{}

type Task struct {
	Type     int		// 任务的类型，包括 Map 和 Reduce
	Filename string		// 当 Type 为 Map 时字段有效，表示需要作用 Map 函数的文件的文件名
	TaskId   int		// 用于标记这个 Task 的 Id，在对 Coordinator 回馈结果的时候作为参数返回
	ReduceId int		// Coordinator 要求这个 Worker 统计的 Reduce 的 Id
}

const (
	Map = iota
	Reduce
)

// 创建一个对应 Map 的 Task 结构
func NewMapTask(filename string, taskId int) Task {
	return Task{
		Type:     Map,
		Filename: filename,
		TaskId:   taskId,
	}
}

// 创建一个对应 Reduce 的 Task 的结构
func NewReduceTask(reduceId, taskId int) Task {
	return Task{
		Type:     Reduce,
		TaskId:   taskId,
		ReduceId: reduceId,
	}
}

// NeedWork RPC: 用于请求 Coordinator 分配新的任务
// NeedWork 的参数
type NeedWorkArgs struct {
}

// NeedWork 的返回值
type NeedWorkReply struct {
	T         Task
	ReduceCnt int
}

// FinishWork RPC: 用于告知 Coordinator 任务已完成
// FinishWork 的参数
type FinishWorkArgs struct {
	TaskId int
}

// FinishWork 的返回值
type FinishWorkReply struct {
}

// 添加的 import
import (
	"encoding/json"
	"fmt"
	"io"
	"os"
	"sort"
	"strings"
)

// 引入的排序 interface，从 mrsequential.go copy 来
// 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) {
	for {
		needWordReply := NeedWorkReply{}
		ok := call("Coordinator.NeedWork", &NeedWorkArgs{}, &needWordReply)
		if !ok {
			// 当 RPC 调用失败，我们认为 Coordinator 任务完成退出，
			// 所以 Worker 理应退出终止
			// Coordinator finish its work
			break
		}
		if needWordReply.T.Type == Map {
			// 处理 Map 任务的逻辑
			filename := needWordReply.T.Filename
			file, err := os.Open(filename)
			if err != nil {
				log.Fatalf("cannot open %v", filename)
			}
			content, err := io.ReadAll(file)
			if err != nil {
				log.Fatalf("cannot read %v", filename)
			}
			file.Close()
			// 得到单个文件的 Key-Value
			kva := mapf(filename, string(content))

			// 生成 intermediate 文件的逻辑
			// 首先计算每个 Key 的 reduceId，保存在 intermediate 这个二维切片中
			intermediate := make([][]KeyValue, needWordReply.ReduceCnt)
			for _, kv := range kva {
				reduceTask := ihash(kv.Key) % needWordReply.ReduceCnt
				intermediate[reduceTask] = append(intermediate[reduceTask], kv)
			}
			// intermediate[i] 对应 mr-{taskId}-{i} 这个中间文件
			// 以 paper 中提示的使用 json 的方式写入
			for i := 0; i < needWordReply.ReduceCnt; i++ {
				ofilename := fmt.Sprintf("mr-%d-%d", needWordReply.T.TaskId, i)
				// ofile, _ := os.Create(ofilename)
				tf, _ := os.CreateTemp("./", ofilename)
				enc := json.NewEncoder(tf)
				for _, kv := range intermediate[i] {
					enc.Encode(&kv)
				}
				tf.Close()
				os.Rename(tf.Name(), ofilename)
			}
		} else if needWordReply.T.Type == Reduce {
			// 处理 Reduce 任务的逻辑
			// 找出目录下所有对应该 Reduce 的文件名
			// find all files corresponding to this reduce task
			var filenames []string
			files, err := os.ReadDir(".")
			if err != nil {
				log.Fatalf("cannot read current directory")
			}
			for _, file := range files {
				if file.IsDir() {
					continue
				}
				filename := file.Name()
				prefix := "mr-"
				suffix := fmt.Sprintf("-%d", needWordReply.T.ReduceId)
				if strings.HasPrefix(filename, prefix) && strings.HasSuffix(filename, suffix) {
					filenames = append(filenames, filename)
				}
			}

			// 对所有已找到的文件进行读取
			// do reduce job
			var kva []KeyValue
			for _, filename := range filenames {
				file, err := os.Open(filename)
				if err != nil {
					log.Fatalf("cannot open %v", filename)
				}
				dec := json.NewDecoder(file)
				for {
					var kv KeyValue
					if err := dec.Decode(&kv); err != nil {
						break
					}
					kva = append(kva, kv)
				}
			}

			// 对该任务进行 Reduce 调用，逻辑和 mrsequential.go 完全一致，代码直接 copy
			// copy from mrsequential.go
			sort.Sort(ByKey(kva))
			oname := fmt.Sprintf("mr-out-%d", needWordReply.T.ReduceId)
			ofile, _ := os.Create(oname)
			i := 0
			for i < len(kva) {
				j := i + 1
				for j < len(kva) && kva[j].Key == kva[i].Key {
					j++
				}
				values := []string{}
				for k := i; k < j; k++ {
					values = append(values, kva[k].Value)
				}
				output := reducef(kva[i].Key, values)
				fmt.Fprintf(ofile, "%v %v\n", kva[i].Key, output)
				i = j
			}
			ofile.Close()
		} else {
			// unknown task type
			log.Fatalf("unknown task type: %v", needWordReply.T.Type)
		}

		// 运行到此代表该任务成功完成，通过 RPC 告知 Coordinator 任务完成
		// make FinishWork RPC call
		call("Coordinator.FinishWork", &FinishWorkArgs{TaskId: needWordReply.T.TaskId}, &FinishWorkReply{})
	}
	// log.Printf("Worker: %v exit", os.Getpid())
}

import (
	"fmt"
	"log"
	"sync"
	"time"
)

type Coordinator struct {
	// Your definitions here.
	// 保存一下 nReduce
	nReduce int

	// 定义任务队列和任务的计数器，
	// 用于任务的发送，判断 Map 和 Reduce 任务的交界以及何时关闭 Coordinator
	// task sending definition
	taskQueue chan Task
	taskWg    sync.WaitGroup

	// 任务 Id 计数器，自增
	// task id counter
	taskIdCounter int

	// 保存用于通知任务完成的 chan 的字段，
	// 若 Worker 完成并调用 RPC，处理函数会向对应的 channel 发送一个信号
	// task notification record
	taskNotifyMap  map[int]chan struct{}
	taskNotifyLock sync.Mutex

	// 用于标记所有任务是否完成
	// done flag
	done     bool
	doneLock sync.Mutex
}

// 获得一个新的 taskId
func (c *Coordinator) NewTaskId() int {
	c.taskIdCounter++
	return c.taskIdCounter
}

// Done 函数，返回所有任务是否已经完成
func (c *Coordinator) Done() bool {
	c.doneLock.Lock()
	ret := c.done
	c.doneLock.Unlock()

	// Your code here.
	return ret
}

func MakeCoordinator(files []string, nReduce int) *Coordinator {
	// 初始化一个 Coordinator
	c := Coordinator{
		nReduce:       nReduce,
		taskQueue:     make(chan Task, 100),
		taskWg:        sync.WaitGroup{},
		taskIdCounter: 0,
		taskNotifyMap: make(map[int]chan struct{}),
		done:          false,
	}
	// log.Printf("%d file found\n", len(files))

	// 由于 Reduce 任务必须在所有 Map 任务完成后才能开始
	// 添加 Map 任务计数到 WaitGroup，用于标识所有的 Map 任务是否完成
	c.taskWg.Add(len(files))

	// Your code here.

	// 开启一个新的 goroutine 向 c.taskQueue channel 中发送任务
	// make send task goroutine
	go func(files []string) {
		// 发送 Map 任务
		// send map task
		for _, filename := range files {
			c.taskQueue <- NewMapTask(filename, c.NewTaskId())
		}
		// log.Println("all task sent, waiting for next reduce task")
		// 等待所有 Map 任务完成，这时才能发送 Reduce 任务
		// wait for all map task done
		c.taskWg.Wait()
		// log.Println("start sending reduce task")

		// 添加 Reduce 任务的计数
		// send reduce task
		c.taskWg.Add(nReduce)
        
		// 下面这个 goroutine 的作用是
		// 等待 WaitGroup 值为 0 时，此时表示所有的 Reduce 任务也完成了
		// 所以可以将 c.done 置为 true，表示 Coordinator 可以结束了
		// make Done() check goroutine
		// log.Println("waiting exit goroutine created")
		go func() {
			c.taskWg.Wait()
			c.doneLock.Lock()
			c.done = true
			c.doneLock.Unlock()
		}()
		// log.Println("sending reduce task")

		// 发送 Reduce 任务
		for i := 0; i < nReduce; i++ {
			c.taskQueue <- NewReduceTask(i, c.NewTaskId())
		}
	}(files)

	c.server()
	return &c
}

func (c *Coordinator) NeedWork(args *NeedWorkArgs, reply *NeedWorkReply) error {
	// 从任务队列 c.taskQueue 中取出一个任务作为 RPC 返回值
	task, ok := <-c.taskQueue
	if !ok {
		return fmt.Errorf("cannot get task from task queue")
	}
	reply.T = task
	reply.ReduceCnt = c.nReduce

	// 使用 taskId 作为唯一任务标识，创建一个 struct{} channel
	// 用于 FinishWork RPC 函数通知任务已被 Worker 完成
	// 注意锁的临界区，否则会导致 condition race
	// make task notification channel
	c.taskNotifyLock.Lock()
	c.taskNotifyMap[task.TaskId] = make(chan struct{})
	// 启动 goroutine 定时
	// set timer for task
	go func(taskChan chan struct{}) {
		select {
		case <-taskChan:
			// 任务完成
			// task done
			return
		case <-time.After(10 * time.Second):
			// 任务超时，重新将任务放回任务队列中
			c.taskQueue <- task
		}
	}(c.taskNotifyMap[task.TaskId])
	c.taskNotifyLock.Unlock()

	return nil
}

func (c *Coordinator) FinishWork(args *FinishWorkArgs, reply *FinishWorkReply) error {
	taskId := args.TaskId
	c.taskNotifyLock.Lock()
	// 思考这里为什么要检查 taskId 所对应的值是否有效
	// 若考虑网络延迟之类的不确定因素，可能导致同一个 taskId 调用两次 FinishWork 函数
	notifyChan, ok := c.taskNotifyMap[taskId]
	if !ok {
		c.taskNotifyLock.Unlock()
		return nil
	}
	// 向对应 channel 发送信号表示任务完成，同时任务计数 -1
	// notify task done
	notifyChan <- struct{}{}
	c.taskWg.Done()
	// log.Printf("task %d done\n", taskId)

	// 删除对应的 channel
	// delete task notification channel
	delete(c.taskNotifyMap, taskId)
	c.taskNotifyLock.Unlock()
	return nil
}