/**
 * Tau Client Extension
 * 
 * Lightweight client that connects to pi-tau-mux-server.
 * - Registers this pi instance with the mux server
 * - Forwards events (messages, model changes, etc.) to the server
 * - Unregisters on shutdown
 * - Optionally prompts to start the mux server if not running
 */

import type { ExtensionAPI, ExtensionContext } from "@mariozechner/pi-coding-agent";
import { WebSocket } from "ws";
import * as fs from "node:fs";
import * as path from "node:path";
import { exec } from "node:child_process";

// Configuration
const MUX_PORT = parseInt(process.env.TAU_PORT || "3001");
const MUX_HOST = process.env.TAU_HOST || "localhost";
const MUX_URL = `ws://${MUX_HOST}:${MUX_PORT}/pi`;
const AUTO_CONNECT = process.env.TAU_AUTO_CONNECT !== "0";

let ws: WebSocket | null = null;
let latestCtx: ExtensionContext | null = null;
let reconnectTimer: NodeJS.Timeout | null = null;
let sessionId: string | null = null;
let serverUrls: { local?: string; tailscale?: string; magicDns?: string } | null = null;

// ─────────────────────────────────────────────────────────────
// Connection management
// ─────────────────────────────────────────────────────────────

function connect(pi: ExtensionAPI, ctx: ExtensionContext, promptIfDown = true) {
  if (ws) return; // Already connected
  
  try {
    ws = new WebSocket(MUX_URL);
    
    ws.on("open", () => {
      // Register this instance
      sessionId = ctx.sessionManager.getSessionFile()?.split("/").pop()?.replace(".jsonl", "") || Date.now().toString();
      send({ type: "register", sessionId, cwd: ctx.cwd || process.cwd() });
      
      // Send initial state
      sendState(ctx, pi);
      
      // Notify user of successful connection
      // Use server's Tailscale URL if available, otherwise fall back to connection URL
      const displayUrl = serverUrls?.tailscale || serverUrls?.magicDns || `http://${MUX_HOST}:${MUX_PORT}`;
      const location = serverUrls?.tailscale ? 'Tailscale' : serverUrls?.magicDns ? 'MagicDNS' : (MUX_HOST === 'localhost' ? 'local' : MUX_HOST);
      ctx.ui.notify(`Connected to Tau mux server at ${location}  •  View UI at ${displayUrl}`, "info");
    });
    
    ws.on("message", (data) => {
      try {
        const msg = JSON.parse(data.toString());
        handleMessage(msg, ctx, pi);
      } catch {}
    });
    
    ws.on("close", () => {
      ws = null;
      
      // Attempt reconnect after delay
      if (!reconnectTimer) {
        reconnectTimer = setTimeout(() => {
          reconnectTimer = null;
          if (latestCtx) connect(pi, latestCtx, false);
        }, 5000);
      }
    });
    
    ws.on("error", (e) => {
      // Prompt to start server if not running
      if (promptIfDown && e.message.includes("ECONNREFUSED")) {
        promptStartMux(pi, ctx);
      }
      
      ws = null;
    });
  } catch {
    // Connection failed silently
  }
}

function disconnect() {
  if (ws) {
    if (sessionId) {
      send({ type: "unregister", sessionId });
    }
    ws.close();
    ws = null;
  }
  
  if (reconnectTimer) {
    clearTimeout(reconnectTimer);
    reconnectTimer = null;
  }
}

function send(msg: any) {
  if (ws && ws.readyState === WebSocket.OPEN) {
    ws.send(JSON.stringify(msg));
  }
}

function sendState(ctx: ExtensionContext, pi: ExtensionAPI) {
  if (!sessionId) return;
  
  const entries = ctx.sessionManager.getEntries();
  const model = ctx.model;
  const thinkingLevel = pi.getThinkingLevel();
  const sessionName = pi.getSessionName();
  const sessionFile = ctx.sessionManager.getSessionFile();
  
  send({
    type: "state",
    sessionId,
    state: {
      entries,
      model,
      thinkingLevel,
      sessionName,
      sessionFile,
      isStreaming: !ctx.isIdle(),
    }
  });
}

function promptStartMux(pi: ExtensionAPI, ctx: ExtensionContext) {
  ctx.ui.confirm(
    "Tau mux server not running. Start it now?",
    "Yes, start server",
    "No, skip"
  ).then((response) => {
    if (response === "Yes, start server") {
      exec("npx pi-tau-mux-server", (err) => {
        if (err) {
          ctx.ui.notify("Failed to start tau mux server: " + err.message, "error");
        }
      });
      
      // Wait a moment for server to start, then connect
      setTimeout(() => connect(pi, ctx, false), 2000);
    }
  });
}

function handleMessage(msg: any, ctx: ExtensionContext, pi: ExtensionAPI) {
  // Handle commands from mux server (e.g., from browser)
  switch (msg.type) {
    case "registered":
      // Store server URLs from registration response
      serverUrls = msg.serverUrls || null;
      break;
    
    case "prompt":
      pi.sendUserMessage(msg.message);
      break;
    
    case "set_model":
      pi.setModel(msg.model);
      break;
    
    case "set_thinking_level":
      pi.setThinkingLevel(msg.level);
      break;
    
    // Add more command handlers as needed
  }
}

// ─────────────────────────────────────────────────────────────
// Event forwarding
// ─────────────────────────────────────────────────────────────

const EVENT_TYPES = [
  "agent_start", "agent_end",
  "turn_start", "turn_end",
  "message_start", "message_update", "message_end",
  "tool_execution_start", "tool_execution_update", "tool_execution_end",
  "auto_compaction_start", "auto_compaction_end",
  "auto_retry_start", "auto_retry_end",
  "model_select",
] as const;

// ─────────────────────────────────────────────────────────────
// Extension entry point
// ─────────────────────────────────────────────────────────────

export default function (pi: ExtensionAPI) {
  // Register commands
  pi.registerCommand("tauconnect", {
    description: "Connect to tau mux server",
    handler: async (_args, ctx) => {
      if (ws) {
        ctx.ui.notify("Already connected to tau mux", "warning");
        return;
      }
      connect(pi, ctx, true);
    },
  });
  
  pi.registerCommand("taudisconnect", {
    description: "Disconnect from tau mux server",
    handler: async (_args, ctx) => {
      disconnect();
      ctx.ui.notify("Disconnected from tau mux", "info");
    },
  });
  
  // Auto-connect on session start
  pi.on("session_start", async (_event, ctx) => {
    latestCtx = ctx;
    
    if (AUTO_CONNECT) {
      ctx.ui.notify(`Connecting to Tau mux server...`, "info");
      connect(pi, ctx, true);
    }
  });
  
  // Forward events
  for (const eventType of EVENT_TYPES) {
    pi.on(eventType as any, async (event: any, ctx: ExtensionContext) => {
      latestCtx = ctx;
      send({ type: "event", sessionId, event: { type: eventType, ...event } });
    });
  }
  
  // Update state periodically
  pi.on("turn_end", async (_event, ctx) => {
    sendState(ctx, pi);
  });
  
  // Cleanup on shutdown
  pi.on("session_shutdown", async () => {
    disconnect();
  });
}