package main

import (
	"bytes"
	"encoding/json"
	"fmt"
	"strconv"

	"github.com/hyperledger/fabric/core/chaincode/shim"
	sc "github.com/hyperledger/fabric/protos/peer"
	// "github.com/hyperledger/fabric-chaincode-go/shim"
	// sc "github.com/hyperledger/fabric-protos-go/peer"
)

type SmartContract struct{}

type SharedData struct {
	PatientId     string `json:patientId`
	DeviceId      string `json:deviceId`
	EncryptedData string `json:encryptedData`
	Step          string `json:step`
}

type Doctor struct {
	Id            string       `json:id`
	Type          string       `json:type`
	PrivateMatrix string       `json:privateMatrix`
	SharedData    []SharedData `json:sharedPatients`
	Department    string       `json:department`
}

type Device struct {
	DeviceID    string   `json:deviceid`
	MasterKey   string   `json:masterkey`
	SharedWith  []string `json:sharedwith`
	RequestedBy []string `json:requestedby`
}

type Patient struct {
	Id            string   `json:id`
	Type          string   `json;type`
	Devices       []Device `json:devices`
	PrivateMatrix string   `json:privateMatrix`
}

func (s *SmartContract) Init(APIstub shim.ChaincodeStubInterface) sc.Response {
	return shim.Success(nil)
}

func (s *SmartContract) Invoke(APIstub shim.ChaincodeStubInterface) sc.Response {

	// Retrieve the requested Smart Contract function and arguments
	function, args := APIstub.GetFunctionAndParameters()
	// Route to the appropriate handler function to interact with the ledger appropriately
	if function == "addDoctor" {
		return s.addDoctor(APIstub, args)
	} else if function == "addPatient" {
		return s.addPatient(APIstub, args)
	} else if function == "addDevice" {
		return s.addDevice(APIstub, args)
	} else if function == "queryData" {
		return s.queryData(APIstub, args)
	} else if function == "shareData" {
		return s.shareData(APIstub, args)
	} else if function == "requestData" {
		return s.requestData(APIstub, args)
	} else if function == "getAllDoctors" {
		return s.getAllDoctors(APIstub, args)
	} else if function == "getAllPatients" {
		return s.getAllPatients(APIstub, args)
	} else if function == "queryShareData" {
		return s.queryShareData(APIstub, args)
	}
	fmt.Println(function, args)
	return shim.Error("Invalid smart contract function name recived : " + function)
}

func (s *SmartContract) addDoctor(APIstub shim.ChaincodeStubInterface, args []string) sc.Response {

	if len(args) != 3 {
		return shim.Error("Incorrect number of arguments. Expecting 3")
	}

	Id := args[0]
	privateMatrix := args[1]
	dept := args[2]
	sd := []SharedData{}
	doctorData := Doctor{Id: Id, Type: "Doctor", PrivateMatrix: privateMatrix, SharedData: sd, Department: dept}

	dataAsBytes, err := json.Marshal(doctorData)
	if err != nil {
		return shim.Error("Error marshaling data.")
	}
	APIstub.PutState(Id, dataAsBytes)

	return shim.Success([]byte("Doctor Enrolled Successfully."))
}

func (s *SmartContract) addPatient(APIstub shim.ChaincodeStubInterface, args []string) sc.Response {

	if len(args) != 2 {
		return shim.Error("Incorrect number of arguments. Expecting 2")
	}

	fmt.Printf("New chaincode is running")

	Id := args[0]
	d := []Device{}
	patientData := Patient{Id: args[0], Type: "Patient", Devices: d, PrivateMatrix: args[1]}
	dataAsBytes, err := json.Marshal(patientData)
	if err != nil {
		return shim.Error("Error marshaling data.")
	}
	APIstub.PutState(Id, dataAsBytes)

	return shim.Success([]byte("Patient Enrolled Successfully."))
}

