From f4272cbc0f272cdc9921ed2159c18eda341348c3 Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?Do=C4=9Fan=20Can=20Bak=C4=B1r?=
 <dogancanbakir@protonmail.com>
Date: Mon, 16 Feb 2026 13:33:42 +0300
Subject: [PATCH] refactor: lazy init dit classifier, prefer headless body for
 classification

---
 runner/runner.go | 20 +++++++++++++-------
 1 file changed, 13 insertions(+), 7 deletions(-)

diff --git a/runner/runner.go b/runner/runner.go
index da37983..7df5c1c 100644
--- a/runner/runner.go
+++ b/runner/runner.go
@@ -430,11 +430,13 @@ func New(options *Options) (*Runner, error) {
 	}
 
 	runner.simHashes = gcache.New[uint64, struct{}](1000).ARC().Build()
-	ditClassifier, err := dit.New()
-	if err != nil {
-		gologger.Warning().Msgf("Could not initialize page classifier: %s", err)
+	if options.JSONOutput || options.CSVOutput || len(options.OutputFilterPageType) > 0 {
+		ditClassifier, err := dit.New()
+		if err != nil {
+			gologger.Warning().Msgf("Could not initialize page classifier: %s", err)
+		}
+		runner.ditClassifier = ditClassifier
 	}
-	runner.ditClassifier = ditClassifier
 
 	if options.SecretFile != "" {
 		authProviderOpts := &authprovider.AuthProviderOptions{
@@ -652,12 +654,16 @@ func (r *Runner) duplicate(result *Result) bool {
 	return false
 }
 
-func (r *Runner) classifyPage(body string, pHash uint64) map[string]interface{} {
+func (r *Runner) classifyPage(headlessBody, body string, pHash uint64) map[string]any {
 	kb := map[string]any{"pHash": pHash}
 	if r.ditClassifier == nil {
 		return kb
 	}
-	result, err := r.ditClassifier.ExtractPageType(body)
+	html := body
+	if headlessBody != "" {
+		html = headlessBody
+	}
+	result, err := r.ditClassifier.ExtractPageType(html)
 	if err != nil {
 		return kb
 	}
@@ -2636,7 +2642,7 @@ retry:
 		ExtractRegex:     extractRegex,
 		ScreenshotBytes:  screenshotBytes,
 		HeadlessBody:     headlessBody,
-		KnowledgeBase: r.classifyPage(respData, pHash),
+		KnowledgeBase: r.classifyPage(headlessBody, respData, pHash),
 		TechnologyDetails: technologyDetails,
 		Resolvers:         resolvers,
 		RequestRaw:        requestDump,