package mr

import (
	"crypto/rand"
	"errors"
	"fmt"
	"log"
	"math/big"
	"net"
	"net/http"
	"net/rpc"
	"os"
	"strconv"
	"strings"
	"sync"
	"time"
)

const (
	JOB_TIMEOUT_S = 10 * time.Second
	DEBUG         = true
)

type Coordinator struct {
	MapJobs    []Job
	ReduceJobs []Job
	NReduce    int
	mu         sync.Mutex
	mapWg      sync.WaitGroup
}

type Job struct {
	WorkerId  string
	Filenames []string
	Status    bool
	Active    bool
	LastPing  time.Time
}

// Your code here -- RPC handlers for the worker to call.
//
// the RPC argument and reply types are defined in rpc.go.
func (c *Coordinator) GetTasks(args *RpcArgument, reply *RpcReply) error {
	c.mu.Lock()
	defer c.mu.Unlock()

	if args.Method != "request_task" {
		return errors.New("argument .Method not valid for this procedure")
	}
	// check for any non-completed OR non-active map jobs
	for i, mj := range c.MapJobs {
		// assign job that not already done OR not actively being run by other workers
		// status = false AND active = false
		if !mj.Status && !mj.Active {
			wId, err := genWorkerId()
			if err != nil {
				fmt.Println(err)
			}
			reply.NReduce = c.NReduce
			reply.TaskType = "map"
			reply.Filenames = mj.Filenames
			reply.WorkerId = wId
			// set job status, taken by a worke
			c.MapJobs[i].WorkerId = wId
			c.MapJobs[i].Active = true
			c.MapJobs[i].LastPing = time.Now()
			if DEBUG {
				fmt.Printf("[Map Job Assigned] %s | file: %s | status: %t | active: %t |\n",
					c.MapJobs[i].WorkerId,
					c.MapJobs[i].Filenames[0],
					c.MapJobs[i].Status, c.MapJobs[i].Active)
			}
			return nil
		}
	}
	c.mapWg.Wait()
	// check for reduce jobs
	for i, rj := range c.ReduceJobs {
		if !rj.Status && !rj.Active {
			wId, err := genWorkerId()
			if err != nil {
				fmt.Println(err)
			}
			reply.WorkerId = wId
			reply.NReduce = c.NReduce
			reply.TaskType = "reduce"
			reply.Filenames = rj.Filenames
			// set job status, taken by a worker
			c.ReduceJobs[i].WorkerId = wId
			c.ReduceJobs[i].Active = true
			c.ReduceJobs[i].LastPing = time.Now()
			if DEBUG {
				fmt.Printf("[Reduce Job Assigned] %s | status: %t | active: %t |\n",
					c.ReduceJobs[i].WorkerId, c.ReduceJobs[i].Status, c.ReduceJobs[i].Active)
			}
			return nil
		}
	}

	reply.TaskType = "done"
	reply.NReduce = c.NReduce
	reply.Filenames = []string{}
	return nil
}

// Set task complete
func (c *Coordinator) SetTaskDone(args *RpcArgument, reply *RpcReply) error {
	if args.Method != "set_task_done" {
		return errors.New("argument .Method not valid for this procedure")
	}
	if args.TaskType == "map" {
		wId := args.WorkerId
		for i, _ := range c.MapJobs {
			if c.MapJobs[i].WorkerId == wId {
				c.MapJobs[i].Status = true
				c.MapJobs[i].Active = false
				c.mapWg.Done()
				if DEBUG {
					fmt.Printf("[DONE] %s | file: %s | status: %t | active: %t\n",
						c.MapJobs[i].WorkerId,
						c.MapJobs[i].Filenames[0],
						c.MapJobs[i].Status, c.MapJobs[i].Active)
				}
			}
		}
		// Check again
		for _, ifn := range args.IntermediateFiles {
			reduceJobIdx := strings.Split(ifn, "-")[2]
			idx, err := strconv.Atoi(reduceJobIdx)
			if err != nil {
				fmt.Printf("Error converting string to integer: %v\n", err)
			}
			c.ReduceJobs[idx].Filenames = append(c.ReduceJobs[idx].Filenames, ifn)
		}

	} else if args.TaskType == "reduce" {
		wId := args.WorkerId
		for i, _ := range c.ReduceJobs {
			if c.ReduceJobs[i].WorkerId == wId {
				c.ReduceJobs[i].Status = true
			}
		}
	} else {
		return errors.New("argument .TaskType not valid for this procedure")
	}
	return nil
}

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

// check scheduled tasks; if their age > JOB_TIMEOUT_S
// reset the job status, making it available for scheduling
func (c *Coordinator) CheckTaskValidity() {
	c.mu.Lock()
	defer c.mu.Unlock()
	for i, rj := range c.ReduceJobs {
		timeNow := time.Now()
		if rj.Active && !rj.Status {
			duration := timeNow.Sub(rj.LastPing)
			if duration >= JOB_TIMEOUT_S {
				c.ReduceJobs[i].Active = false
				c.ReduceJobs[i].Status = false
			}
		}
	}
}

// main/mrcoordinator.go calls Done() periodically to find out
// if the entire job has finished.
func (c *Coordinator) Done() bool {
	ret := false
	// Check any non-completed jobs
	for _, mj := range c.MapJobs {
		if !mj.Status {
			return ret
		}
	}
	for _, mj := range c.ReduceJobs {
		if !mj.Status {
			return ret
		}
	}
	// all mj completed
	ret = true
	return ret
}

// create a Coordinator.
// main/mrcoordinator.go calls this function.
// nReduce is the number of reduce tasks to use.
func MakeCoordinator(files []string, nReduce int) *Coordinator {
	c := Coordinator{}

	c.NReduce = nReduce

	for i := range files {
		mj := Job{}
		mj.Filenames = []string{files[i]}
		mj.Status = false
		mj.Active = false
		c.mapWg.Add(1)
		c.MapJobs = append(c.MapJobs, mj)
	}

	for j := 0; j < nReduce; j++ {
		reduceJob := Job{}
		reduceJob.Active = false
		reduceJob.Filenames = []string{}
		reduceJob.Status = false
		c.ReduceJobs = append(c.ReduceJobs, reduceJob)
	}

	c.server()
	return &c
}

func genWorkerId() (string, error) {
	const charset = "abcdefghijklmnopqrstuvwxyzABCDEFGHIJKLMNOPQRSTUVWXYZ0123456789"
	result := make([]byte, 8)
	for i := range result {
		index, err := rand.Int(rand.Reader, big.NewInt(int64(len(charset))))
		if err != nil {
			return "", fmt.Errorf("error generating random number: %v", err)
		}
		result[i] = charset[index.Int64()]
	}
	return string(result), nil
}