func (s *SmartContract) addDevice(APIstub shim.ChaincodeStubInterface, args []string) sc.Response {

	if len(args) != 3 {
		return shim.Error("Incorrect number of arguments. Expecting 3")
	}

	Id := args[0]
	dataAsBytes, err := APIstub.GetState(Id)
	if err != nil {
		return shim.Error("Error getting data for given key.")
	}
	patientData := Patient{}
	err = json.Unmarshal(dataAsBytes, &patientData)
	if err != nil {
		return shim.Error("Error unmarshaling data.")
	}
	d := Device{DeviceID: args[1], MasterKey: args[2], SharedWith: []string{}}
	patientData.Devices = append(patientData.Devices, d)

	dataAsBytes, err = json.Marshal(patientData)
	if err != nil {
		return shim.Error("Error marshaling data.")
	}
	APIstub.PutState(Id, dataAsBytes)

	return shim.Success([]byte("Device enrolled successfully"))
}

func (s *SmartContract) shareData(APIstub shim.ChaincodeStubInterface, args []string) sc.Response {

	if len(args) != 5 {
		return shim.Error("Incorrect number of arguments. Expecting 5 got : " + strconv.Itoa(len(args)))
	}
	PatientID := args[0]
	doctorid := args[1]
	DeviceID := args[2]
	encryptedData := args[3]
	Step := args[4]
	dataAsBytes, err := APIstub.GetState(PatientID)
	if err != nil {
		return shim.Error("Error getting data for given key.")
	}
	patientData := Patient{}
	err = json.Unmarshal(dataAsBytes, &patientData)
	if err != nil {
		return shim.Error("Error unmarshaling data.")
	}
	for i, device := range patientData.Devices {
		if device.DeviceID == DeviceID {
			for _, doctor := range patientData.Devices[i].SharedWith {
				if doctor == doctorid {
					doctorAsBytes, err := APIstub.GetState(doctorid)
					if err != nil {
						return shim.Error("Error getting data for given key.")
					}
					doctor := Doctor{}
					err = json.Unmarshal(doctorAsBytes, &doctor)
					if err != nil {
						return shim.Error("Error unmarshaling data.")
					}
					for k, shared := range doctor.SharedData {
						if (shared.PatientId == PatientID) && (shared.DeviceId == DeviceID) {
							doctor.SharedData[k].EncryptedData = encryptedData
							doctor.SharedData[k].Step = Step
							dataAsBytes, err := json.Marshal(doctor)
							if err != nil {
								return shim.Error("Error marshaling data.")
							}
							APIstub.PutState(doctorid, dataAsBytes)
							return shim.Success([]byte("successfully updated data share"))
						}
					}
				}
			}

			patientData.Devices[i].SharedWith = append(patientData.Devices[i].SharedWith, doctorid)
			dataAsBytes, err := APIstub.GetState(doctorid)
			if err != nil {
				return shim.Error("Error getting data for given key.")
			}
			doctorData := Doctor{}
			err = json.Unmarshal(dataAsBytes, &doctorData)
			if err != nil {
				return shim.Error("Error unmarshaling data.")
			}
			s := SharedData{PatientId: PatientID, DeviceId: DeviceID, EncryptedData: encryptedData, Step: Step}
			doctorData.SharedData = append(doctorData.SharedData, s)
			dataAsBytes, err = json.Marshal(patientData)
			if err != nil {
				return shim.Error("Error marshaling data.")
			}
			APIstub.PutState(PatientID, dataAsBytes)
			dataAsBytes, err = json.Marshal(doctorData)
			if err != nil {
				return shim.Error("Error marshaling data.")
			}
			APIstub.PutState(doctorid, dataAsBytes)
			return shim.Success([]byte("Successfully started data share."))
		}
	}
	return shim.Error("Error sharing data. Device/Patient not found.")
}

func (s *SmartContract) requestData(APIstub shim.ChaincodeStubInterface, args []string) sc.Response {

	if len(args) != 3 {
		return shim.Error("Incorrect number of arguments. Expecting 3")
	}
	doctorid := args[0]
	PatientID := args[1]
	DeviceID := args[2]
	dataAsBytes, err := APIstub.GetState(PatientID)
	if err != nil {
		return shim.Error("Error getting data for given key.")
	}
	patientData := Patient{}
	err = json.Unmarshal(dataAsBytes, &patientData)
	if err != nil {
		return shim.Error("Error unmarshaling data.")
	}
	for i, device := range patientData.Devices {
		if device.DeviceID == DeviceID {
			patientData.Devices[i].RequestedBy = append(patientData.Devices[i].RequestedBy, doctorid)
			dataAsBytes, err = json.Marshal(patientData)
			if err != nil {
				return shim.Error("Error marshaling data.")
			}
			APIstub.PutState(PatientID, dataAsBytes)
			return shim.Success([]byte("Successfully sent request."))
		}
	}
	return shim.Error("Error sending request. Device/Patient not found.")
}

