import { WebSocket } from 'ws';
import type {
  NodeLocationSetParams,
  NodePushTokenSetParams,
  NodeStatusSetParams,
} from '../gateway/protocol.js';
import { GATEWAY_PROTOCOL_VERSION } from '../gateway/protocol.js';

interface RpcSuccess {
  id: number;
  result: unknown;
}

interface RpcFailure {
  id: number;
  error: {
    code: number;
    message: string;
  };
}

interface RpcEvent {
  id: number;
  event: string;
  data: unknown;
}

type RpcMessage = RpcSuccess | RpcFailure | RpcEvent;

type PendingRequest = {
  resolve: (value: unknown) => void;
  reject: (error: Error) => void;
  timeout: NodeJS.Timeout;
};

export interface CompanionRuntimeClientOptions {
  url: string;
  token?: string;
  requestTimeoutMs?: number;
  autoConnect?: boolean;
  websocketFactory?: (url: string) => WebSocket;
}

export interface WaitForIdleOptions {
  timeoutMs?: number;
  pollIntervalMs?: number;
  signal?: AbortSignal;
}

export type CompanionEventHandler = (event: string, data: unknown) => void;
export type CompanionTypedEventHandler<TData = unknown> = (data: TData) => void;
export type CompanionEventPredicate<TData = unknown> = (data: TData) => boolean;
export type CompanionEventEnvelope<TData = unknown> = {
  event: string;
  data: TData;
};

export const COMPANION_EVENT_NAMES = {
  agentStream: 'agent.stream',
  agentTyping: 'agent.typing',
  contextWarning: 'context_warning',
} as const;
export type CompanionEventName =
  (typeof COMPANION_EVENT_NAMES)[keyof typeof COMPANION_EVENT_NAMES];

export interface RegisterNodeInput {
  nodeId: string;
  role: string;
  capabilities: string[];
  protocolVersion?: number;
}

export interface NodeIdentitySummary {
  id: string;
  role: string;
}

export interface NodeRegisterResult {
  registered: boolean;
  node: NodeIdentitySummary;
  protocol: {
    serverVersion: number;
    clientVersion: number;
    negotiatedVersion: number;
  };
  capabilities: {
    declared: string[];
    enabled: string[];
  };
}

export interface NodeCapabilitiesResult {
  protocol: {
    serverVersion: number;
    nodeVersion: number;
    negotiatedVersion: number;
  };
  node: {
    id: string;
    role: string;
    registeredAt: number;
  };
  capabilities: {
    declared: string[];
    enabled: string[];
    featureGates: Record<string, boolean>;
  };
}

export interface NodeBootstrapResult {
  register: NodeRegisterResult;
  capabilities: NodeCapabilitiesResult;
  systemCapabilities?: SystemCapabilitiesResult;
}

export interface NodeStatus {
  platform: 'macos' | 'ios' | 'android' | 'linux' | 'windows' | 'unknown';
  appVersion?: string;
  deviceName?: string;
  statusText?: string;
  batteryPct?: number;
  powerSource: 'ac' | 'battery' | 'unknown';
  reportedAt: number;
}

export interface NodeLocation {
  latitude: number;
  longitude: number;
  accuracyMeters?: number;
  altitudeMeters?: number;
  headingDegrees?: number;
  speedMps?: number;
  source: 'gps' | 'network' | 'manual' | 'unknown';
  capturedAt: number;
  receivedAt: number;
}

export interface NodePushSummary {
  provider: 'apns' | 'fcm';
  tokenPreview: string;
  topic?: string;
  environment?: 'sandbox' | 'production';
  registeredAt: number;
}

export interface NodeStatusSetResult {
  updated: boolean;
  node: NodeIdentitySummary;
  status: NodeStatus;
}

export interface NodeLocationSetResult {
  updated: boolean;
  node: NodeIdentitySummary;
  location: NodeLocation;
}

export interface NodePushTokenSetResult {
  updated: boolean;
  node: NodeIdentitySummary;
  push: NodePushSummary;
}

export interface NodeLocationGetResult {
  node: NodeIdentitySummary;
  location: NodeLocation | null;
}

