package iam

import (
	"encoding/json"
	"errors"
)

type (
	// Policy grammar https://docs.aws.amazon.com/IAM/latest/UserGuide/reference_policies_grammar.html
	// Currently 'NotPrincipal', 'NotAction' and 'NotResource' are not supported (so cannot be unmarshalled).
	Policy struct {
		Version   string     `json:"Version,omitempty"`
		ID        string     `json:"Id,omitempty"`
		Statement Statements `json:"Statement"`
	}

	Statements []Statement

	Statement struct {
		SID        string     `json:"Sid,omitempty"`
		Principal  Principal  `json:"Principal,omitempty"`
		Effect     Effect     `json:"Effect"`
		Action     Action     `json:"Action"`
		Resource   Resource   `json:"Resource"`
		Conditions Conditions `json:"Condition,omitempty"`
	}

	Principal map[string][]string

	Effect string

	Action []string

	Resource []string

	Conditions map[string]Condition

	Condition map[string][]string
)

const Wildcard = "*"

const (
	AllowEffect Effect = "Allow"
	DenyEffect  Effect = "Deny"
)

func (s *Statements) UnmarshalJSON(data []byte) error {
	var list []Statement
	if err := json.Unmarshal(data, &list); err == nil {
		*s = list
		return nil
	}

	var elem Statement
	if err := json.Unmarshal(data, &elem); err != nil {
		return err
	}

	*s = []Statement{elem}

	return nil
}

func (p *Principal) UnmarshalJSON(data []byte) error {
	*p = make(Principal)

	var str string

	if err := json.Unmarshal(data, &str); err == nil {
		if str != Wildcard {
			return errors.New("invalid IAM string principal")
		}
		(*p)[Wildcard] = nil
		return nil
	}

	m := make(map[string]interface{})
	if err := json.Unmarshal(data, &m); err != nil {
		return err
	}

	for key, val := range m {
		element, ok := val.(string)
		if ok {
			(*p)[key] = []string{element}
			continue
		}

		list, ok := val.([]interface{})
		if !ok {
			return errors.New("invalid principal format")
		}

		resList := make([]string, len(list))
		for i := range list {
			val, ok := list[i].(string)
			if !ok {
				return errors.New("invalid principal format")
			}
			resList[i] = val
		}

		(*p)[key] = resList
	}

	return nil
}

func (a *Action) UnmarshalJSON(data []byte) error {
	var list []string
	if err := json.Unmarshal(data, &list); err == nil {
		*a = list
		return nil
	}

	var elem string
	if err := json.Unmarshal(data, &elem); err != nil {
		return err
	}

	*a = []string{elem}

	return nil
}

func (r *Resource) UnmarshalJSON(data []byte) error {
	var list []string
	if err := json.Unmarshal(data, &list); err == nil {
		*r = list
		return nil
	}

	var elem string
	if err := json.Unmarshal(data, &elem); err != nil {
		return err
	}

	*r = []string{elem}

	return nil
}

func (c *Condition) UnmarshalJSON(data []byte) error {
	*c = make(Condition)

	m := make(map[string]interface{})
	if err := json.Unmarshal(data, &m); err != nil {
		return err
	}

	for key, val := range m {
		element, ok := val.(string)
		if ok {
			(*c)[key] = []string{element}
			continue
		}

		list, ok := val.([]interface{})
		if !ok {
			return errors.New("invalid principal format")
		}

		resList := make([]string, len(list))
		for i := range list {
			val, ok := list[i].(string)
			if !ok {
				return errors.New("invalid principal format")
			}
			resList[i] = val
		}

		(*c)[key] = resList
	}

	return nil
}