import fs from 'fs-extra';
import path from 'path';
import { Parser, Language } from 'web-tree-sitter';
import chalk from 'chalk';
import os from 'os';
import { fileURLToPath } from 'url';

const __dirname = path.dirname(fileURLToPath(import.meta.url));
// Go up two levels from dist/languages/ to get to the root where web-tree-sitter.wasm is
const PACKAGE_ROOT = path.resolve(__dirname, '../../');
const USER_CACHE_DIR = path.join(os.homedir(), '.owstra', 'grammars');
const PACKAGE_GRAMMARS_DIR = path.join(PACKAGE_ROOT, 'grammars');

export class WASMLoader {
  private static languageCache: Map<string, Language> = new Map();
  private static globalLoadLock: Promise<any> = Promise.resolve();
  private static isInitialized = false;

  static async init() {
    if (this.isInitialized) return;

    let wasmPath = path.join(PACKAGE_ROOT, 'web-tree-sitter.wasm');
    
    if (!fs.existsSync(wasmPath)) {
      // Fallback: Check if we are in dist/ and assets are in the true root
      const fallbackPath = path.resolve(PACKAGE_ROOT, '../web-tree-sitter.wasm');
      if (fs.existsSync(fallbackPath)) {
        wasmPath = fallbackPath;
      } else {
        throw new Error(`Critical: web-tree-sitter.wasm not found at ${wasmPath} or ${fallbackPath}`);
      }
    }

    // Parser.init() must be called once. We point it to the local web-tree-sitter.wasm
    try {
      await Parser.init({
        locateFile(filePath: string) {
          // Only redirect the core WASM file, not the grammars
          if (filePath.includes('web-tree-sitter.wasm')) {
            return wasmPath;
          }
          return filePath;
        }
      });
      this.isInitialized = true;
    } catch (e: any) {
      throw new Error(`Failed to initialize Parser with wasmPath: ${wasmPath}. Error: ${e.message}`);
    }
  }

  static async getLanguage(lang: string): Promise<Language> {
    const normalizedLang = this.normalizeLanguageName(lang);
    
    // Check in-memory cache first
    if (this.languageCache.has(normalizedLang)) {
      return this.languageCache.get(normalizedLang)!;
    }

    // Use a global promise chain to ensure serial loading of all grammars.
    // This is the most robust way to prevent concurrent WASM module loads.
    // We append a catch to ensure that a failure in one grammar doesn't block others.
    const promise = this.globalLoadLock.catch(() => {}).then(async () => {
      // Check cache again in case it was populated while we were waiting in the queue
      if (this.languageCache.has(normalizedLang)) {
        return this.languageCache.get(normalizedLang)!;
      }

      const fileName = `tree-sitter-${normalizedLang}.wasm`;
      const userCachePath = path.join(USER_CACHE_DIR, fileName);
      const packagePath = path.join(PACKAGE_GRAMMARS_DIR, fileName);

      let finalWasmPath = '';

      if (fs.existsSync(packagePath)) {
        finalWasmPath = packagePath;
      } else if (fs.existsSync(userCachePath)) {
        finalWasmPath = userCachePath;
      } else {
        // Fallback: Check legacy .shield directory
        const legacyDir = path.join(os.homedir(), '.shield', 'grammars');
        const legacyPath = path.join(legacyDir, fileName);
        if (fs.existsSync(legacyPath)) {
          await fs.ensureDir(USER_CACHE_DIR);
          await fs.copy(legacyPath, userCachePath);
          finalWasmPath = userCachePath;
        } else {
          console.log(chalk.yellow(`\n⬇️  Grammar for ${chalk.bold(normalizedLang)} not found locally.`));
          console.log(chalk.gray(`   Attempting to download from official sources...`));
          await this.downloadGrammar(normalizedLang, userCachePath);
          finalWasmPath = userCachePath;
        }
      }

      try {
        const language = await Language.load(finalWasmPath);
        this.languageCache.set(normalizedLang, language);
        return language;
      } catch (e: any) {
        // If it failed to load, it might be an incompatible ABI or corrupt file.
        // If it's in the user cache, we can try to redownload it once.
        if (finalWasmPath === userCachePath) {
          console.log(chalk.yellow(`   ⚠️  Cached grammar for ${chalk.bold(normalizedLang)} is incompatible. Attempting to redownload...`));
          await fs.remove(userCachePath);
          await this.downloadGrammar(normalizedLang, userCachePath);
          try {
            const language = await Language.load(userCachePath);
            this.languageCache.set(normalizedLang, language);
            return language;
          } catch (retryError: any) {
            throw new Error(`Failed to load grammar for ${normalizedLang} even after redownload: ${retryError.message}`);
          }
        }
        throw new Error(`Failed to load grammar for ${normalizedLang} from ${finalWasmPath}: ${e.message}`);
      }
    });

    this.globalLoadLock = promise;
    return promise;
  }

  private static normalizeLanguageName(lang: string): string {
    const maps: Record<string, string> = {
      ts: 'typescript',
      js: 'javascript',
      py: 'python',
      rb: 'ruby',
      kt: 'kotlin',
      dart: 'dart',
      go: 'go',
      java: 'java',
      cs: 'csharp',
      swift: 'swift',
      m: 'objc',
      h: 'objc',
    };
    return maps[lang.toLowerCase()] || lang.toLowerCase();
  }

  private static async downloadGrammar(lang: string, dest: string) {
    await fs.ensureDir(USER_CACHE_DIR);
    
    // Normalize language name for common CDN patterns
    const variations = [lang];
    if (lang === 'csharp') variations.push('c-sharp', 'c_sharp');
    if (lang === 'typescript') variations.push('typescript');

    const urls: string[] = [];
    for (const v of variations) {
      urls.push(`https://unpkg.com/tree-sitter-${v}@latest/tree-sitter-${v}.wasm`);
      // Special case for C# which often uses underscore in the WASM filename
      if (v === 'c-sharp' || v === 'c_sharp') {
        urls.push(`https://unpkg.com/tree-sitter-c-sharp@latest/tree-sitter-c_sharp.wasm`);
      }
      urls.push(`https://unpkg.com/@tree-sitter/tree-sitter-${v}@latest/tree-sitter-${v}.wasm`);
      urls.push(`https://unpkg.com/tree-sitter-wasms@latest/out/tree-sitter-${v}.wasm`);
    }

    let lastError: Error | null = null;

    for (const url of urls) {
      try {
        const res = await fetch(url);
        if (!res.ok) continue;

        const buffer = await res.arrayBuffer();
        if (buffer.byteLength < 5000) continue; // Likely a small error page or empty file

        await fs.writeFile(dest, Buffer.from(buffer));
        const sizeKB = Math.round(buffer.byteLength / 1024);
        console.log(chalk.green(`   ✅ Successfully downloaded ${chalk.bold(lang)} grammar from ${url.includes('@tree-sitter') ? '@tree-sitter/' : ''}${path.basename(url)} (${sizeKB} KB).`));
        return;
      } catch (e: any) {
        lastError = e;
      }
    }
    
    console.log(chalk.red(`   ❌ Failed to download ${lang} grammar.`));
    throw new Error(`Failed to download grammar for ${lang}. Last error: ${lastError?.message}`);
  }
}
