package main

import (
	"context"
	"encoding/base64"
	"encoding/json"
	"errors"
	"fmt"
	"strings"

	"enclave/internal/keybase"

	"github.com/extism/go-pdk"
)

// GenerateInput represents the input for the generate function
type GenerateInput struct {
	Credential string `json:"credential"` // Base64-encoded PublicKeyCredential
}

// GenerateOutput represents the output of the generate function
type GenerateOutput struct {
	DID      string `json:"did"`
	Database []byte `json:"database"`
}

// LoadInput represents the input for the load function
type LoadInput struct {
	Database []byte `json:"database"`
}

// LoadOutput represents the output of the load function
type LoadOutput struct {
	Success bool   `json:"success"`
	DID     string `json:"did,omitempty"`
	Error   string `json:"error,omitempty"`
}

// ExecInput represents the input for the exec function
type ExecInput struct {
	Filter string `json:"filter"` // GitHub-style filter: "resource:accounts action:sign"
	Token  string `json:"token"`  // UCAN token for authorization
}

// ExecOutput represents the output of the exec function
type ExecOutput struct {
	Success bool            `json:"success"`
	Result  json.RawMessage `json:"result,omitempty"`
	Error   string          `json:"error,omitempty"`
}

// QueryInput represents the input for the query function
type QueryInput struct {
	DID string `json:"did"`
}

// QueryOutput represents the output of the query function
type QueryOutput struct {
	DID                 string               `json:"did"`
	Controller          string               `json:"controller"`
	VerificationMethods []VerificationMethod `json:"verification_methods"`
	Accounts            []Account            `json:"accounts"`
	Credentials         []Credential         `json:"credentials"`
}

// VerificationMethod represents a DID verification method
type VerificationMethod struct {
	ID         string `json:"id"`
	Type       string `json:"type"`
	Controller string `json:"controller"`
	PublicKey  string `json:"public_key"`
	Purpose    string `json:"purpose"`
}

// Account represents a derived blockchain account
type Account struct {
	Address      string `json:"address"`
	ChainID      string `json:"chain_id"`
	CoinType     int    `json:"coin_type"`
	AccountIndex int    `json:"account_index"`
	AddressIndex int    `json:"address_index"`
	Label        string `json:"label"`
	IsDefault    bool   `json:"is_default"`
}

// Credential represents a WebAuthn credential
type Credential struct {
	CredentialID  string   `json:"credential_id"`
	DeviceName    string   `json:"device_name"`
	DeviceType    string   `json:"device_type"`
	Authenticator string   `json:"authenticator"`
	Transports    []string `json:"transports"`
	CreatedAt     string   `json:"created_at"`
	LastUsed      string   `json:"last_used"`
}

// FilterParams parsed from GitHub-style filter syntax
type FilterParams struct {
	Resource string
	Action   string
	Subject  string
}

// Enclave holds the plugin state
type Enclave struct {
	initialized bool
	did         string
}

var enclave = &Enclave{}

func main() {}

// PingInput represents the input for the ping function
type PingInput struct {
	Message string `json:"message"`
}

// PingOutput represents the output of the ping function
type PingOutput struct {
	Success bool   `json:"success"`
	Message string `json:"message"`
	Echo    string `json:"echo"`
}

//go:wasmexport ping
func ping() int32 {
	pdk.Log(pdk.LogInfo, "ping: received request")

	var input PingInput
	if err := pdk.InputJSON(&input); err != nil {
		output := PingOutput{
			Success: false,
			Message: fmt.Sprintf("failed to parse input: %s", err),
		}
		pdk.OutputJSON(output)
		return 0
	}

	output := PingOutput{
		Success: true,
		Message: "pong",
		Echo:    input.Message,
	}

	if err := pdk.OutputJSON(output); err != nil {
		pdk.Log(pdk.LogError, fmt.Sprintf("ping: failed to output: %s", err))
		return 1
	}

	pdk.Log(pdk.LogInfo, fmt.Sprintf("ping: responded with echo=%s", input.Message))
	return 0
}

//go:wasmexport generate
func generate() int32 {
	pdk.Log(pdk.LogInfo, "generate: starting database initialization")

	var input GenerateInput
	if err := pdk.InputJSON(&input); err != nil {
		pdk.SetError(fmt.Errorf("generate: failed to parse input: %w", err))
		return 1
	}

	if input.Credential == "" {
		pdk.SetError(errors.New("generate: credential is required"))
		return 1
	}

	credentialBytes, err := base64.StdEncoding.DecodeString(input.Credential)
	if err != nil {
		pdk.SetError(fmt.Errorf("generate: invalid base64 credential: %w", err))
		return 1
	}

	did, err := initializeDatabase(credentialBytes)
	if err != nil {
		pdk.SetError(fmt.Errorf("generate: failed to initialize database: %w", err))
		return 1
	}

	enclave.initialized = true
	enclave.did = did

	dbBytes, err := serializeDatabase()
	if err != nil {
		pdk.SetError(fmt.Errorf("generate: failed to serialize database: %w", err))
		return 1
	}

	output := GenerateOutput{
		DID:      did,
		Database: dbBytes,
	}

	if err := pdk.OutputJSON(output); err != nil {
		pdk.SetError(fmt.Errorf("generate: failed to output result: %w", err))
		return 1
	}

	pdk.Log(pdk.LogInfo, fmt.Sprintf("generate: created DID %s", did))
	return 0
}

