#!/usr/bin/env node

import * as http from "http";
import * as https from "https";
import * as child_process from "child_process";
import * as fs from "fs";
import * as path from "path";
import * as os from "os";
import { pipeline, env } from "@xenova/transformers";

// Disable local model downloading if we only want to fetch from HF, 
// or let it download once and cache locally (preferred).
env.allowLocalModels = true; 

const pkgPath = path.join(__dirname, "..", "package.json");
const CURRENT_VERSION = fs.existsSync(pkgPath) ? JSON.parse(fs.readFileSync(pkgPath, "utf-8")).version : "1.0.4";

function parseArgs() {
  const args = process.argv.slice(2);
  let token = "";
  let host = "https://aclade.com";

  for (let i = 0; i < args.length; i++) {
    if (args[i] === "--token" && args[i + 1]) {
      token = args[i + 1];
      i++;
    } else if (args[i] === "--host" && args[i + 1]) {
      host = args[i + 1];
      i++;
    } else if (args[i] === "connect") {
      // ignore
    }
  }
  
  const configPath = path.join(os.homedir(), ".aclade-agent.json");
  if (!token && fs.existsSync(configPath)) {
    try {
      const config = JSON.parse(fs.readFileSync(configPath, "utf-8"));
      if (config.token) token = config.token;
      if (config.host) host = config.host;
    } catch {}
  } else if (token) {
    try {
      fs.writeFileSync(configPath, JSON.stringify({ token, host }));
    } catch {}
  }
  
  return { token, host };
}

function makeRequest(urlStr: string, options: any, bodyData?: string): Promise<string> {
  return new Promise((resolve, reject) => {
    const url = new URL(urlStr);
    const client = url.protocol === "https:" ? https : http;
    const req = client.request(url, options, (res) => {
      let data = "";
      res.on("data", (chunk) => (data += chunk));
      res.on("end", () => {
        if (res.statusCode && res.statusCode >= 200 && res.statusCode < 300) {
          resolve(data);
        } else {
          reject(new Error(`HTTP Error ${res.statusCode}: ${data}`));
        }
      });
    });
    req.on("error", reject);
    if (bodyData) {
      req.write(bodyData);
    }
    req.end();
  });
}

// Global ML model cache
let embedder: any = null;
async function getEmbedder() {
  if (!embedder) {
    console.log("=> Initializing ONNX Vector Embedding Engine (Downloading Xenova/all-MiniLM-L6-v2 on first run...)");
    embedder = await pipeline('feature-extraction', 'Xenova/all-MiniLM-L6-v2');
    console.log("=> Embedding Engine Ready.");
  }
  return embedder;
}

async function executeTask(task: any): Promise<any> {
  const { toolName, input } = task;
  
  if (toolName === "execute_bash") {
    return new Promise((resolve) => {
      console.log(`=> Running command: ${input.command}`);
      // Use spawn for streaming, chunked execution so it doesn't crash on huge stdout
      const proc = child_process.spawn(input.command, [], {
        shell: true,
        stdio: ['ignore', 'pipe', 'pipe']
      });

      let stdout = "";
      let stderr = "";
      let isTruncated = false;
      const MAX_LEN = 1024 * 1024 * 2; // 2MB max buffer

      proc.stdout.on("data", (chunk) => {
        if (stdout.length < MAX_LEN) stdout += chunk.toString();
        else isTruncated = true;
      });

      proc.stderr.on("data", (chunk) => {
        if (stderr.length < MAX_LEN) stderr += chunk.toString();
        else isTruncated = true;
      });

      proc.on("close", (code) => {
        if (isTruncated) {
          stdout += "\n...[OUTPUT TRUNCATED: EXCEEDED 2MB BUFFER]...";
          stderr += "\n...[OUTPUT TRUNCATED: EXCEEDED 2MB BUFFER]...";
        }
        resolve({
          output: stdout || "",
          error: stderr || "",
          exitCode: code
        });
      });
      
      proc.on("error", (err) => {
        resolve({ error: err.message });
      });
    });
  } else if (toolName === "read_file") {
    try {
      const content = fs.readFileSync(input.path, "utf-8");
      return { output: content };
    } catch (e: any) {
      return { error: e.message };
    }
  } else if (toolName === "write_file") {
    try {
      fs.writeFileSync(input.path, input.content, "utf-8");
      return { output: "File written successfully." };
    } catch (e: any) {
      return { error: e.message };
    }
  } else if (toolName === "embed_texts") {
    try {
      const ext = await getEmbedder();
      const results = [];
      for (const text of input.texts) {
        const out = await ext(text, { pooling: 'mean', normalize: true });
        results.push(Array.from(out.data));
      }
      return { output: results };
    } catch (e: any) {
      return { error: e.message };
    }
  } else if (toolName === "list_directory_recursive") {
    try {
      const execSync = child_process.execSync;
      let out = "";
      if (process.platform === "win32") {
        out = execSync(`dir "${input.path}" /s /b`, { encoding: 'utf-8', maxBuffer: 1024*1024*10 });
      } else {
        out = execSync(`find "${input.path}" -type f`, { encoding: 'utf-8', maxBuffer: 1024*1024*10 });
      }
      return { output: out.slice(0, 500000) }; // cap output at 500kb
    } catch(e:any) {
      return { error: e.message };
    }
  } else {
    return { error: `Unknown tool: ${toolName}` };
  }
}

