// AGPL-3.0-only
//
// This is free software for the public good of a permacomputer
// hosted at permacomputer.com, an always-on computer by the people, for the
// people. One which is durable, easy to repair, and distributed like tap water
// for machine learning intelligence.
//
// The permacomputer is community-owned infrastructure optimized around
// four values:
//
//   TRUTH    - First principles, math & science, open source code
//   FREEDOM  - Voluntary partnerships, freedom from tyranny & corporate control
//   HARMONY  - Minimal waste, self-renewing systems with diverse connections
//   LOVE     - Be yourself without hurting others, cooperation through natural law
//
// This software contributes to that vision by making machine learning
// accessible to everyone through a free, open, embeddable chat interface.
// Code is seeds to sprout on any abandoned technology.
//
// Learn more: https://www.permacomputer.com
//
// NO WARRANTY. THE SOFTWARE IS PROVIDED "AS IS" WITHOUT WARRANTY OF ANY KIND.
//
// Copyright 2025-2026 TimeHexOn & foxhop & russell@unturf
// https://www.permacomputer.com

// Simplified BPE tokenizer based on tiktoken's core logic
// MIT License - Derived from tiktoken (https://github.com/dqbd/tiktoken)
// Copyright (c) 2023 David Bezdek

// Core regex pattern simplified from cl100k_base encoding
// Original tiktoken pattern: r"""'(?i:[sdmt]|ll|ve|re)|[^\r\n\p{L}\p{N}]?+\p{L}++|\p{N}{1,3}+| ?[^\s\p{L}\p{N}]++[\r\n]*+|\s++$|\s*[\r\n]|\s+(?!\S)|\s"""

export class SimpleBPETokenizer {
	constructor() {
		this.initializePatterns();
		this.initializeCommonTokens();
	}

	initializePatterns() {
		// Simplified tokenization patterns based on tiktoken's approach
		this.patterns = [
			// Contractions (case insensitive)
			/'(?:s|t|m|d|ll|ve|re)/gi,

			// Letters (word boundaries)
			/[a-zA-Z]+/g,

			// Numbers (1-3 digits groups)
			/\d{1,3}/g,

			// Punctuation and symbols
			/[^\s\w]/g,

			// Whitespace (including newlines)
			/\s+/g,
		];

		// Compiled regex for faster processing
		this.mainPattern = /'(?:s|t|m|d|ll|ve|re)|[a-zA-Z]+|\d{1,3}|[^\s\w]|\s+/gi;
	}

	initializeCommonTokens() {
		// Common subword patterns from BPE training (simplified set)
		this.commonSubwords = new Set([
			// Common prefixes
			"un",
			"re",
			"in",
			"dis",
			"en",
			"non",
			"pre",
			"over",
			"mis",
			"sub",
			"anti",
			"auto",
			"co",
			"de",
			"ex",
			"inter",
			"multi",
			"out",
			"under",

			// Common suffixes
			"ing",
			"ed",
			"er",
			"est",
			"ly",
			"tion",
			"ness",
			"ment",
			"able",
			"ible",
			"ful",
			"less",
			"ship",
			"ward",
			"wise",
			"like",
			"ous",
			"ive",
			"ate",

			// Common words (high frequency)
			"the",
			"and",
			"for",
			"are",
			"but",
			"not",
			"you",
			"all",
			"can",
			"had",
			"her",
			"was",
			"one",
			"our",
			"out",
			"day",
			"get",
			"use",
			"man",
			"new",
			"now",
			"way",
			"may",
			"say",
			"see",
			"him",
			"two",
			"how",
			"its",
			"who",

			// Common letter combinations
			"th",
			"er",
			"on",
			"an",
			"re",
			"he",
			"in",
			"ed",
			"nd",
			"ha",
			"at",
			"en",
			"es",
			"of",
			"or",
			"nt",
			"ea",
			"ti",
			"to",
			"it",
			"st",
			"io",
		]);

		// Token length estimates for common patterns
		this.tokenWeights = {
			word: 1,
			number: 0.8,
			punctuation: 1,
			whitespace: 0.3,
			contraction: 0.5,
			subword: 0.6,
		};
	}

	// Main tokenization method
	encode(text) {
		if (!text || typeof text !== "string") return [];

		// Split text using main pattern
		const rawTokens = text.match(this.mainPattern) || [];
		const processedTokens = [];

		for (const token of rawTokens) {
			processedTokens.push(...this.processToken(token));
		}

		return processedTokens;
	}