//go:wasmexport load
func load() int32 {
	pdk.Log(pdk.LogInfo, "load: loading database from buffer")

	var input LoadInput
	if err := pdk.InputJSON(&input); err != nil {
		pdk.SetError(fmt.Errorf("load: failed to parse input: %w", err))
		return 1
	}

	if len(input.Database) == 0 {
		pdk.SetError(errors.New("load: database buffer is required"))
		return 1
	}

	did, err := loadDatabase(input.Database)
	if err != nil {
		output := LoadOutput{
			Success: false,
			Error:   err.Error(),
		}
		pdk.OutputJSON(output)
		return 1
	}

	enclave.initialized = true
	enclave.did = did

	output := LoadOutput{
		Success: true,
		DID:     did,
	}

	if err := pdk.OutputJSON(output); err != nil {
		pdk.SetError(fmt.Errorf("load: failed to output result: %w", err))
		return 1
	}

	pdk.Log(pdk.LogInfo, fmt.Sprintf("load: loaded database for DID %s", did))
	return 0
}

//go:wasmexport exec
func exec() int32 {
	pdk.Log(pdk.LogInfo, "exec: executing action")

	if !enclave.initialized {
		output := ExecOutput{Success: false, Error: "database not initialized, call generate or load first"}
		pdk.OutputJSON(output)
		return 0
	}

	var input ExecInput
	if err := pdk.InputJSON(&input); err != nil {
		output := ExecOutput{Success: false, Error: fmt.Sprintf("failed to parse input: %s", err)}
		pdk.OutputJSON(output)
		return 0
	}

	if input.Filter == "" {
		output := ExecOutput{Success: false, Error: "filter is required"}
		pdk.OutputJSON(output)
		return 0
	}

	params, err := parseFilter(input.Filter)
	if err != nil {
		output := ExecOutput{Success: false, Error: fmt.Sprintf("invalid filter: %s", err)}
		pdk.OutputJSON(output)
		return 0
	}

	if input.Token != "" {
		if err := validateUCAN(input.Token, params); err != nil {
			output := ExecOutput{
				Success: false,
				Error:   fmt.Sprintf("authorization failed: %s", err.Error()),
			}
			pdk.OutputJSON(output)
			return 1
		}
	}

	result, err := executeAction(params)
	if err != nil {
		output := ExecOutput{
			Success: false,
			Error:   err.Error(),
		}
		pdk.OutputJSON(output)
		return 1
	}

	output := ExecOutput{
		Success: true,
		Result:  result,
	}

	pdk.OutputJSON(output)
	pdk.Log(pdk.LogInfo, fmt.Sprintf("exec: completed %s on %s", params.Action, params.Resource))
	return 0
}

//go:wasmexport query
func query() int32 {
	pdk.Log(pdk.LogInfo, "query: resolving DID document")

	if !enclave.initialized {
		pdk.SetError(errors.New("database not initialized, call generate or load first"))
		return 1
	}

	var input QueryInput
	if err := pdk.InputJSON(&input); err != nil {
		pdk.SetError(fmt.Errorf("query: failed to parse input: %w", err))
		return 1
	}

	if input.DID == "" {
		input.DID = enclave.did
	}

	if !strings.HasPrefix(input.DID, "did:") {
		pdk.SetError(errors.New("query: invalid DID format"))
		return 1
	}

	output, err := resolveDID(input.DID)
	if err != nil {
		pdk.SetError(fmt.Errorf("query: failed to resolve DID: %w", err))
		return 1
	}

	if err := pdk.OutputJSON(output); err != nil {
		pdk.SetError(fmt.Errorf("query: failed to output result: %w", err))
		return 1
	}

	pdk.Log(pdk.LogInfo, fmt.Sprintf("query: resolved DID %s", input.DID))
	return 0
}

func initializeDatabase(credentialBytes []byte) (string, error) {
	kb, err := keybase.Open()
	if err != nil {
		return "", fmt.Errorf("open database: %w", err)
	}

	ctx := context.Background()
	did, err := kb.Initialize(ctx, credentialBytes)
	if err != nil {
		return "", fmt.Errorf("initialize: %w", err)
	}

	pdk.Log(pdk.LogDebug, "initializeDatabase: created schema and initial records")
	return did, nil
}

func serializeDatabase() ([]byte, error) {
	kb := keybase.Get()
	if kb == nil {
		return nil, errors.New("database not initialized")
	}
	return kb.Serialize()
}

func loadDatabase(data []byte) (string, error) {
	if len(data) < 10 {
		return "", errors.New("invalid database format")
	}

	kb, err := keybase.Open()
	if err != nil {
		return "", fmt.Errorf("open database: %w", err)
	}

	ctx := context.Background()
	did, err := kb.Load(ctx, data)
	if err != nil {
		return "", fmt.Errorf("load DID: %w", err)
	}

	pdk.Log(pdk.LogDebug, "loadDatabase: database loaded successfully")
	return did, nil
}