export interface SystemCapabilitiesResult {
  protocol: {
    version: number;
  };
  nodes: {
    enabled: boolean;
    locationEnabled: boolean;
    pushEnabled: boolean;
    allowedRoles: string[];
    registered: boolean;
    role?: string;
    nodeId?: string;
  };
  featureGates: Record<string, boolean>;
}

export interface SystemNodeEntry {
  connectionId: string;
  nodeId: string;
  role: string;
  identity?: string;
  protocolVersion: number;
  capabilities: string[];
  registeredAt: number;
  location?: NodeLocation;
  status?: NodeStatus;
  push?: NodePushSummary;
}

export interface SystemNodesResult {
  nodes: SystemNodeEntry[];
  summary: {
    total: number;
  };
}

export interface ListNodesInput {
  role?: string;
  platform?: string;
  limit?: number;
}

export interface SetNodeStatusInput extends Omit<NodeStatusSetParams, 'connectionId'> {}

export interface SetNodeLocationInput extends Omit<NodeLocationSetParams, 'connectionId'> {}

export interface SetNodePushTokenInput extends Omit<NodePushTokenSetParams, 'connectionId'> {}

export interface CanvasArtifact {
  id: string;
  type: string;
  title?: string;
  content: unknown;
  metadata?: Record<string, unknown>;
  createdAt: number;
  updatedAt: number;
}

export interface PutCanvasArtifactInput {
  sessionId: string;
  artifactId?: string;
  type: string;
  title?: string;
  content: unknown;
  metadata?: Record<string, unknown>;
}

export interface GetCanvasArtifactInput {
  sessionId: string;
  artifactId: string;
}

export interface DeleteCanvasArtifactInput {
  sessionId: string;
  artifactId: string;
}

export interface CanvasPutResult {
  artifact: CanvasArtifact;
  upserted: boolean;
}

export interface CanvasGetResult {
  artifact: CanvasArtifact;
}

export interface CanvasListResult {
  artifacts: CanvasArtifact[];
}

export interface CanvasDeleteResult {
  deleted: boolean;
}

export interface CanvasClearResult {
  cleared: number;
}

export class GatewayRpcError extends Error {
  readonly code: number;

  constructor(code: number, message: string) {
    super(message);
    this.name = 'GatewayRpcError';
    this.code = code;
  }
}

export class CompanionRuntimeClient {
  private readonly url: string;
  private readonly token?: string;
  private readonly requestTimeoutMs: number;
  private readonly autoConnect: boolean;
  private readonly websocketFactory: (url: string) => WebSocket;

  private ws: WebSocket | null = null;
  private connectPromise: Promise<void> | null = null;
  private nextId = 1;
  private pending = new Map<number, PendingRequest>();
  private readonly eventHandlers = new Set<CompanionEventHandler>();
  private readonly pendingEventWaits = new Set<(error: Error) => void>();

  constructor(options: CompanionRuntimeClientOptions) {
    const requestTimeoutMs = options.requestTimeoutMs ?? 15_000;
    if (!Number.isFinite(requestTimeoutMs) || requestTimeoutMs <= 0) {
      throw new Error('requestTimeoutMs must be a positive number');
    }
    this.url = options.url;
    this.token = options.token;
    this.requestTimeoutMs = requestTimeoutMs;
    this.autoConnect = options.autoConnect ?? false;
    this.websocketFactory = options.websocketFactory ?? ((url) => new WebSocket(url));
  }

  get connected(): boolean {
    return this.ws?.readyState === WebSocket.OPEN;
  }

  get eventSubscriptionCount(): number {
    return this.eventHandlers.size;
  }

  get pendingRequestCount(): number {
    return this.pending.size;
  }

  get pendingEventWaitCount(): number {
    return this.pendingEventWaits.size;
  }

  get hasPendingWork(): boolean {
    return this.pendingRequestCount > 0 || this.pendingEventWaitCount > 0;
  }

  async connect(): Promise<void> {
    if (this.connected) {
      return;
    }

    if (this.connectPromise) {
      return this.connectPromise;
    }

    this.connectPromise = this.openConnection();
    try {
      await this.connectPromise;
    } finally {
      this.connectPromise = null;
    }
  }

