package main

import (
	"crypto/x509"
	"encoding/json"
	"encoding/pem"
	"fmt"
	"io/ioutil"
	"log"
	"net/http"
	"text/template"
	"time"

	"github.com/golang-jwt/jwt"
	"github.com/google/uuid"
	"github.com/lestrrat-go/jwx/v2/jwk"
)

type AuthorizationRequest struct {
	ResponseType string
	Scope        string
	ClientId     string
	State        string
	RedirectUri  string
	Nonce        string
}

type TokenRequest struct {
	GrantType   string
	Code        string
	RedirectUri string
}

type TokenResponse struct {
	AccessToken  string
	TokenType    string
	RefreshToken string
	ExpiresIn    int
	IdToken      string
}

var (
	tokens = map[string]string{}
)

func createIdToken(clientId string) string {
	now := time.Now()
	expiresAt := now.Add(time.Hour * time.Duration(1))
	idToken := jwt.NewWithClaims(jwt.SigningMethodRS256, &jwt.StandardClaims{
		Issuer:    "https://example.com",
		Subject:   "1",
		Audience:  clientId,
		ExpiresAt: expiresAt.Unix(),
		NotBefore: now.Unix(),
		IssuedAt:  now.Unix(),
		Id:        uuid.NewString(),
	})

	keyData, _ := ioutil.ReadFile("insecure.pem")
	key, _ := jwt.ParseRSAPrivateKeyFromPEM(keyData)
	signedIdToken, _ := idToken.SignedString(key)
	return signedIdToken
}

func handler(w http.ResponseWriter, r *http.Request) {
	if r.URL.Path == "/" && r.Method == "GET" {
		w.WriteHeader(http.StatusOK)
		fmt.Fprintf(w, "Hello, world!\n")
	} else if r.URL.Path == "/authorize" && r.Method == "GET" {
		responseType := r.FormValue("response_type")
		if responseType == "code" {
			// Authorization Code Flow https://openid.net/specs/openid-connect-core-1_0.html#CodeFlowAuth
			ar := &AuthorizationRequest{
				ResponseType: r.FormValue("response_type"),
				Scope:        r.FormValue("scope"),
				ClientId:     r.FormValue("client_id"),
				State:        r.FormValue("state"),
				RedirectUri:  r.FormValue("redirect_uri"),
			}
			code := uuid.NewString()
			tokens[code] = uuid.NewString()
			url := fmt.Sprintf("%s?code=%s&state=%s", ar.RedirectUri, code, ar.State)
			http.Redirect(w, r, url, 302)
		} else if responseType == "id_token token" || responseType == "id_token" {
			// Implicit Flow https://openid.net/specs/openid-connect-core-1_0.html#ImplicitFlowAuth
			ar := &AuthorizationRequest{
				ResponseType: r.FormValue("response_type"),
				RedirectUri:  r.FormValue("redirect_uri"),
				Nonce:        r.FormValue("nonce"),
			}
			idToken := createIdToken(r.FormValue("client_id"))
			url := fmt.Sprintf("%s?access_token=example&token_type=bearer&id_token=%s&expires_in=3600&state=%s", ar.RedirectUri, idToken, ar.State)
			http.Redirect(w, r, url, 302)
		} else if responseType == "code id_token" || responseType == "code token" || responseType == "code id_token token" {
			// Hybrid Flow https://openid.net/specs/openid-connect-core-1_0.html#HybridFlowAuth
			w.WriteHeader(http.StatusNotImplemented)
		} else {
			w.WriteHeader(http.StatusNotFound)
			fmt.Fprintf(w, "Not Found\n")
		}
	} else if r.URL.Path == "/token" && r.Method == "POST" {
		tr := &TokenRequest{
			GrantType:   r.FormValue("grant_type"),
			Code:        r.FormValue("code"),
			RedirectUri: r.FormValue("redirect_uri"),
		}
		if tr.GrantType == "authorization_code" {
			// Authorization Code Flow https://openid.net/specs/openid-connect-core-1_0.html#CodeFlowAuth
			r := &TokenResponse{
				AccessToken:  tokens[tr.Code],
				TokenType:    "Bearer",
				RefreshToken: "TODO::",
				ExpiresIn:    3600,
				IdToken:      createIdToken(r.FormValue("client_id")),
			}

			w.Header().Set("Content-Type", "application/json")
			w.Header().Set("Cache-Control", "no-store")
			w.Header().Set("Pragma", "no-cache")
			fmt.Fprintf(w, `{"access_token": "%s","token_type": "%s","refresh_token": "%s","expires_in": %d,"id_token": "%s"}`, r.AccessToken, r.TokenType, r.RefreshToken, r.ExpiresIn, r.IdToken)
		} else {
			w.WriteHeader(http.StatusNotFound)
			fmt.Fprintf(w, "Not Found\n")
		}
	} else if r.URL.Path == "/.well-known/openid-configuration" {
		w.Header().Set("Content-Type", "application/json")
		data, _ := ioutil.ReadFile("openid-configuration.json")
		tmpl, _ := template.New("test").Parse(string(data))
		tmpl.Execute(w, struct{ Host string }{Host: "http://localhost:8282"})
	} else if r.URL.Path == "/userinfo" {
		w.WriteHeader(http.StatusNotImplemented)
	} else if r.URL.Path == "/.well-known/jwks.json" {
		w.Header().Set("Content-Type", "application/json")
		keyData, _ := ioutil.ReadFile("insecure.pem")
		privatePem, _ := pem.Decode(keyData)
		parsedKey, _ := x509.ParsePKCS1PrivateKey(privatePem.Bytes)
		key, _ := jwk.FromRaw(parsedKey)
		pubKey, _ := jwk.PublicKeyOf(key)
		pubKey.Set(jwk.KeyIDKey, "X")
		pubKey.Set(jwk.KeyUsageKey, "sig")

		set := jwk.NewSet()
		set.Add(pubKey)
		json.NewEncoder(w).Encode(set)
	} else if r.URL.Path == "/revoke" {
		w.WriteHeader(http.StatusNotImplemented)
	} else {
		w.WriteHeader(http.StatusNotFound)
		fmt.Fprintf(w, "Not Found\n")
	}
}

func main() {
	log.Println("Starting server, listening on port 8282.")

	server := &http.Server{
		Addr:         ":8282",
		Handler:      http.HandlerFunc(handler),
		ReadTimeout:  0,
		WriteTimeout: 0,
		IdleTimeout:  0,
	}

	log.Fatal(server.ListenAndServe())
}