diff --git a/src/media.php b/src/media.php
index e6af199..c574021 100644
--- a/src/media.php
+++ b/src/media.php
@@ -4,37 +4,6 @@ namespace TheaterGF\Core;
 
 require_once __DIR__ . "/tag.php";
 
-function array_contains( $array, $element ) {
-	return null != array_find($array, function ( $e ) use ($element) { return $e == $element; });
-}
-
-function build_tag_name( $term ) {
-
-	if (is_array( $term )) {
-		$out = [];
-
-		foreach ( $term  as $t ) {
-			$out[] = build_tag_name($t);
-		}
-		sort($out);
-
-		return $out;
-	}
-
-    $parts = [];
-
-    while ( $term && $term->parent != 0 ) {
-        $parts[] = $term->name;
-        $term = get_term($term->parent, 'ttgf_media_tags');
-    }
-
-    if ( $term ) {
-        $parts[] = $term->name;
-    }
-
-    return implode('/', array_reverse($parts));
-}
-
 class MediaManager {
 
 	public $tags;
@@ -45,7 +14,7 @@ class MediaManager {
 
 		add_filter('attachment_fields_to_edit', [ $this, 'create_edit_media_control' ], 10, 2);
 		add_action('add_meta_boxes', [ $this, 'create_tag_management_box' ]);
-		add_action('save_post', [ $this, 'save_tags' ]);
+		add_action('edit_attachment', [ $this, 'save_tags' ]);
 		
 		// Dont know what it does
 		add_filter( 'update_post_term_count_statuses', function( $statuses, $taxonomy ) {
@@ -77,21 +46,29 @@ class MediaManager {
 	}
 
 	public function tag_management_meta_box_html( $post ) {
-		$terms = get_the_terms($post, 'ttgf_media_tags');
-		$tags = $this->tags->get_multiple($terms ? $terms : []);
+		$tags = $this->tags->get_from_post($post);
 ?>
 	<h3>Assigned</h3>
 
 	<div class="ttgf_tags_container">
 		<?php foreach ( $tags as $tag ): ?>
-			<div class="ttgf_tag"><?= $tag->name() ?></div>
-		<?php endforeach; ?>
+			<div class="ttgf_tag"><?= $tag->full_name() ?></div>
+		<?php endforeach ?><br/>
+		<?php foreach ( get_the_terms( $post, 'ttgf_media_tags') ?: [] as $term): ?>
+			<?= $term->name ?> - <?= $term->term_id ?><br/>
+		<?php endforeach ?><br/>
+		<?php foreach ( $this->tags->get_multiple(get_the_terms( $post, 'ttgf_media_tags') ?: []) as $term): ?>
+			<?= $term->full_name() ?> - <?= $term->id ?><br/>
+		<?php endforeach ?>
 	</div>
 
 	<h3>Available</h3>
 	<div>
 		<?php foreach ( $this->tags->get_all() as $tag ): ?>
-			<?= $tag->name() ?>: <?= $tag->id ?><br/>
+			<?= $tag->full_name() ?>: <?= $tag->id ?><br/>
+		<?php endforeach ?><br/>
+		<?php foreach ( get_terms([ 'taxonomy' => 'ttgf_media_tags', 'hide_empty' => false ]) as $term): ?>
+			<?= $term->name ?> - <?= $term->term_id ?><br/>
 		<?php endforeach ?>
 	</div>
 
@@ -101,7 +78,7 @@ class MediaManager {
 		<input type="text" list="ttgf_available_tags" name="ttgf_new_tag">
 		<datalist id="ttgf_available_tags">
 			<?php foreach ( $this->tags->get_all() as $tag ): ?>
-				<option value="<?= $tag->name() ?>">
+				<option value="<?= $tag->full_name() ?>">
 			<?php endforeach ?>
 		</datalist>
 	</div>
@@ -109,8 +86,8 @@ class MediaManager {
 	}
 
 	public function save_tags( $post ) {
-		$is_autosave = wp_is_post_autosave( $post_id );
-		$is_revision = wp_is_post_revision( $post_id );
+		$is_autosave = wp_is_post_autosave( $post );
+		$is_revision = wp_is_post_revision( $post );
 		$is_valid_nonce = isset($_POST['ttgf_core_admin_nonce']) && wp_verify_nonce($_POST['ttgf_core_admin_nonce'], basename(__FILE__));
  
 		if ( $is_autosave || $is_revision || !$is_valid_nonce ) {
@@ -118,35 +95,15 @@ class MediaManager {
 		}
 		
 		$new_tag = sanitize_text_field($_POST['ttgf_new_tag']);
-		if ( isset($new_tag) ) {	
-
-			$all_terms = get_terms([
-				'taxonomy' => 'ttgf_media_tags',
-				'hide_empty' => false,
-			]);
-			$all_tags = build_tag_name($all_terms);
-
-			$terms = get_the_terms($post, 'ttgf_media_tags');
-			$tags = build_tag_name($terms);
+		if ( isset($new_tag) ) {
 			
-			if ( array_find($tags, function ( $tag ) { return $tag == $new_tag; }) ) {
-				return;
-			}
+			$tag = $this->tags->get_by_name($new_tag);
 
-			if ( ! array_find($all_tags, function ( $tag ) { return $tag == $new_tag; }) ) {
-
-				$parts = explode("/", $new_tag);
-				
-				$index = count($parts);
-				for ( ; $index > 0; --$index ) {
-					if ( array_contains($all_tags, implode("/", array_slice($parts, 0, $index))) ) {
-						break;
-					}
-				}
-				
-				
-				
+			if ( ! $tag ) {
+				$tag = $this->tags->create_tag($new_tag);
 			}
+			
+			wp_set_post_terms($post, [ $tag->id ], 'ttgf_media_tags', true);
 		}
 	}
 
diff --git a/src/tag.php b/src/tag.php
index f8d5f36..3e2ddec 100644
--- a/src/tag.php
+++ b/src/tag.php
@@ -7,6 +7,7 @@ class Tag {
 	public $term;
 
 	public $id;
+	public $name;
 
 	public $parent;
 	public $children;
@@ -14,11 +15,12 @@ class Tag {
 	public function __construct( $term ) {
 		$this->term      = $term;
 		$this->id        = $term->term_id;
+		$this->name      = $term->name;
 		$this->parent    = null;
 		$this->children  = [];
 	}
 
-	public function name() {
+	public function full_name() {
 		return implode('/', array_map(function ($segment) { return $segment->term->name; }, $this->path()));
 	}
 
@@ -137,7 +139,16 @@ class TagTree {
 		$term_ids = array_map(function ($term) { return $term->term_id; }, $terms);
 		$term_ids_set = array_flip($term_ids);
 
+		$tags = [];
+
 		foreach ( $this->root_tags as $rt ) {
+			if ( isset($term_ids_set[$rt->id]) ) {
+				$tags[] = $rt;
+			}
+			if ( count($tags) == $num_terms ) {
+				return $tags;
+			}
+
 			foreach ( $rt->get_all_children() as $child ) {
 				if ( isset($term_ids_set[$child->id]) ) {
 					$tags[] = $child;
@@ -163,17 +174,92 @@ class TagTree {
 		return null;
 	}
 
-	public function create_tag( $parent, $tag_name, $description = '') {
-		$new_term = wp_insert_term( $tag_name, $this->taxonomy, [
-			'description' => $description,
-			'parent'      => $parent->id
-		]);
+	public function get_from_post( $post ) {
+		$terms = get_the_terms($post, $this->taxonomy);
+		return $this->get_multiple($terms ?: []);
+	}
 
-		if ( is_wp_error($new_term) ) {
-			return $new_term;
+	public function get_by_name( $name, $return_path = false ) {
+		$segments = array_filter(explode('/', trim($name, '/')));
+
+		$node = null;
+		$path = [];
+
+		foreach ( $segments as $segment ) {
+			$current_nodes = $node?->children ?: $this->root_tags;
+
+			$node = array_find($current_nodes, fn ($tag) => $tag->name === $segment);
+
+			if ( ! $node ) {
+				break;
+			}
+
+			$path[] = $node;
 		}
 
-		return new Tag($new_term, $parent);
+		if ( $return_path ) {
+			return $path;
+		}
+
+		if ( count($path) === count($segments) ) {
+			return $node;
+		}
+
+		return null;	
+	}
+
+	public function create_tag( $tag_name, $description = '') {	
+		$segments = array_filter(explode('/', trim($tag_name, '/')));
+
+		if ( empty($segments) ) {
+			return new WP_Error('invalid_tag', 'Tag name is empty');
+		}
+
+		$path = $this->get_by_name($tag_name, true);
+		
+		if ( count($segments) === count($path) ) {
+			return end($path);
+		}
+
+		$parent_tag = ( ! empty($path) ) ? end($path) : null;
+
+		for ( $i = count($path); $i < count($segments); ++$i ) {
+			$segment = $segments[$i];
+
+			$args = [];
+
+			if ( $parent_tag ) {
+				$args['parent'] = $parent_tag->id;
+			}
+
+			if ( $description && $i === count($segments) - 1 ) {
+				$args['description'] = $description;
+			}
+
+			$result = wp_insert_term($segment, $this->taxonomy, $args);
+
+			if ( is_wp_error($result) ) {
+				return $result;
+			}
+
+			$term = get_term($result['term_id'], $this->taxonomy);
+
+			if ( is_wp_error($term) || !$term ) {
+				return new WP_Error('term_fetch_failed', 'Failed to retrieve created term');
+			}
+
+			$tag = new Tag($term);
+
+			if ( $parent_tag ) {
+				$parent_tag->add_child($tag);
+			} else {
+				$this->root_tags[] = $tag;
+			}
+
+			$parent_tag = $tag;
+		}
+
+		return $parent_tag;
 	}
 
 	public function move_tag( $tag, $new_parent ) {