  private async openConnection(): Promise<void> {
    const ws = this.websocketFactory(withToken(this.url, this.token));

    await new Promise<void>((resolve, reject) => {
      let settled = false;

      const onOpen = () => {
        cleanup();
        settled = true;
        this.ws = ws;
        this.ws.on('message', (raw) => this.handleMessage(raw.toString()));
        this.ws.on('close', () => {
          if (this.ws === ws) {
            this.ws = null;
          }
          this.rejectAllPending(new Error('WebSocket closed'));
          this.rejectEventWaits(new Error('WebSocket closed'));
        });
        this.ws.on('error', () => {
          // close event handles pending rejection
        });
        resolve();
      };

      const onError = (err: Error) => {
        cleanup();
        settled = true;
        reject(err);
      };

      const onClose = () => {
        cleanup();
        if (!settled) {
          settled = true;
          reject(new Error('WebSocket closed before connection established'));
        }
      };

      const cleanup = () => {
        ws.off('open', onOpen);
        ws.off('error', onError);
        ws.off('close', onClose);
      };

      ws.once('open', onOpen);
      ws.once('error', onError);
      ws.once('close', onClose);
    });
  }

  disconnect(code?: number, reason?: string): void {
    if (!this.ws) {
      this.rejectEventWaits(new Error('Disconnected'));
      return;
    }

    const ws = this.ws;
    this.ws = null;
    this.rejectAllPending(new Error('Disconnected'));
    this.rejectEventWaits(new Error('Disconnected'));
    ws.close(code, reason);
  }

  dispose(code?: number, reason?: string): void {
    this.disconnect(code, reason);
    this.clearEventSubscriptions();
  }

  subscribeEvents(handler: CompanionEventHandler): () => void {
    this.eventHandlers.add(handler);
    return () => {
      this.eventHandlers.delete(handler);
    };
  }

  clearEventSubscriptions(): void {
    this.eventHandlers.clear();
    this.rejectEventWaits(new Error('Event subscriptions cleared'));
  }

  subscribeEvent<TData = unknown>(
    eventName: CompanionEventName | string,
    handler: CompanionTypedEventHandler<TData>,
  ): () => void {
    return this.subscribeEvents((event, data) => {
      if (event !== eventName) {
        return;
      }
      handler(data as TData);
    });
  }

  subscribeAgentStream<TData = unknown>(
    handler: CompanionTypedEventHandler<TData>,
  ): () => void {
    return this.subscribeEvent<TData>(COMPANION_EVENT_NAMES.agentStream, handler);
  }

  subscribeAgentTyping<TData = unknown>(
    handler: CompanionTypedEventHandler<TData>,
  ): () => void {
    return this.subscribeEvent<TData>(COMPANION_EVENT_NAMES.agentTyping, handler);
  }

  subscribeContextWarning<TData = unknown>(
    handler: CompanionTypedEventHandler<TData>,
  ): () => void {
    return this.subscribeEvent<TData>(COMPANION_EVENT_NAMES.contextWarning, handler);
  }

  waitForEvent<TData = unknown>(
    eventName: CompanionEventName | string,
    options?: {
      timeoutMs?: number;
      predicate?: CompanionEventPredicate<TData>;
      signal?: AbortSignal;
    },
  ): Promise<TData> {
    if (typeof eventName !== 'string' || eventName.trim().length === 0) {
      throw new Error('eventName must be a non-empty string');
    }
    const timeoutMs = options?.timeoutMs ?? this.requestTimeoutMs;
    const predicate = options?.predicate;
    const signal = options?.signal;

    return new Promise<TData>((resolve, reject) => {
      let settled = false;
      let abortCleanup: (() => void) | null = null;

      const finish = (fn: () => void) => {
        if (settled) {
          return;
        }
        settled = true;
        clearTimeout(timeout);
        unsubscribe();
        if (abortCleanup) {
          abortCleanup();
          abortCleanup = null;
        }
        this.pendingEventWaits.delete(cancelWait);
        fn();
      };

      const cancelWait = (error: Error) => {
        finish(() => reject(error));
      };
      this.pendingEventWaits.add(cancelWait);

      const unsubscribe = this.subscribeEvent<TData>(eventName, (data) => {
        if (predicate && !predicate(data)) {
          return;
        }
        finish(() => resolve(data));
      });

      const timeout = setTimeout(() => {
        finish(() => reject(new Error(`Timed out waiting for event ${eventName}`)));
      }, timeoutMs);

      if (signal) {
        const onAbort = () => {
          cancelWait(new Error(`Aborted while waiting for event ${eventName}`));
        };
        signal.addEventListener('abort', onAbort, { once: true });
        abortCleanup = () => {
          signal.removeEventListener('abort', onAbort);
        };
        if (signal.aborted) {
          onAbort();
        }
      }
    });
  }

