import { serve } from "@hono/node-server";
import { serveStatic } from "@hono/node-server/serve-static";
import {
  CopilotRuntime,
  ExperimentalEmptyAdapter,
  copilotRuntimeNextJSAppRouterEndpoint,
} from "@copilotkit/runtime";
import { HttpAgent, Middleware, EventType } from "@ag-ui/client";
import { Observable } from "rxjs";
import { Hono } from "hono";

const BACKEND_URL =
  process.env.BACKEND_URL ?? "http://localhost:8001";
// Prod sets AGENT_URL explicitly (compose network); the default derives from
// BACKEND_URL so local dev reaches the docker backend published on :8001.
const AGENT_URL =
  process.env.AGENT_URL ?? `${BACKEND_URL}/api/v1/agent/agui`;

const isProd = process.env.NODE_ENV === "production";
const PORT = parseInt(process.env.PORT ?? (isProd ? "3000" : "3001"));

// ── pairOrphanToolCalls ──────────────────────────────────────────────────────
// CopilotKit's browser-side store sometimes drops the tool-role message
// between turns, producing assistant(toolCalls=[X]) → assistant(text) which
// OpenAI rejects. Inject a synthetic empty tool response for each orphan id.

interface AGUIMessage {
  id?: string;
  role?: string;
  content?: unknown;
  toolCalls?: Array<{ id?: string }>;
  tool_calls?: Array<{ id?: string }>;
  toolCallId?: string;
  tool_call_id?: string;
}

function pairOrphanToolCalls(messages: AGUIMessage[]): AGUIMessage[] {
  const out: AGUIMessage[] = [];
  let pending: string[] = [];

  const flush = () => {
    for (const callId of pending) {
      out.push({ id: `synth-${Math.random().toString(36).slice(2)}`, role: "tool", toolCallId: callId, content: "" });
    }
    pending = [];
  };

  for (const msg of messages) {
    const role = String(msg?.role ?? "");
    if (role === "tool") {
      const callId = msg.toolCallId ?? msg.tool_call_id;
      if (callId) pending = pending.filter((p) => p !== callId);
      out.push(msg);
      continue;
    }
    if (role === "assistant") {
      flush();
      const toolCalls = msg.toolCalls ?? msg.tool_calls ?? [];
      out.push(msg);
      for (const tc of toolCalls) {
        if (tc?.id) pending.push(String(tc.id));
      }
      continue;
    }
    flush();
    out.push(msg);
  }
  flush();
  return out;
}

// ── DeduplicateToolCallMiddleware ─────────────────────────────────────────────
// MAF re-emits TOOL_CALL_START for declaration-only calls when they are invoked
// in parallel with other tools, causing verifyEvents to throw an AGUIError.
// This middleware tracks completed tool calls and silently drops any duplicate
// START/ARGS/END events for a call ID that has already been closed.
//
// NOTE: runNextWithState emits { event, messages, state } wrappers; always
// extract `.event` and forward the raw BaseEvent — never the wrapper.

// ── DebugLogMiddleware ──────────────────────────────────────────────────────
// Temporary: logs every event type flowing out to diagnose double-response bug.

// eslint-disable-next-line @typescript-eslint/no-explicit-any
class DebugLogMiddleware extends (Middleware as any) {
  // eslint-disable-next-line @typescript-eslint/no-explicit-any
  run(input: any, next: any): Observable<any> {
    // eslint-disable-next-line @typescript-eslint/no-explicit-any
    return new Observable<any>((observer) => {
      // eslint-disable-next-line @typescript-eslint/no-explicit-any
      const sub = (this as any).runNextWithState(input, next).subscribe({
        // eslint-disable-next-line @typescript-eslint/no-explicit-any
        next: (eventWithState: any) => {
          const event = eventWithState.event;
          const type: string = event?.type ?? "";
          const extra = event?.toolCallName ? ` [${event.toolCallName}]`
            : event?.activityType ? ` [${event.activityType}]`
            : event?.messageId ? ` [msg:${event.messageId.slice(0, 8)}]`
            : "";
          console.log(`[DBG] ${type}${extra}`);
          if (type === EventType.MESSAGES_SNAPSHOT) {
            const msgs = (event as any)?.messages ?? [];
            console.log(`[DBG] SNAPSHOT msgs: ${msgs.map((m: any) => `${m.role}:${(m.id ?? "").slice(0,8)}:${typeof m.content === "string" ? m.content.length : "?"}c`).join(" | ")}`);
          }
          observer.next(event);
        },
        // eslint-disable-next-line @typescript-eslint/no-explicit-any
        error: (err: any) => observer.error(err),
        complete: () => observer.complete(),
      });
      return () => sub.unsubscribe();
    });
  }
}

