import { z } from 'zod';
import {
  Tools,
  SSEOptionsSchema,
  MCPOptionsSchema,
  MCPServersSchema,
  StdioOptionsSchema,
  WebSocketOptionsSchema,
  StreamableHTTPOptionsSchema,
} from 'librechat-data-provider';
import type {
  EmbeddedResource,
  ListToolsResult,
  ImageContent,
  AudioContent,
  TextContent,
  Tool,
} from '@modelcontextprotocol/sdk/types.js';
import type { SearchResultData, UIResource, TPlugin } from 'librechat-data-provider';
import type { TokenMethods, IUser } from '@librechat/data-schemas';
import type { LCTool } from '@librechat/agents';
import type { FlowStateManager } from '~/flow/manager';
import type { RequestBody } from '~/types/http';
import type * as o from '~/mcp/oauth/types';

export type StdioOptions = z.infer<typeof StdioOptionsSchema>;
export type WebSocketOptions = z.infer<typeof WebSocketOptionsSchema>;
export type SSEOptions = z.infer<typeof SSEOptionsSchema>;
export type StreamableHTTPOptions = z.infer<typeof StreamableHTTPOptionsSchema>;
export type MCPOptions = z.infer<typeof MCPOptionsSchema> & {
  customUserVars?: Record<
    string,
    {
      title: string;
      description: string;
    }
  >;
};
export type MCPServers = z.infer<typeof MCPServersSchema>;
export interface MCPResource {
  uri: string;
  name: string;
  description?: string;
  mimeType?: string;
}

export interface LCFunctionTool {
  type: 'function';
  ['function']: LCTool;
}

export type LCAvailableTools = Record<string, LCFunctionTool>;
export type LCManifestTool = TPlugin;
export type LCToolManifest = TPlugin[];
export interface MCPPrompt {
  name: string;
  description?: string;
  arguments?: Array<{ name: string }>;
}

export type ConnectionState = 'disconnected' | 'connecting' | 'connected' | 'error';

export type MCPTool = Tool;
export type MCPToolListResponse = ListToolsResult;
export type ToolContentPart = TextContent | ImageContent | EmbeddedResource | AudioContent;
export type { TextContent, ImageContent, EmbeddedResource, AudioContent };
export type MCPToolCallResponse =
  | undefined
  | {
      _meta?: Record<string, unknown>;
      content?: Array<ToolContentPart>;
      isError?: boolean;
    };

export type Provider =
  | 'google'
  | 'anthropic'
  | 'openai'
  | 'azureopenai'
  | 'openrouter'
  | 'xai'
  | 'deepseek'
  | 'ollama'
  | 'bedrock';

export type FormattedContent =
  | {
      type: 'text';
      text: string;
    }
  | {
      type: 'image';
      inlineData: {
        mimeType: string;
        data: string;
      };
    }
  | {
      type: 'image';
      source: {
        type: 'base64';
        media_type: string;
        data: string;
      };
    }
  | {
      type: 'image_url';
      image_url: {
        url: string;
      };
    };

export type FileSearchSource = {
  fileId: string;
  relevance: number;
  fileName?: string;
  metadata?: {
    storageType?: string;
    [key: string]: unknown;
  };
  [key: string]: unknown;
};

export type Artifacts =
  | {
      content?: FormattedContent[];
      [Tools.ui_resources]?: {
        data: UIResource[];
      };
      [Tools.file_search]?: {
        sources: FileSearchSource[];
        fileCitations?: boolean;
      };
      [Tools.web_search]?: SearchResultData;
      files?: Array<{ id: string; name: string }>;
      session_id?: string;
      file_ids?: string[];
    }
  | undefined;

export type FormattedContentResult = [string, Artifacts | undefined];

export type ImageFormatter = (item: ImageContent) => FormattedContent;

export type FormattedToolResponse = FormattedContentResult;

/**
 * Origin of an MCP server definition.
 * - `'yaml'`   — operator-defined in librechat.yaml, full trust, boot-time init
 * - `'config'` — admin-defined via Config override, full trust, lazy init
 * - `'user'`   — user-provided via UI, sandboxed (restricted placeholder resolution)
 */
export type MCPServerSource = 'yaml' | 'config' | 'user';

export type ParsedServerConfig = MCPOptions & {
  url?: string;
  requiresOAuth?: boolean;
  oauthMetadata?: Record<string, unknown> | null;
  capabilities?: string;
  tools?: string;
  toolFunctions?: LCAvailableTools;
  initDuration?: number;
  updatedAt?: number;
  dbId?: string;
  /** Origin of this server definition — determines trust level and placeholder resolution */
  source?: MCPServerSource;
  /** True if access is only via agent (not directly shared with user) */
  consumeOnly?: boolean;
  /** True when inspection failed at startup; the server is known but not fully initialized */
  inspectionFailed?: boolean;
};

export type AddServerResult = {
  serverName: string;
  config: ParsedServerConfig;
};

export interface BasicConnectionOptions {
  serverName: string;
  serverConfig: MCPOptions;
  useSSRFProtection?: boolean;
  allowedDomains?: string[] | null;
  /** When true, only resolve customUserVars in processMCPEnv (for DB-stored servers) */
  dbSourced?: boolean;
}

/** User context for placeholder resolution in MCP connections (non-OAuth and OAuth alike) */
export interface UserConnectionContext {
  user?: IUser;
  customUserVars?: Record<string, string>;
  requestBody?: RequestBody;
  connectionTimeout?: number;
}

export interface OAuthConnectionOptions extends UserConnectionContext {
  useOAuth: true;
  flowManager: FlowStateManager<o.MCPOAuthTokens | null>;
  tokenMethods?: TokenMethods;
  signal?: AbortSignal;
  oauthStart?: (authURL: string) => Promise<void>;
  oauthEnd?: () => Promise<void>;
  returnOnOAuth?: boolean;
}

export interface ToolDiscoveryOptions {
  serverName: string;
  user?: IUser;
  flowManager?: FlowStateManager<o.MCPOAuthTokens | null>;
  tokenMethods?: TokenMethods;
  signal?: AbortSignal;
  oauthStart?: (authURL: string) => Promise<void>;
  customUserVars?: Record<string, string>;
  requestBody?: RequestBody;
  connectionTimeout?: number;
  /** Pre-resolved config-source servers for tenant-scoped lookup */
  configServers?: Record<string, ParsedServerConfig>;
}

export interface ToolDiscoveryResult {
  tools: Tool[] | null;
  oauthRequired: boolean;
  oauthUrl: string | null;
}