  waitForAnyEvent<TData = unknown>(
    eventNames: readonly (CompanionEventName | string)[],
    options?: {
      timeoutMs?: number;
      predicate?: (event: string, data: TData) => boolean;
      signal?: AbortSignal;
    },
  ): Promise<CompanionEventEnvelope<TData>> {
    if (eventNames.length === 0) {
      throw new Error('eventNames must contain at least one event name');
    }
    if (
      eventNames.some(
        (eventName) => typeof eventName !== 'string' || eventName.trim().length === 0,
      )
    ) {
      throw new Error('eventNames must not contain empty values');
    }
    const eventNameSet = new Set(eventNames);
    const timeoutMs = options?.timeoutMs ?? this.requestTimeoutMs;
    const predicate = options?.predicate;
    const signal = options?.signal;

    return new Promise<CompanionEventEnvelope<TData>>((resolve, reject) => {
      let settled = false;
      let abortCleanup: (() => void) | null = null;

      const finish = (fn: () => void) => {
        if (settled) {
          return;
        }
        settled = true;
        clearTimeout(timeout);
        unsubscribe();
        if (abortCleanup) {
          abortCleanup();
          abortCleanup = null;
        }
        this.pendingEventWaits.delete(cancelWait);
        fn();
      };

      const cancelWait = (error: Error) => {
        finish(() => reject(error));
      };
      this.pendingEventWaits.add(cancelWait);

      const unsubscribe = this.subscribeEvents((event, data) => {
        if (!eventNameSet.has(event)) {
          return;
        }
        const castData = data as TData;
        if (predicate && !predicate(event, castData)) {
          return;
        }
        finish(() => resolve({ event, data: castData }));
      });

      const timeout = setTimeout(() => {
        cancelWait(new Error(`Timed out waiting for any event in [${eventNames.join(', ')}]`));
      }, timeoutMs);

      if (signal) {
        const onAbort = () => {
          cancelWait(new Error(`Aborted while waiting for events [${eventNames.join(', ')}]`));
        };
        signal.addEventListener('abort', onAbort, { once: true });
        abortCleanup = () => {
          signal.removeEventListener('abort', onAbort);
        };
        if (signal.aborted) {
          onAbort();
        }
      }
    });
  }

  waitForAgentStream<TData = unknown>(options?: {
    timeoutMs?: number;
    predicate?: CompanionEventPredicate<TData>;
    signal?: AbortSignal;
  }): Promise<TData> {
    return this.waitForEvent<TData>(COMPANION_EVENT_NAMES.agentStream, options);
  }

  waitForAgentTyping<TData = unknown>(options?: {
    timeoutMs?: number;
    predicate?: CompanionEventPredicate<TData>;
    signal?: AbortSignal;
  }): Promise<TData> {
    return this.waitForEvent<TData>(COMPANION_EVENT_NAMES.agentTyping, options);
  }

  waitForContextWarning<TData = unknown>(options?: {
    timeoutMs?: number;
    predicate?: CompanionEventPredicate<TData>;
    signal?: AbortSignal;
  }): Promise<TData> {
    return this.waitForEvent<TData>(COMPANION_EVENT_NAMES.contextWarning, options);
  }

