import { Bot, Context } from "grammy";
import Anthropic from "@anthropic-ai/sdk";
import { readFileSync } from "fs";
import { resolve } from "path";

// ─── Config ───
const BOT_TOKEN = process.env.TELEGRAM_BOT_TOKEN!;
const ANTHROPIC_KEY = process.env.ANTHROPIC_API_KEY!;
const SSH_USER = process.env.SSH_USER || "root";
const SSH_HOST = process.env.SSH_HOST!;
const SSH_PORT = process.env.SSH_PORT || "22";
const ALLOWED_USERS = process.env.TELEGRAM_ALLOWED_USERS?.split(",").map(Number) || [];

if (!BOT_TOKEN) throw new Error("TELEGRAM_BOT_TOKEN missing from .env");
if (!ANTHROPIC_KEY) throw new Error("ANTHROPIC_API_KEY missing from .env");
if (!SSH_HOST) throw new Error("SSH_HOST missing from .env");

// ─── Load infra context ───
let infraContext = "";
try {
  infraContext = readFileSync(resolve(__dirname, "../.pi/infra.md"), "utf-8");
} catch {
  console.warn("⚠️  Could not load .pi/infra.md");
}

// ─── Init ───
const bot = new Bot(BOT_TOKEN);
const anthropic = new Anthropic({ apiKey: ANTHROPIC_KEY });

// Per-user conversation history
const conversations = new Map<number, Anthropic.MessageParam[]>();
const MAX_HISTORY = 30;

// ─── Auth middleware ───
function isAuthorized(ctx: Context): boolean {
  // Reject if no whitelist configured (fail-closed)
  if (ALLOWED_USERS.length === 0) return false;
  return ALLOWED_USERS.includes(ctx.from?.id || 0);
}

// ─── Rate limiting ───
const rateLimits = new Map<number, number[]>();
const RATE_LIMIT_WINDOW = 60_000; // 1 minute
const RATE_LIMIT_MAX = 15; // max messages per window

function isRateLimited(userId: number): boolean {
  const now = Date.now();
  const timestamps = (rateLimits.get(userId) || []).filter(t => now - t < RATE_LIMIT_WINDOW);
  timestamps.push(now);
  rateLimits.set(userId, timestamps);
  return timestamps.length > RATE_LIMIT_MAX;
}