// eslint-disable-next-line @typescript-eslint/no-explicit-any
class DeduplicateToolCallMiddleware extends (Middleware as any) {
  // eslint-disable-next-line @typescript-eslint/no-explicit-any
  run(input: any, next: any): Observable<any> {
    const open = new Set<string>();   // tool call IDs currently in progress
    const closed = new Set<string>(); // tool call IDs that already received END

    // eslint-disable-next-line @typescript-eslint/no-explicit-any
    return new Observable<any>((observer) => {
      // eslint-disable-next-line @typescript-eslint/no-explicit-any
      const sub = (this as any).runNextWithState(input, next).subscribe({
        // eslint-disable-next-line @typescript-eslint/no-explicit-any
        next: (eventWithState: any) => {
          const event = eventWithState.event; // raw BaseEvent
          const type: string = event?.type ?? "";
          const id: string | undefined = event?.toolCallId;

          if (type === EventType.TOOL_CALL_START) {
            if (!id || closed.has(id) || open.has(id)) return; // duplicate
            open.add(id);
          } else if (type === EventType.TOOL_CALL_ARGS || type === EventType.TOOL_CALL_END) {
            if (id && closed.has(id)) return; // already completed, drop
            if (type === EventType.TOOL_CALL_END && id) {
              open.delete(id);
              closed.add(id);
            }
          }

          observer.next(event); // emit raw BaseEvent
        },
        // eslint-disable-next-line @typescript-eslint/no-explicit-any
        error: (err: any) => observer.error(err),
        complete: () => observer.complete(),
      });
      return () => sub.unsubscribe();
    });
  }
}

// ── StripModelArtifactsMiddleware ────────────────────────────────────────────
// Some OpenAI models occasionally leak training-data special tokens such as
// `<|ipynb_marker|>` into completion text. Filter them out of streamed deltas
// before they reach the chat UI.

const ARTIFACT_RE = /<\|[^|>]+\|>/g;

// eslint-disable-next-line @typescript-eslint/no-explicit-any
class StripModelArtifactsMiddleware extends (Middleware as any) {
  // eslint-disable-next-line @typescript-eslint/no-explicit-any
  run(input: any, next: any): Observable<any> {
    // eslint-disable-next-line @typescript-eslint/no-explicit-any
    return new Observable<any>((observer) => {
      // eslint-disable-next-line @typescript-eslint/no-explicit-any
      const sub = (this as any).runNextWithState(input, next).subscribe({
        // eslint-disable-next-line @typescript-eslint/no-explicit-any
        next: (eventWithState: any) => {
          const event = eventWithState.event;
          if (
            event?.type === EventType.TEXT_MESSAGE_CONTENT &&
            typeof event?.delta === "string" &&
            ARTIFACT_RE.test(event.delta)
          ) {
            const cleaned = event.delta.replace(ARTIFACT_RE, "");
            if (cleaned.length === 0) return;
            observer.next({ ...event, delta: cleaned });
            return;
          }
          observer.next(event);
        },
        // eslint-disable-next-line @typescript-eslint/no-explicit-any
        error: (err: any) => observer.error(err),
        complete: () => observer.complete(),
      });
      return () => sub.unsubscribe();
    });
  }
}

// ── SuppressRenderToolTextMiddleware ──────────────────────────────────────────
// When the LLM emits text content in the same response as a render tool call
// (render_a2ui or render_spending_summary), the text appears as a duplicate
// below the card. This middleware buffers TEXT_MESSAGE_* events and discards
// them if any render tool call is detected in the same turn.