func (s *SmartContract) queryData(APIstub shim.ChaincodeStubInterface, args []string) sc.Response {
	if len(args) != 1 {
		return shim.Error("Incorrect number of arguments. Expecting 1")
	}
	dataAsBytes, err := APIstub.GetState(args[0])
	if err != nil {
		return shim.Error("Error getting data for given key.")
	}
	return shim.Success(dataAsBytes)
}
func (s *SmartContract) queryShareData(APIstub shim.ChaincodeStubInterface, args []string) sc.Response {
	if len(args) != 3 {
		return shim.Error("Incorrect number of arguments. Expecting 3")
	}
	PatientID := args[0]
	doctorid := args[1]
	DeviceID := args[2]
	doctorAsBytes, err := APIstub.GetState(doctorid)
	if err != nil {
		return shim.Error("Error getting data for given key.")
	}
	doctor := Doctor{}
	err = json.Unmarshal(doctorAsBytes, &doctor)
	if err != nil {
		return shim.Error("Error unmarshaling data.")
	}
	for _, shared := range doctor.SharedData {
		if (shared.PatientId == PatientID) && (shared.DeviceId == DeviceID) {
			dataAsBytes, err := json.Marshal(shared)
			if err != nil {
				return shim.Error("Error marshaling data.")
			}
			return shim.Success(dataAsBytes)
		}
	}
	return shim.Error("Error getting data. Device/Patient/Doctor not found.")
}
func (s *SmartContract) getAllDoctors(APIstub shim.ChaincodeStubInterface, args []string) sc.Response {
	if len(args) != 1 {
		return shim.Error("Incorrect number of arguments. Expecting 1")
	}

	if args[0] == "" {
		return shim.Error("Department missing.")
	}
	query := `{"selector": {"Type": "Doctor","Department": "` + args[0] + `"}}`
	dataAsBytes, _ := getQueryResultForQueryString(APIstub, query)
	return shim.Success(dataAsBytes)
}

func (s *SmartContract) getAllPatients(APIstub shim.ChaincodeStubInterface, args []string) sc.Response {
	query := `{"selector": {"Type": "Patient"}}`
	dataAsBytes, _ := getQueryResultForQueryString(APIstub, query)
	return shim.Success(dataAsBytes)
}

func getQueryResultForQueryString(APIstub shim.ChaincodeStubInterface, queryString string) ([]byte, error) {
	fmt.Printf("- getQueryResultForQueryString queryString:\n%s\n", queryString)
	resultsIterator, err := APIstub.GetQueryResult(queryString)
	defer resultsIterator.Close()
	if err != nil {
		return nil, err
	}
	// buffer is a JSON array containing QueryRecords
	var buffer bytes.Buffer
	buffer.WriteString("[")
	bArrayMemberAlreadyWritten := false
	for resultsIterator.HasNext() {
		queryResponse,
			err := resultsIterator.Next()
		if err != nil {
			return nil, err
		}
		// Add a comma before array members, suppress it for the first array member
		if bArrayMemberAlreadyWritten == true {
			buffer.WriteString(",")
		}
		buffer.WriteString("{\"Key\":")
		buffer.WriteString("\"")
		buffer.WriteString(queryResponse.Key)
		buffer.WriteString("\"")
		buffer.WriteString(", \"Record\":")
		// Record is a JSON object, so we write as-is
		buffer.WriteString(string(queryResponse.Value))
		buffer.WriteString("}")
		bArrayMemberAlreadyWritten = true
	}
	buffer.WriteString("]")
	fmt.Printf("- getQueryResultForQueryString queryResult:\n%s\n", buffer.String())
	return buffer.Bytes(), nil
}

// The main function is only relevant in unit test mode. Only included here for completeness.
func main() {

	// Create a new Smart Contract
	err := shim.Start(new(SmartContract))
	if err != nil {
		fmt.Printf("Error creating new Smart Contract: %s", err)
	}
}