// ─── Blocked command patterns ───
const BLOCKED_PATTERNS = [
  /rm\s+-rf\s+\/(?!\w)/i,   // rm -rf /
  /mkfs\./i,                  // format disk
  /dd\s+if=.*of=\/dev/i,     // overwrite disk
  /:\(\)\s*\{\s*:\|:/,       // fork bomb
];

function isDangerous(cmd: string): string | null {
  for (const pattern of BLOCKED_PATTERNS) {
    if (pattern.test(cmd)) return `⛔ Blocked: matches dangerous pattern \`${pattern.source}\``;
  }
  return null;
}

// ─── SSH exec helper ───
async function sshExec(command: string, timeoutSec = 30): Promise<string> {
  const proc = Bun.spawn(
    ["ssh", "-o", "ConnectTimeout=5", "-o", "StrictHostKeyChecking=accept-new",
     "-p", SSH_PORT, `${SSH_USER}@${SSH_HOST}`, command],
    { stdout: "pipe", stderr: "pipe" }
  );

  const timer = setTimeout(() => proc.kill(), timeoutSec * 1000);
  const [stdout, stderr] = await Promise.all([
    new Response(proc.stdout).text(),
    new Response(proc.stderr).text(),
  ]);
  clearTimeout(timer);
  await proc.exited;

  const output = (stdout + stderr).trim();
  return output.length > 3500 ? output.slice(0, 3500) + "\n...(truncated)" : output;
}

// ─── System prompt ───
const SYSTEM_PROMPT = `You are a DevOps assistant connected to a live server via Telegram.
You have access to run SSH commands on the server. You prioritize uptime.

Infrastructure context:
${infraContext}

RULES:
- If the user asks to run a command, use the run_command tool.
- For destructive commands (rm -rf, DROP, reboot, etc.), WARN the user and ask for confirmation before executing.
- Keep responses concise — this is Telegram, not a terminal.
- You can chain multiple commands in one SSH call using && or ;
- Format output with monospace blocks for readability.`;

// ─── Claude tools ───
const tools: Anthropic.Tool[] = [
  {
    name: "run_command",
    description: "Execute a command on the server via SSH. Use for any server inspection, management, or debugging task.",
    input_schema: {
      type: "object" as const,
      properties: {
        command: { type: "string", description: "Shell command to run on the server" },
        timeout: { type: "number", description: "Timeout in seconds (default 30)" },
      },
      required: ["command"],
    },
  },
];

// ─── Process message through Claude ───
async function chat(userId: number, message: string): Promise<string> {
  // Get or init conversation history
  if (!conversations.has(userId)) conversations.set(userId, []);
  const history = conversations.get(userId)!;

  history.push({ role: "user", content: message });

  // Trim history if too long
  while (history.length > MAX_HISTORY) history.shift();

  let response = await anthropic.messages.create({
    model: "claude-sonnet-4-20250514",
    max_tokens: 1024,
    system: SYSTEM_PROMPT,
    tools,
    messages: history,
  });

  // Tool use loop
  while (response.stop_reason === "tool_use") {
    const assistantContent = response.content;
    history.push({ role: "assistant", content: assistantContent });

    const toolResults: Anthropic.ToolResultBlockParam[] = [];

    for (const block of assistantContent) {
      if (block.type === "tool_use") {
        const input = block.input as { command: string; timeout?: number };
        console.log(`🔧 [${userId}] run_command: ${input.command}`);
        const danger = isDangerous(input.command);
        if (danger) {
          toolResults.push({ type: "tool_result", tool_use_id: block.id, content: danger, is_error: true });
          continue;
        }
        try {
          const output = await sshExec(input.command, input.timeout || 30);
          toolResults.push({ type: "tool_result", tool_use_id: block.id, content: output || "(no output)" });
        } catch (err: any) {
          toolResults.push({ type: "tool_result", tool_use_id: block.id, content: `ERROR: ${err.message}`, is_error: true });
        }
      }
    }

    history.push({ role: "user", content: toolResults });

    response = await anthropic.messages.create({
      model: "claude-sonnet-4-20250514",
      max_tokens: 1024,
      system: SYSTEM_PROMPT,
      tools,
      messages: history,
    });
  }

  // Extract text response
  const text = response.content
    .filter((b): b is Anthropic.TextBlock => b.type === "text")
    .map((b) => b.text)
    .join("\n");

  history.push({ role: "assistant", content: text });

  return text || "(no response)";
}

// ─── Security: block groups, unauthorized, rate-limited ───
bot.use(async (ctx, next) => {
  // Block group/channel messages — private only
  if (ctx.chat?.type !== "private") return;

  // Silent block for unauthorized users — no response, no leak
  if (!isAuthorized(ctx)) {
    console.log(`🚫 Blocked user: ${ctx.from?.id} (@${ctx.from?.username})`);
    return;
  }

  // Rate limit
  if (ctx.from && isRateLimited(ctx.from.id)) {
    return ctx.reply("⏳ Slow down — rate limited.");
  }

  await next();
});

// ─── Bot handlers ───
bot.command("start", (ctx) => {
  return ctx.reply(
    "🤖 *CR Management Bot*\n\n" +
    "I'm connected to your server. You can:\n" +
    "• Ask me anything about your infra\n" +
    "• Ask me to run commands\n" +
    "• `/status` — quick server health\n" +
    "• `/clear` — reset conversation\n" +
    "• `/ssh <command>` — run a command directly\n",
    { parse_mode: "Markdown" }
  );
});

bot.command("status", async (ctx) => {
  await ctx.reply("🔍 Checking server...");
  try {
    const output = await sshExec("uptime && echo '---' && df -h / && echo '---' && free -h | head -2 && echo '---' && incus list --format csv -c ns 2>/dev/null");
    return ctx.reply(`\`\`\`\n${output}\n\`\`\``, { parse_mode: "Markdown" });
  } catch (err: any) {
    return ctx.reply(`❌ SSH failed: ${err.message}`);
  }
});

bot.command("ssh", async (ctx) => {
  const cmd = ctx.match;
  if (!cmd) return ctx.reply("Usage: `/ssh <command>`", { parse_mode: "Markdown" });
  const danger = isDangerous(cmd);
  if (danger) return ctx.reply(danger);
  await ctx.reply(`⚡ Running: \`${cmd}\``, { parse_mode: "Markdown" });
  try {
    const output = await sshExec(cmd);
    return ctx.reply(`\`\`\`\n${output || "(no output)"}\n\`\`\``, { parse_mode: "Markdown" });
  } catch (err: any) {
    return ctx.reply(`❌ ${err.message}`);
  }
});

bot.command("clear", (ctx) => {
  conversations.delete(ctx.from!.id);
  return ctx.reply("🧹 Conversation cleared.");
});

// Free text → Claude
bot.on("message:text", async (ctx) => {
  const userId = ctx.from!.id;
  const msg = ctx.message!.text;

  await ctx.replyWithChatAction("typing");

  try {
    const reply = await chat(userId, msg);
    console.log(`📤 [${userId}] reply length: ${reply.length}`);
    // Split long messages (Telegram 4096 char limit)
    const chunks = reply.match(/[\s\S]{1,4000}/g) || ["(empty)"];
    for (const chunk of chunks) {
      try {
        await ctx.reply(chunk);
      } catch (sendErr: any) {
        console.error(`❌ [${userId}] send failed:`, sendErr.message);
        await ctx.reply("⚠️ Failed to format response. Check logs.");
      }
    }
  } catch (err: any) {
    console.error(`❌ [${userId}]`, err);
    await ctx.reply(`❌ Error: ${err.message}`);
  }
});

// ─── Launch ───
console.log("🚀 Bot starting...");
bot.start({
  onStart: (info) => console.log(`✅ @${info.username} is live`),
});