  waitForIdle(options?: WaitForIdleOptions): Promise<void> {
    const pollIntervalMs = options?.pollIntervalMs ?? 25;
    if (!Number.isFinite(pollIntervalMs) || pollIntervalMs <= 0) {
      throw new Error('pollIntervalMs must be a positive number');
    }
    if (!this.hasPendingWork) {
      return Promise.resolve();
    }
    const timeoutMs = options?.timeoutMs ?? this.requestTimeoutMs;
    const signal = options?.signal;

    return new Promise<void>((resolve, reject) => {
      let settled = false;
      let timeout: NodeJS.Timeout | null = null;
      let poll: NodeJS.Timeout | null = null;
      let abortCleanup: (() => void) | null = null;

      const cleanup = () => {
        if (timeout) {
          clearTimeout(timeout);
          timeout = null;
        }
        if (poll) {
          clearInterval(poll);
          poll = null;
        }
        if (abortCleanup) {
          abortCleanup();
          abortCleanup = null;
        }
      };

      const finish = (fn: () => void) => {
        if (settled) {
          return;
        }
        settled = true;
        cleanup();
        fn();
      };

      const check = () => {
        if (!this.hasPendingWork) {
          finish(() => resolve());
        }
      };

      timeout = setTimeout(() => {
        finish(() => reject(new Error('Timed out waiting for runtime idle state')));
      }, timeoutMs);
      poll = setInterval(check, pollIntervalMs);
      check();

      if (signal) {
        const onAbort = () => {
          finish(() => reject(new Error('Aborted while waiting for runtime idle state')));
        };
        signal.addEventListener('abort', onAbort, { once: true });
        abortCleanup = () => {
          signal.removeEventListener('abort', onAbort);
        };
        if (signal.aborted) {
          onAbort();
        }
      }
    });
  }

  listKnownEventNames(): CompanionEventName[] {
    return Object.values(COMPANION_EVENT_NAMES);
  }

  async call<T>(method: string, params?: Record<string, unknown>): Promise<T> {
    if (!this.connected) {
      if (!this.autoConnect) {
        throw new Error('WebSocket is not connected');
      }
      await this.connect();
    }

    if (!this.ws || this.ws.readyState !== WebSocket.OPEN) {
      throw new Error('WebSocket is not connected');
    }

    const id = this.nextId++;

    return new Promise<T>((resolve, reject) => {
      const timeout = setTimeout(() => {
        this.pending.delete(id);
        reject(new Error(`RPC timeout for method ${method}`));
      }, this.requestTimeoutMs);

      this.pending.set(id, {
        resolve: (value) => {
          clearTimeout(timeout);
          resolve(value as T);
        },
        reject: (error) => {
          clearTimeout(timeout);
          reject(error);
        },
        timeout,
      });

      this.ws?.send(JSON.stringify({ id, method, params: params ?? {} }), (err) => {
        if (!err) {
          return;
        }

        const pending = this.pending.get(id);
        if (!pending) {
          return;
        }
        clearTimeout(pending.timeout);
        this.pending.delete(id);
        reject(err);
      });
    });
  }

  registerNode(input: RegisterNodeInput): Promise<NodeRegisterResult> {
    return this.call<NodeRegisterResult>('node.register', {
      nodeId: input.nodeId,
      role: input.role,
      protocolVersion: input.protocolVersion ?? GATEWAY_PROTOCOL_VERSION,
      capabilities: input.capabilities,
    });
  }

  async bootstrapNode(
    input: RegisterNodeInput,
    options?: { includeSystemCapabilities?: boolean },
  ): Promise<NodeBootstrapResult> {
    const register = await this.registerNode(input);
    const capabilities = await this.getNodeCapabilities();
    if (options?.includeSystemCapabilities) {
      const systemCapabilities = await this.getSystemCapabilities();
      return {
        register,
        capabilities,
        systemCapabilities,
      };
    }
    return {
      register,
      capabilities,
    };
  }

  getNodeCapabilities(): Promise<NodeCapabilitiesResult> {
    return this.call<NodeCapabilitiesResult>('node.capabilities.get');
  }