	processToken(token) {
		const trimmed = token.trim();
		if (!trimmed) return [token]; // Keep whitespace as-is

		// Check token type and split if needed
		if (this.isWord(trimmed)) {
			return this.splitWord(trimmed);
		}
		if (this.isNumber(trimmed)) {
			return this.splitNumber(trimmed);
		}
		return [token]; // Punctuation, symbols, etc.
	}

	isWord(token) {
		return /^[a-zA-Z]+$/.test(token);
	}

	isNumber(token) {
		return /^\d+$/.test(token);
	}

	splitWord(word) {
		if (word.length <= 3) return [word];

		const tokens = [];
		let remaining = word.toLowerCase();

		// Try to find common subwords
		while (remaining.length > 0) {
			let found = false;

			// Look for longest matching subword first
			for (let len = Math.min(remaining.length, 6); len >= 2; len--) {
				const substr = remaining.substring(0, len);
				if (this.commonSubwords.has(substr)) {
					tokens.push(substr);
					remaining = remaining.substring(len);
					found = true;
					break;
				}
			}

			if (!found) {
				// Split remaining into chunks
				if (remaining.length <= 4) {
					tokens.push(remaining);
					break;
				}
				// Take first 3-4 characters
				const chunkSize = Math.min(4, Math.ceil(remaining.length / 2));
				tokens.push(remaining.substring(0, chunkSize));
				remaining = remaining.substring(chunkSize);
			}
		}

		return tokens;
	}

	splitNumber(number) {
		// Split numbers into groups of 1-3 digits
		const chunks = [];
		for (let i = 0; i < number.length; i += 3) {
			chunks.push(number.substring(i, i + 3));
		}
		return chunks;
	}

	// Quick token count (main method for estimation)
	count(text) {
		if (!text || typeof text !== "string") return 0;

		// Handle special content first
		const specialTokens = this.countSpecialContent(text);

		// Remove special content and count regular tokens
		const cleanText = this.removeSpecialContent(text);
		const regularTokens = this.encode(cleanText);

		return {
			totalTokens: specialTokens.count + regularTokens.length,
			regularTokens: regularTokens.length,
			specialTokens: specialTokens.count,
			breakdown: {
				...specialTokens.breakdown,
				words: regularTokens.filter((t) => this.isWord(t.trim())).length,
				numbers: regularTokens.filter((t) => this.isNumber(t.trim())).length,
				punctuation: regularTokens.filter((t) => /^[^\s\w]+$/.test(t.trim()))
					.length,
				whitespace: regularTokens.filter((t) => /^\s+$/.test(t)).length,
			},
			textLength: text.length,
		};
	}

	countSpecialContent(text) {
		let count = 0;
		const breakdown = { htmlTags: 0, codeBlocks: 0, inlineCode: 0, urls: 0 };

		// HTML tags
		const htmlMatches = text.match(/<[^>]+>/g) || [];
		breakdown.htmlTags = htmlMatches.length;
		count += htmlMatches.length * 3; // ~3 tokens per tag

		// Code blocks
		const codeBlockMatches = text.match(/```[\s\S]*?```/g) || [];
		breakdown.codeBlocks = codeBlockMatches.length;
		codeBlockMatches.forEach((block) => {
			const content = block.replace(/```[\w]*\n?/g, "").replace(/```$/g, "");
			count += Math.ceil(content.length / 2.5); // Code is denser
		});

		// Inline code
		const inlineMatches = text.match(/`[^`]+`/g) || [];
		breakdown.inlineCode = inlineMatches.length;
		inlineMatches.forEach((code) => {
			count += Math.ceil(code.length / 3);
		});

		// URLs
		const urlMatches = text.match(/https?:\/\/[^\s<>"{}|\\^`[\]]+/g) || [];
		breakdown.urls = urlMatches.length;
		count += urlMatches.length * 4; // URLs are typically 2-6 tokens

		return { count, breakdown };
	}

	removeSpecialContent(text) {
		return text
			.replace(/```[\s\S]*?```/g, " ") // Replace code blocks
			.replace(/`[^`]+`/g, " ") // Replace inline code
			.replace(/<[^>]+>/g, " ") // Replace HTML tags
			.replace(/https?:\/\/[^\s<>"{}|\\^`[\]]+/g, " "); // Replace URLs
	}
}

// Export singleton instance
export const tokenizer = new SimpleBPETokenizer();

// Cache endpoints that don't support tokenization
const failedTokenizeEndpoints = new Set();

