#include <vector>
#include <iostream>
#include <string>
using namespace std;
typedef const char* String;

class Iterator;	//Forard Declaration

class Component {
public:
    String name;
    Component(String name) :name(name) { }
    virtual void Add(Component* c) = 0;
    virtual void Remove(Component* c) = 0;
    virtual void Display(int depth = 1) = 0;
    virtual void Register(vector<Component*>& vec) = 0;
    virtual ~Component() { }
};

class Composite :public Component {
    vector<Component*> children;
public:
    Iterator* CreateIterator();
    Composite(String name) :Component(name) { }
    ~Composite() {
        for (auto child : children) delete child;
    }
    void Add(Component* component) {
        children.push_back(component);
    }
    void Remove(Component* component) {
        for (int i = 0; i < children.size(); i++) {
            if (children[i] == component) {
                children.erase(children.begin() + i);
                break;
            }
        }
    }
    void Display(int depth = 1) {
        for (int i = 0; i < depth; i++) cout << '-';
        cout << name << endl;
        for (auto child : children) child->Display(depth + 2);
    }
    void Register(vector<Component*>& vec) {
        vec.push_back(this);
        for (auto child : children) child->Register(vec);
    }
};

class Leaf :public Component {
public:
    Leaf(String name) :Component(name) {}

    void Add(Component* c) {}

    void Remove(Component* c) {}

    void Display(int depth = 1) {
        for (int i = 0; i < depth; i++) cout << '-';
        cout << name << endl;
    }
    ~Leaf() {
        //cout<<name<<" is destroyed.."<<endl;
    }
    void Register(vector<Component*>& vec) {
        vec.push_back(this);
    }
};

class Iterator {
public:
    virtual Component* Next() = 0;
    virtual bool IsDone() = 0;
    virtual Component* CurrentItem() = 0;
};

class TreeIterator : public Iterator {
    Composite* tree;
    int current;
    vector<Component*> elements;
public:
    TreeIterator(Composite* tree) :tree(tree), current(0) {
        tree->Register(elements);
    }
    Component* Next() {
        return (current < elements.size()) ? elements[current++] : 0;
    }
    Component* CurrentItem() {
        return elements[current];
    }
    bool IsDone() {
        return (current == elements.size());
    }
};

Iterator* Composite::CreateIterator() {
    return new TreeIterator(this);
}

int main() {
    // Create a tree structure 
    Composite* tree = new Composite("root");//Beware that I am using Composite* here
                                            //Only composite needs iterator, not component in general
    tree->Add(new Leaf("A"));
    tree->Add(new Leaf("B"));

    Composite* comp = new Composite("X");
    comp->Add(new Leaf("XA"));
    comp->Add(new Leaf("XB"));
    //Leaf* zz = new Leaf("ZZ");
    //comp->Add(zz);
    //comp->Remove(zz);
    //delete zz;

    tree->Add(comp);
    tree->Add(new Leaf("C"));
    //tree->Display();

    Iterator* i = tree->CreateIterator();
    Component* element = i->Next();
    while (element) {
        cout << element->name << '\t';
        element = i->Next();
    }
    cout << endl;
    delete i;
    delete tree;
    return 0;
}