package visitor

import (
	"fmt"
	"io"
	"os"
)

type MessageA struct {
	Msg    string
	Output io.Writer
}

func (m *MessageA) Accept(v Visitor) {
	//nothing
	v.VisitA(m)
}

func (m *MessageA) Print() {
	//nothing
	if m.Output == nil {
		m.Output = os.Stdout
	}
	fmt.Printf("A: %s\n", m.Msg)
	fmt.Fprintf(m.Output, "A: %s", m.Msg)
}

type MessageB struct {
	Msg    string
	Output io.Writer
}

func (m *MessageB) Accept(v Visitor) {
	//nothing
	v.VisitB(m)
}

func (m *MessageB) Print() {
	//nothing
	if m.Output == nil {
		m.Output = os.Stdout
	}
	fmt.Printf("B: %s\n", m.Msg)
	fmt.Fprintf(m.Output, "B: %s", m.Msg)
}

type Visitor interface {
	VisitA(*MessageA)
	VisitB(*MessageB)
}

type Visitable interface {
	Accept(Visitor)
}

type MessageVisitor struct{}

func (mf *MessageVisitor) VisitA(m *MessageA) {
	//nothing
	m.Msg = fmt.Sprintf("%s %s", m.Msg, "(Visited A)")
}

func (mf *MessageVisitor) VisitB(m *MessageB) {
	//nothing
	m.Msg = fmt.Sprintf("%s %s", m.Msg, "(Visited B)")
}

package visitor

import (
	"fmt"
	"testing"
)

type TestHelper struct {
	Received string
}

func (t TestHelper) Write(p []byte) (int, error) {
	t.Received = string(p)
	return len(p), nil
}

func Test_Overall(t *testing.T) {
	testHelper := &TestHelper{}
	visitor := &MessageVisitor{}

	t.Run("MessageA test", func(t *testing.T) {
		msg := MessageA{
			Msg:    "Hello World",
			Output: testHelper,
		}

		msg.Accept(visitor)
		msg.Print()

		fmt.Print(testHelper)

		expected := "A: Hello World (Visited A)"
		if testHelper.Received != expected {
			t.Errorf("Expected result was incorrect. %s != %s", testHelper.Received, expected)
		}
	})

	t.Run("MessageB test", func(t *testing.T) {
		msg := MessageB{
			Msg:    "Hello World",
			Output: testHelper,
		}

		msg.Accept(visitor)
		msg.Print()

		expected := "B: Hello World (Visited B)"
		if testHelper.Received != expected {
			t.Errorf("Expected result was incorrect. %s != %s", testHelper.Received, expected)
		}
	})
}