diff --git a/static/style.css b/static/style.css
index cc6862d..a9d2258 100644
--- a/static/style.css
+++ b/static/style.css
@@ -7,9 +7,12 @@ html {
 #chat-messages .message-content pre {
     overflow-x: auto;
     white-space: pre;
-    max-width: 662px;
     position: relative;
-    border-radius: 8px;
+    max-width: 685px;
+    border-radius: 0px 0px 8px 8px;
+    background-color: #2c2c2c;
+    margin-top: 0;
+    /* Remove top margin to align with header */
 }
 
 #chat-messages .message-content pre code {
@@ -18,10 +21,24 @@ html {
     white-space: inherit;
 }
 
+.code-header {
+    display: flex;
+    justify-content: space-between;
+    align-items: center;
+    background-color: #515151;
+    border-radius: 8px 8px 0 0;
+    /* Rounded corners on top */
+    padding: 5px 10px 10px 10px;
+    margin-bottom: -8px;
+    /* Overlap with pre element */
+}
+
 .copy-button {
-    position: absolute;
-    top: 5px;
-    right: 5px;
+    background: none;
+    border: none;
+    cursor: pointer;
+    font-size: 14px;
+    color: #ffffff;
 }
 
 .content {
@@ -30,6 +47,7 @@ html {
     color: #ffffff;
 }
 
+
 /* Style for the overall chat container */
 
 #chat-messages {
diff --git a/utils.go b/utils.go
index 8703545..94a6761 100644
--- a/utils.go
+++ b/utils.go
@@ -10,34 +10,69 @@ import (
 
 	"github.com/yuin/goldmark"
 	highlighting "github.com/yuin/goldmark-highlighting"
+	"github.com/yuin/goldmark/parser"
+	"github.com/yuin/goldmark/renderer/html"
 )
 
+var md goldmark.Markdown
+
+func init() {
+	md = goldmark.New(goldmark.WithExtensions(highlighting.NewHighlighting(highlighting.WithStyle("monokai"))),
+		goldmark.WithParserOptions(
+			parser.WithAutoHeadingID(),
+		),
+		goldmark.WithRendererOptions(
+			html.WithHardWraps(),
+			html.WithXHTML(),
+		))
+}
+
 func markdownToHTML(markdownText string) string {
 	var buf bytes.Buffer
-	md := goldmark.New(
-		goldmark.WithExtensions(
-			highlighting.NewHighlighting(highlighting.WithStyle("monokai")),
-		),
-	)
+
 	if err := md.Convert([]byte(markdownText), &buf); err != nil {
 		fmt.Println("failed to convert markdown to HTML")
 		panic(err)
 	}
 
-	return addCopyButtonsToCode(buf.String())
+	return addCodeHeader(buf.String(), extractLanguages(markdownText))
 }
 
-func addCopyButtonsToCode(htmlContent string) string {
-	buttonHTML := `<button class="copy-button button is-small is-primary is-outlined" onclick="copyToClipboardCode(this)"><i class="fa-solid fa-copy"></i></button>`
+func extractLanguages(markdownText string) []string {
+	// Regular expression to match code blocks and capture the language identifier
+	re := regexp.MustCompile("(?s)```([a-zA-Z0-9]*)\n.*?```")
 
-	// Regular expression pattern to match <pre> elements and insert the button right before <code>
+	matches := re.FindAllStringSubmatch(markdownText, -1)
+	languages := make([]string, 0, len(matches))
+
+	for _, match := range matches {
+		if len(match) > 1 {
+			languages = append(languages, match[1])
+		} else {
+			languages = append(languages, "")
+		}
+	}
+
+	return languages
+}
+
+func addCodeHeader(htmlContent string, languages []string) string {
+	// Regular expression pattern to match <pre> elements
 	pattern := `(<pre[^>]*>)`
 
 	// Compile the regular expression
 	re := regexp.MustCompile(pattern)
 
 	// Replace each matched <pre> element with the updated HTML
-	updatedHTML := re.ReplaceAllString(htmlContent, "$1"+buttonHTML)
+	updatedHTML := re.ReplaceAllStringFunc(htmlContent, func(match string) string {
+		var language string
+		if len(languages) > 0 {
+			language = languages[0]
+			languages = languages[1:]
+		}
+		headerHTML := fmt.Sprintf(`<div class="code-header"><span>%s</span><button class="copy-button" onclick="copyToClipboardCode(this)"><i class="fa-solid fa-copy"></i></button></div>`, language)
+		return headerHTML + match
+	})
 
 	return updatedHTML
 }
diff --git a/views/layouts/main.html b/views/layouts/main.html
index 52bfdb8..06f329f 100644
--- a/views/layouts/main.html
+++ b/views/layouts/main.html
@@ -26,7 +26,7 @@
     <script>
         function copyToClipboardCode(button) {
             // Get the code element next to the button
-            var codeElement = button.nextElementSibling;
+            var codeElement = button.parentElement.nextElementSibling;
 
             // Create a temporary textarea element
             var tempTextarea = document.createElement('textarea');