async function pollLoop(token: string, host: string) {
  console.log(`=> Connected to ${host}. Waiting for instructions from Aclade Cloud...`);
  
  let lastUpdateCheck = 0; // Check immediately on first loop

  while (true) {
    try {
      if (Date.now() - lastUpdateCheck > 1000 * 60 * 60) {
        lastUpdateCheck = Date.now();
        // Fire-and-forget update check so it doesn't block the poll loop
        new Promise<string>((resolve, reject) => {
          https.get("https://registry.npmjs.org/aclade-agent/latest", (resp) => {
            let data = "";
            resp.on("data", (chunk) => data += chunk);
            resp.on("end", () => resolve(data));
          }).on("error", reject);
        }).then((res) => {
          const latest = JSON.parse(res).version;
          if (latest && latest !== CURRENT_VERSION) {
            console.log(`\n=> [UPDATE] New version ${latest} detected (current: ${CURRENT_VERSION}). Upgrading automatically...`);
            try {
              child_process.execSync("npm install -g aclade-agent@latest", { stdio: "ignore" });
              console.log(`=> [UPDATE] Upgrade complete. Restarting agent...`);
              process.exit(0);
            } catch (e) {
              console.error(`=> [UPDATE] Failed to upgrade automatically. Please run: npm install -g aclade-agent@latest`);
            }
          }
        }).catch(() => {});
      }

      const res = await makeRequest(`${host}/api/connector/poll?token=${encodeURIComponent(token)}`, {
        method: "GET",
        headers: { "Accept": "application/json" }
      });
      const data = JSON.parse(res);
      
      if (data.task) {
        console.log(`\n=> [TASK RECEIVED] ${data.task.toolName}`);
        if (data.task.toolName !== "embed_texts") {
          console.log(data.task.input);
        } else {
          console.log(`=> Embedding ${data.task.input.texts?.length || 0} chunks...`);
        }
        
        const result = await executeTask(data.task);
        
        console.log(`=> [TASK COMPLETE] Sending result back...`);
        await makeRequest(`${host}/api/connector/respond`, {
          method: "POST",
          headers: {
            "Content-Type": "application/json",
            "Authorization": `Bearer ${token}`
          }
        }, JSON.stringify({ taskId: data.task.id, result }));
      } else {
        // No task, sleep for a bit
        await new Promise((r) => setTimeout(r, 1500));
      }
    } catch (e: any) {
      if (e.message.includes("401")) {
        console.error("=> Authentication failed. Invalid token.");
        process.exit(1);
      }
      // Silently retry on network errors
      await new Promise((r) => setTimeout(r, 3000));
    }
  }
}

function main() {
  const { token, host } = parseArgs();
  if (!token) {
    console.error("Usage: aclade-agent connect --token <YOUR_TOKEN>");
    process.exit(1);
  }

  if (!process.env.ACLADE_AGENT_DAEMON) {
    console.log("=> Starting aclade-agent in background daemon mode...");
    const logFile = fs.openSync(path.join(os.homedir(), ".aclade-agent.log"), "a");
    const child = child_process.spawn(process.execPath, [process.argv[1], ...process.argv.slice(2)], {
      detached: true,
      stdio: ["ignore", logFile, logFile],
      env: { ...process.env, ACLADE_AGENT_DAEMON: "1" }
    });
    child.unref();
    console.log("=> Daemon started successfully! It will securely wait for tasks in the background.");
    console.log("=> You can safely close this terminal window.");
    process.exit(0);
  } else {
    pollLoop(token, host);
  }
}

main();