  setNodeStatus(input: SetNodeStatusInput): Promise<NodeStatusSetResult> {
    return this.call<NodeStatusSetResult>('node.status.set', {
      platform: input.platform,
      appVersion: input.appVersion,
      deviceName: input.deviceName,
      statusText: input.statusText,
      batteryPct: input.batteryPct,
      powerSource: input.powerSource,
    });
  }

  setNodeLocation(input: SetNodeLocationInput): Promise<NodeLocationSetResult> {
    return this.call<NodeLocationSetResult>('node.location.set', {
      latitude: input.latitude,
      longitude: input.longitude,
      accuracyMeters: input.accuracyMeters,
      altitudeMeters: input.altitudeMeters,
      headingDegrees: input.headingDegrees,
      speedMps: input.speedMps,
      source: input.source,
      capturedAt: input.capturedAt,
    });
  }

  getNodeLocation(): Promise<NodeLocationGetResult> {
    return this.call<NodeLocationGetResult>('node.location.get');
  }

  setNodePushToken(input: SetNodePushTokenInput): Promise<NodePushTokenSetResult> {
    return this.call<NodePushTokenSetResult>('node.push_token.set', {
      provider: input.provider,
      token: input.token,
      topic: input.topic,
      environment: input.environment,
    });
  }

  getSystemCapabilities(): Promise<SystemCapabilitiesResult> {
    return this.call<SystemCapabilitiesResult>('system.capabilities');
  }

  listSystemNodes(input?: ListNodesInput): Promise<SystemNodesResult> {
    return this.call<SystemNodesResult>('system.nodes', {
      role: input?.role,
      platform: input?.platform,
      limit: input?.limit,
    });
  }

  putCanvasArtifact(input: PutCanvasArtifactInput): Promise<CanvasPutResult> {
    return this.call<CanvasPutResult>('canvas.put', {
      sessionId: input.sessionId,
      artifactId: input.artifactId,
      type: input.type,
      title: input.title,
      content: input.content,
      metadata: input.metadata,
    });
  }

  getCanvasArtifact(input: GetCanvasArtifactInput): Promise<CanvasGetResult> {
    return this.call<CanvasGetResult>('canvas.get', {
      sessionId: input.sessionId,
      artifactId: input.artifactId,
    });
  }

  listCanvasArtifacts(sessionId: string): Promise<CanvasListResult> {
    return this.call<CanvasListResult>('canvas.list', { sessionId });
  }

  deleteCanvasArtifact(input: DeleteCanvasArtifactInput): Promise<CanvasDeleteResult> {
    return this.call<CanvasDeleteResult>('canvas.delete', {
      sessionId: input.sessionId,
      artifactId: input.artifactId,
    });
  }

  clearCanvasArtifacts(sessionId: string): Promise<CanvasClearResult> {
    return this.call<CanvasClearResult>('canvas.clear', { sessionId });
  }

  private handleMessage(raw: string): void {
    let parsed: RpcMessage;
    try {
      parsed = JSON.parse(raw) as RpcMessage;
    } catch {
      return;
    }

    if (!('id' in parsed) || typeof parsed.id !== 'number') {
      return;
    }

    if ('event' in parsed) {
      for (const handler of this.eventHandlers) {
        try {
          handler(parsed.event, parsed.data);
        } catch {
          // Event subscribers are userland callbacks; isolate failures.
        }
      }
      return;
    }

    const pending = this.pending.get(parsed.id);
    if (!pending) {
      return;
    }

    this.pending.delete(parsed.id);

    if ('error' in parsed) {
      pending.reject(new GatewayRpcError(parsed.error.code, parsed.error.message));
      return;
    }

    pending.resolve(parsed.result);
  }

  private rejectAllPending(error: Error): void {
    for (const [, pending] of this.pending) {
      clearTimeout(pending.timeout);
      pending.reject(error);
    }
    this.pending.clear();
  }

  private rejectEventWaits(error: Error): void {
    for (const cancel of this.pendingEventWaits) {
      cancel(error);
    }
    this.pendingEventWaits.clear();
  }
}

function withToken(url: string, token?: string): string {
  if (!token) {
    return url;
  }

  const parsed = new URL(url);
  parsed.searchParams.set('token', token);
  return parsed.toString();
}
