package mr

import (
	"encoding/json"
	"fmt"
	"hash/fnv"
	"io"
	"log"
	"net/rpc"
	"os"
	"sort"
	"sync/atomic"
	"time"
)

// Map functions return a slice of KeyValue.
type KeyValue struct {
	Key   string
	Value string
}

// use ihash(key) % NReduce to choose the reduce
// 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)
}

// main/mrworker.go calls this function.
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()

	machine := MakeMachine(-1, -1)
	var id atomic.Int32

	for {
		reply := Machine{}
		state := machine.State
		switch state {
		case MACHINE_INIT:
			// log.Printf("worker init %d\n", machine.Id)
			if ok := call("Coordinator.UpdateState", machine, &reply); !ok {
				reply.State = MACHINE_DONE
			} else {
				id.Store(int32(reply.Id))
				go func() {
					for {
						time.Sleep(time.Second)
						rid := id.Load()
						if ok := call("Coordinator.FlushTimeout", rid, &Empty{}); !ok {
							log.Printf("FlushTimeout error %v\n", ok)
						}
					}
				}()
			}
			machine = &reply
		case MACHINE_READY:
			if ok := runOnceTask(machine, mapf, reducef); !ok {
				machine.State = MACHINE_IDLE
			}
		case MACHINE_IDLE:
			// log.Printf("worker idle %d\n", machine.Id)
			machine.Task.TaskType = TASK_NONE
			if ok := call("Coordinator.GetTask", machine, &reply); !ok {
				reply.State = MACHINE_DONE
			}
			machine = &reply
		case MACHINE_DONE:
			return
		default:
			goto SLEEP
		}
		continue
	SLEEP:
		time.Sleep(time.Second)
	}
}

func runOnceTask(machine *Machine,
	mapf func(string, string) []KeyValue,
	reducef func(string, []string) string) bool {
	machine.State = MACHINE_RUNNING
	task := &machine.Task

	switch task.TaskType {
	case TASK_MAP:
		kva := workMap(mapf, task)
		saved_fname := saveMap(kva, task)
		if saved_fname == nil {
			return false
		}
	case TASK_REDUCE:
		workRecude(reducef, task)
	default:
		log.Printf("unknown task type %d", task.TaskType)
		return false
	}

	// log.Printf("submit task %d %+v %+v", task.TaskId, task.MapTask, task.ReduceTask)
	reply := Machine{}
	if ok := call("Coordinator.SubmitTask", machine, &reply); !ok {
		return false
	}
	*machine = reply

	return true
}

func workMap(mapf func(string, string) []KeyValue, params *Task) []KeyValue {
	filename := params.MapTask.Fname

	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()
	kva := mapf(filename, string(content))
	return kva
}

func saveMap(kva []KeyValue, params *Task) []string {
	nReduce := params.MapTask.NReduce
	fnames := make([]string, nReduce)
	files := make([]*os.File, nReduce)
	encodes := make([]*json.Encoder, nReduce)

	for i := range fnames {
		fnames[i] = fmt.Sprintf("mr-%d-%d", params.TaskId, i)
		outfname := fnames[i]
		ofile, err := os.Create(outfname)
		if err != nil {
			log.Fatalf("cannot create %v", outfname)
		}
		files[i] = ofile
		enc := json.NewEncoder(ofile)
		encodes[i] = enc
	}
	defer func() {
		for _, f := range files {
			if f != nil {
				f.Close()
			}
		}
	}()

	for _, kv := range kva {
		enc := encodes[ihash(kv.Key)%nReduce]
		if err := enc.Encode(&kv); err != nil {
			log.Fatalf("cannot encode %v", kv)
			break
		}
	}

	return fnames
}

// 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 workRecude(reducef func(string, []string) string, params *Task) {
	reduceId := params.TaskId
	fnames := params.ReduceTask.Fnames
	intermediate := []KeyValue{}

	for _, fname := range fnames {
		file, err := os.Open(fname)
		if err != nil {
			log.Fatal("can't open file")
		}

		dec := json.NewDecoder(file)
		for {
			var kv KeyValue
			if err := dec.Decode(&kv); err != nil {
				break
			}
			intermediate = append(intermediate, kv)
		}

		file.Close()
	}

	sort.Sort(ByKey(intermediate))

	oname := fmt.Sprintf("mr-out-%d", reduceId)
	ofile, err := os.Create(oname)
	if err != nil {
		log.Fatal("out file open error")
	}

	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()
}

// example function to show how to make an RPC call to the coordinator.
//
// the RPC argument and reply types are defined in rpc.go.
func CallExample() {

	// declare an argument structure.
	args := ExampleArgs{}

	// fill in the argument(s).
	args.X = 99

	// declare a reply structure.
	reply := ExampleReply{}

	// send the RPC request, wait for the reply.
	// the "Coordinator.Example" tells the
	// receiving server that we'd like to call
	// the Example() method of struct Coordinator.
	ok := call("Coordinator.Example", &args, &reply)
	if ok {
		// reply.Y should be 100.
		fmt.Printf("reply.Y %v\n", reply.Y)
	} else {
		fmt.Printf("call failed!\n")
	}
}

// send an RPC request to the coordinator, wait for the response.
// usually returns true.
// returns false if something goes wrong.
func call(rpcname string, args interface{}, reply interface{}) bool {
	// c, err := rpc.DialHTTP("tcp", "127.0.0.1"+":1234")
	sockname := coordinatorSock()
	c, err := rpc.DialHTTP("unix", sockname)
	if err != nil {
		log.Fatal("dialing:", err)
	}
	defer c.Close()

	err = c.Call(rpcname, args, reply)
	if err == nil {
		return true
	}

	fmt.Println(err)
	return false
}