// Use upstream tokenizer API for accurate token counting
export async function countTokensWithUpstream(text, model) {
	const { API_KEY } = await import("./config.js");
	const { getSelectedModel, getSelectedModelEndpoint } = await import("./models.js");
	
	try {
		const endpoint = getSelectedModelEndpoint().replace('/v1', '');
		
		// Skip if this endpoint previously failed tokenization
		if (failedTokenizeEndpoints.has(endpoint)) {
			console.log(`Endpoint ${endpoint} previously failed tokenization, using JS tokenizer`);
			return tokenizer.count(text);
		}
		
		console.log(`Trying tokenization at: ${endpoint}/tokenize`);
		
		const response = await fetch(`${endpoint}/tokenize`, {
			method: "POST",
			headers: {
				"Content-Type": "application/json",
				Authorization: `Bearer ${API_KEY}`,
			},
			body: JSON.stringify({
				model: model || getSelectedModel(),
				prompt: text,
				add_special_tokens: true,
			}),
		});

		if (!response.ok) {
			console.warn(`Upstream tokenizer failed (${response.status}) at ${endpoint}/tokenize, falling back to JS tokenizer`);
			// Cache this endpoint as failed for future requests
			failedTokenizeEndpoints.add(endpoint);
			return tokenizer.count(text);
		}

		const tokens = await response.json();
		console.log("Tokenizer response:", tokens);
		// Simply count the length of the token array
		const tokenCount = Array.isArray(tokens) ? tokens.length : (tokens.count || Object.keys(tokens).length);
		
		console.log(`Upstream tokenizer: ${tokenCount.toLocaleString()} tokens for ${text.length.toLocaleString()} chars`);
		
		// Return in same format as JS tokenizer for compatibility
		return {
			totalTokens: tokenCount,
			regularTokens: tokenCount,
			specialTokens: 0,
			breakdown: {
				words: Math.ceil(tokenCount * 0.7),
				numbers: Math.ceil(tokenCount * 0.1),
				punctuation: Math.ceil(tokenCount * 0.1),
				whitespace: Math.ceil(tokenCount * 0.1),
			},
			textLength: text.length,
		};
	} catch (error) {
		console.warn("Error with upstream tokenizer, falling back to JS tokenizer:", error);
		// Cache this endpoint as failed for future requests
		const endpoint = getSelectedModelEndpoint().replace('/v1', '');
		failedTokenizeEndpoints.add(endpoint);
		return tokenizer.count(text);
	}
}

// Get just the token count as a number from upstream
export async function getTokenCount(text, model) {
	const { API_KEY } = await import("./config.js");
	const { getSelectedModel, getSelectedModelEndpoint } = await import("./models.js");
	
	try {
		const endpoint = getSelectedModelEndpoint().replace('/v1', '');
		
		// Skip if this endpoint previously failed tokenization
		if (failedTokenizeEndpoints.has(endpoint)) {
			console.log(`Endpoint ${endpoint} previously failed tokenization, using JS tokenizer`);
			return tokenizer.count(text).totalTokens;
		}
		
		console.log(`Trying tokenization at: ${endpoint}/tokenize`);
		
		const response = await fetch(`${endpoint}/tokenize`, {
			method: "POST",
			headers: {
				"Content-Type": "application/json",
				Authorization: `Bearer ${API_KEY}`,
			},
			body: JSON.stringify({
				model: model || getSelectedModel(),
				prompt: text,
				add_special_tokens: true,
			}),
		});

		if (!response.ok) {
			// Cache this endpoint as failed for future requests
			failedTokenizeEndpoints.add(endpoint);
			throw new Error(`Tokenizer API error: ${response.status} at ${endpoint}/tokenize`);
		}

		const tokens = await response.json();
		console.log("Tokenizer response:", tokens);
		return Array.isArray(tokens) ? tokens.length : (tokens.count || Object.keys(tokens).length);
	} catch (error) {
		console.warn("Upstream tokenizer error:", error);
		// Cache this endpoint as failed for future requests
		const endpoint = getSelectedModelEndpoint().replace('/v1', '');
		failedTokenizeEndpoints.add(endpoint);
		// Fallback to JS estimate
		return tokenizer.count(text).totalTokens;
	}
}

// Convenience functions
export function countTokens(text) {
	return tokenizer.count(text);
}

export function estimateTokens(text) {
	const result = tokenizer.count(text);
	return result.totalTokens;
}

// Format token count as a string with commas (e.g., "8,432 tokens")
export function formatTokenCount(count) {
	return `${count.toLocaleString()} tokens`;
}