const RENDER_TOOLS = new Set(["render_a2ui", "render_spending_summary"]);

// eslint-disable-next-line @typescript-eslint/no-explicit-any
class SuppressRenderToolTextMiddleware extends (Middleware as any) {
  // eslint-disable-next-line @typescript-eslint/no-explicit-any
  run(input: any, next: any): Observable<any> {
    let renderToolSeen = false;
    // eslint-disable-next-line @typescript-eslint/no-explicit-any
    const textBuffer: any[] = [];

    // eslint-disable-next-line @typescript-eslint/no-explicit-any
    return new Observable<any>((observer) => {
      // eslint-disable-next-line @typescript-eslint/no-explicit-any
      const sub = (this as any).runNextWithState(input, next).subscribe({
        // eslint-disable-next-line @typescript-eslint/no-explicit-any
        next: (eventWithState: any) => {
          const event = eventWithState.event; // raw BaseEvent
          const type: string = event?.type ?? "";

          if (type === EventType.TOOL_CALL_START) {
            const toolName: string = event?.toolCallName ?? "";
            if (RENDER_TOOLS.has(toolName)) {
              renderToolSeen = true;
              textBuffer.length = 0; // discard text buffered before we saw the tool call
            }
          }

          if (
            type === EventType.TEXT_MESSAGE_START ||
            type === EventType.TEXT_MESSAGE_CONTENT ||
            type === EventType.TEXT_MESSAGE_END
          ) {
            if (!renderToolSeen) textBuffer.push(event); // buffer raw event
            return; // always hold — flush at turn end
          }

          if (type === EventType.RUN_FINISHED || type === EventType.RUN_ERROR) {
            if (!renderToolSeen) {
              for (const e of textBuffer) observer.next(e); // flush raw events
            }
            textBuffer.length = 0;
            observer.next(event); // emit raw event
            return;
          }

          observer.next(event); // emit raw event
        },
        // eslint-disable-next-line @typescript-eslint/no-explicit-any
        error: (err: any) => observer.error(err),
        complete: () => {
          if (!renderToolSeen) {
            for (const e of textBuffer) observer.next(e);
          }
          observer.complete();
        },
      });
      return () => sub.unsubscribe();
    });
  }
}

// ── Hono app ─────────────────────────────────────────────────────────────────

const app = new Hono();

// Security headers on every response. Responses returned by fetch() (the
// backend proxy) carry immutable headers, so rebuild the Response instead of
// mutating in place.
const SECURITY_HEADERS: Record<string, string> = {
  "X-Frame-Options": "DENY",
  "X-Content-Type-Options": "nosniff",
  "Referrer-Policy": "strict-origin-when-cross-origin",
};

app.use("*", async (c, next) => {
  await next();
  const res = c.res;
  const headers = new Headers(res.headers);
  for (const [k, v] of Object.entries(SECURITY_HEADERS)) headers.set(k, v);
  c.res = new Response(res.body, {
    status: res.status,
    statusText: res.statusText,
    headers,
  });
});