func parseFilter(filter string) (*FilterParams, error) {
	params := &FilterParams{}
	parts := strings.Fields(filter)

	for _, part := range parts {
		kv := strings.SplitN(part, ":", 2)
		if len(kv) != 2 {
			continue
		}

		key, value := kv[0], kv[1]
		switch key {
		case "resource":
			params.Resource = value
		case "action":
			params.Action = value
		case "subject":
			params.Subject = value
		}
	}

	if params.Resource == "" {
		return nil, errors.New("resource is required")
	}
	if params.Action == "" {
		return nil, errors.New("action is required")
	}

	return params, nil
}

func validateUCAN(token string, params *FilterParams) error {
	// TODO: Decode UCAN token
	// TODO: Verify signature chain
	// TODO: Check capabilities match params
	// TODO: Verify not expired or revoked

	if token == "" {
		return errors.New("token is required")
	}

	pdk.Log(pdk.LogDebug, fmt.Sprintf("validateUCAN: validated token for %s:%s", params.Resource, params.Action))
	return nil
}

func executeAction(params *FilterParams) (json.RawMessage, error) {
	// TODO: Route to appropriate handler based on resource/action
	// TODO: Execute database queries
	// TODO: Return results

	switch params.Resource {
	case "accounts":
		return executeAccountAction(params)
	case "credentials":
		return executeCredentialAction(params)
	case "sessions":
		return executeSessionAction(params)
	case "grants":
		return executeGrantAction(params)
	default:
		return nil, fmt.Errorf("unknown resource: %s", params.Resource)
	}
}

func executeAccountAction(params *FilterParams) (json.RawMessage, error) {
	switch params.Action {
	case "list":
		accounts := []Account{
			{
				Address:      "sonr1abc123...",
				ChainID:      "sonr-mainnet-1",
				CoinType:     118,
				AccountIndex: 0,
				AddressIndex: 0,
				Label:        "Primary",
				IsDefault:    true,
			},
		}
		return json.Marshal(accounts)
	case "sign":
		return json.Marshal(map[string]string{"signature": "placeholder"})
	default:
		return nil, fmt.Errorf("unknown action for accounts: %s", params.Action)
	}
}

func executeCredentialAction(params *FilterParams) (json.RawMessage, error) {
	switch params.Action {
	case "list":
		credentials := []Credential{
			{
				CredentialID:  "cred_abc123",
				DeviceName:    "MacBook Pro",
				DeviceType:    "platform",
				Authenticator: "Touch ID",
				Transports:    []string{"internal"},
				CreatedAt:     "2025-01-01T00:00:00Z",
				LastUsed:      "2025-01-01T00:00:00Z",
			},
		}
		return json.Marshal(credentials)
	default:
		return nil, fmt.Errorf("unknown action for credentials: %s", params.Action)
	}
}

func executeSessionAction(params *FilterParams) (json.RawMessage, error) {
	switch params.Action {
	case "list":
		return json.Marshal([]map[string]interface{}{})
	case "create":
		return json.Marshal(map[string]string{"session_id": "sess_placeholder"})
	case "revoke":
		return json.Marshal(map[string]bool{"revoked": true})
	default:
		return nil, fmt.Errorf("unknown action for sessions: %s", params.Action)
	}
}

func executeGrantAction(params *FilterParams) (json.RawMessage, error) {
	switch params.Action {
	case "list":
		return json.Marshal([]map[string]interface{}{})
	case "create":
		return json.Marshal(map[string]string{"grant_id": "grant_placeholder"})
	case "revoke":
		return json.Marshal(map[string]bool{"revoked": true})
	default:
		return nil, fmt.Errorf("unknown action for grants: %s", params.Action)
	}
}

func resolveDID(did string) (*QueryOutput, error) {
	// TODO: Query database for DID document
	// TODO: Fetch verification methods
	// TODO: Fetch associated accounts
	// TODO: Fetch credentials

	output := &QueryOutput{
		DID:        did,
		Controller: did,
		VerificationMethods: []VerificationMethod{
			{
				ID:         did + "#key-1",
				Type:       "Ed25519VerificationKey2020",
				Controller: did,
				PublicKey:  "placeholder_public_key",
				Purpose:    "authentication",
			},
		},
		Accounts: []Account{
			{
				Address:      "sonr1abc123...",
				ChainID:      "sonr-mainnet-1",
				CoinType:     118,
				AccountIndex: 0,
				AddressIndex: 0,
				Label:        "Primary",
				IsDefault:    true,
			},
		},
		Credentials: []Credential{
			{
				CredentialID:  "cred_abc123",
				DeviceName:    "MacBook Pro",
				DeviceType:    "platform",
				Authenticator: "Touch ID",
				Transports:    []string{"internal"},
				CreatedAt:     "2025-01-01T00:00:00Z",
				LastUsed:      "2025-01-01T00:00:00Z",
			},
		},
	}

	return output, nil
}