app.all("/api/copilotkit/*", async (c) => {
  const cookieHeader = c.req.header("cookie") ?? "";
  const match = cookieHeader.match(/(?:^|;\s*)ws_token=([^;]+)/);
  const token = match?.[1];
  const agentHeaders: Record<string, string> = token
    ? { Authorization: `Bearer ${token}` }
    : {};

  const agent = new HttpAgent({ url: AGENT_URL, headers: agentHeaders });
  // CK_MIDDLEWARES=off runs the stack bare — used to audit which of these
  // workarounds are still needed after CopilotKit/MAF upgrades.
  /* eslint-disable @typescript-eslint/no-explicit-any */
  if (process.env.CK_MIDDLEWARES !== "off") {
    agent.use(new SuppressRenderToolTextMiddleware() as any);
    agent.use(new StripModelArtifactsMiddleware() as any);
    // ReconcileSnapshotMiddleware retired 2026-07-04: MAF 1.9+ reuses the
    // streamed message id in MESSAGES_SNAPSHOT (verified in _agent_run.py at
    // python-1.10.0 and in live event logs), and keeping the surgery caused
    // duplicated user messages on a2ui follow-up runs.
    agent.use(new DeduplicateToolCallMiddleware() as any);
  }
  if (process.env.CK_DEBUG_EVENTS === "1" || process.env.CK_MIDDLEWARES === "off") {
    agent.use(new DebugLogMiddleware() as any);
  }
  /* eslint-enable @typescript-eslint/no-explicit-any */

  const runtime = new CopilotRuntime({
    agents: { wealthysmart: agent },
    a2ui: { injectA2UITool: true },
    beforeRequestMiddleware: async ({ request: outbound }) => {
      if (outbound.method !== "POST") return;
      const ct = outbound.headers.get("content-type") ?? "";
      if (!ct.includes("application/json")) return;
      try {
        const body = (await outbound.clone().json()) as {
          messages?: AGUIMessage[];
          context?: { description?: string; value?: string }[];
        };
        if (!Array.isArray(body.messages)) return;
        let messages = pairOrphanToolCalls(body.messages);

        // MAF's AG-UI endpoint declares `context` on its request model but
        // never feeds it to the model (agent_framework_ag_ui, checked at
        // python-1.10.0). CopilotKit 1.60+ ships the A2UI catalog schema and
        // guidelines in that context — without them the model emits ops the
        // A2UI validator rejects and the surface never paints. Inject the
        // context as a system message ourselves. Remove once MAF consumes
        // RunAgentInput.context.
        if (Array.isArray(body.context) && body.context.length > 0) {
          const contextText = body.context
            .map((c) => `## ${c.description ?? "Context"}\n${c.value ?? ""}`)
            .join("\n\n");
          messages = [
            {
              id: `ctx-${Date.now()}`,
              role: "system",
              content: `Additional run context:\n\n${contextText}`,
            } as AGUIMessage,
            ...messages,
          ];
        }

        if (messages === body.messages) return;
        return new Request(outbound.url, {
          method: outbound.method,
          headers: outbound.headers,
          body: JSON.stringify({ ...body, messages }),
        });
      } catch (err) {
        // Skipping the repair is survivable; doing it silently is not —
        // orphan tool calls then fail downstream with no trace of why.
        console.error("[copilotkit] outbound request repair skipped:", err);
        return;
      }
    },
  });

  const { handleRequest } = copilotRuntimeNextJSAppRouterEndpoint({
    runtime,
    serviceAdapter: new ExperimentalEmptyAdapter(),
    endpoint: "/api/copilotkit",
  });

  return handleRequest(c.req.raw as Parameters<typeof handleRequest>[0]);
});

app.get("/api/health", (c) => c.json({ ok: true }));

// Proxy backend API calls (FastAPI). In dev these hit Vite's proxy directly,
// but in prod the browser talks to this Hono server, which must forward
// `/api/v1/*` and `/api/auth/*` to the FastAPI container — otherwise the SPA
// fallback below swallows them and returns index.html.
const proxyToBackend = async (c: import("hono").Context) => {
  const url = new URL(c.req.url);
  const target = `${BACKEND_URL}${url.pathname}${url.search}`;
  const method = c.req.method;
  const headers = new Headers(c.req.raw.headers);
  headers.delete("host");
  const init: RequestInit = {
    method,
    headers,
    redirect: "manual",
  };
  if (method !== "GET" && method !== "HEAD") {
    init.body = c.req.raw.body;
    // @ts-expect-error undici requires duplex for streamed bodies
    init.duplex = "half";
  }
  try {
    return await fetch(target, init);
  } catch (err) {
    console.error(`[proxy] ${method} ${url.pathname} → backend failed:`, err);
    return c.json({ error: "Backend unavailable" }, 502);
  }
};

app.all("/api/v1/*", proxyToBackend);
app.all("/api/auth/*", proxyToBackend);

// In production, serve the Vite build output.
if (isProd) {
  app.use("/*", serveStatic({ root: "./dist" }));
  // SPA fallback: any non-asset path returns index.html
  app.get("/*", serveStatic({ path: "./dist/index.html" }));
}

serve({ fetch: app.fetch, port: PORT }, (info) => {
  console.log(`CopilotKit server on port ${info.port} [${isProd ? "prod" : "dev"}]`);
});