import { DataAPIClient } from '@datastax/astra-db-ts';

export interface ToolParameter {
  param: string;
  description?: string;
  attribute?: string;
  type?: string;
  required?: boolean | number;
  operator?: string;
  expr?: string;
  enum?: string[];
  embedding_model?: string;
  info?: string;
  value?: any;
  paramMode?: 'tool_param' | 'static' | 'expression';
}

export interface Tool {
  _id?: string;
  tags?: string[];
  type: string;
  name: string;
  description?: string;
  limit?: number;
  method?: string;
  collection_name?: string;
  table_name?: string;
  db_name?: string;
  projection?: Record<string, number | string>;
  parameters?: ToolParameter[];
  enabled?: boolean;
}

class AstraClient {
  private client: DataAPIClient | null = null;
  private db: any = null;
  private collection: any = null;
  private token: string | null = null;
  private endpoint: string | null = null;
  private dbName: string | null = null;
  private collectionName: string = 'tool_catalog';

  async connect(): Promise<void> {
    // Access server-side environment variables
    this.token = process.env.ASTRA_DB_APPLICATION_TOKEN || '';
    this.endpoint = process.env.ASTRA_DB_API_ENDPOINT || '';
    this.dbName = process.env.ASTRA_DB_DB_NAME || '';
    this.collectionName = process.env.ASTRA_DB_CATALOG_COLLECTION || 'tool_catalog';

    if (!this.token || !this.endpoint || !this.dbName) {
      throw new Error('Missing required environment variables: ASTRA_DB_APPLICATION_TOKEN, ASTRA_DB_API_ENDPOINT, and ASTRA_DB_DB_NAME are required');
    }

    try {
      // Initialize the DataAPIClient with the token
      this.client = new DataAPIClient(this.token);
      
      // Get the database using the endpoint and token
      // The endpoint format: https://<database-id>-<region>.apps.astra.datastax.com
      this.db = this.client.db(this.endpoint, { token: this.token });
      
      // Get the collection
      this.collection = this.db.collection(this.collectionName);
    } catch (error) {
      throw new Error(`Failed to connect to Astra DB: ${error instanceof Error ? error.message : String(error)}`);
    }
  }

  async getTools(): Promise<Tool[]> {
    if (!this.collection) {
      await this.connect();
    }

    try {
      // Use find() with filter, which returns a cursor
      const cursor = this.collection.find({ type: 'tool' });
      
      // Convert cursor to array
      const tools: Tool[] = [];
      for await (const doc of cursor) {
        tools.push(doc as Tool);
      }
      
      return tools;
    } catch (error) {
      throw new Error(`Failed to fetch tools: ${error instanceof Error ? error.message : String(error)}`);
    }
  }

  async updateTool(tool: Tool): Promise<void> {
    if (!this.collection) {
      await this.connect();
    }

    try {
      if (tool._id) {
        // Update existing document
        tool.enabled = tool.enabled === false ? false : true;
        const _id = tool._id;
        delete tool._id;
        await this.collection.updateOne(
          { _id: _id },
          { $set: tool }
        );
      } else {
        // Insert new document
        await this.collection.insertOne({ document: tool });
      }
    } catch (error) {
      throw new Error(`Failed to update tool: ${error instanceof Error ? error.message : String(error)}`);
    }
  }

  async getSampleDocuments(tool: Tool, limit: number = 5): Promise<any[]> {
    if (!this.db) {
      await this.connect();
    }

    try {
      const objectType = tool.collection_name ? 'collection' : 'table';
      const objectName = tool.collection_name || tool.table_name;
      const dbName = tool.db_name || this.dbName;

      if (!objectName || !dbName) {
        throw new Error('Collection/table name and database name are required');
      }

      // Get the target database
      const targetDb = this.client!.db(this.endpoint!, { token: this.token });
      
      let targetObject: any;
      if (objectType === 'collection') {
        targetObject = targetDb.collection(objectName);
      } else {
        targetObject = targetDb.table(objectName);
      }

      // Fetch sample documents
      const cursor = targetObject.find({}).limit(limit);
      const documents: any[] = [];
      for await (const doc of cursor) {
        documents.push(doc);
      }

      return documents;
    } catch (error) {
      throw new Error(`Failed to fetch sample documents: ${error instanceof Error ? error.message : String(error)}`);
    }
  }

  async listKeyspaces(): Promise<string[]> {
    if (!this.client) {
      await this.connect();
    }

    try {
      // Get list of databases (keyspaces)
      const admin = this.client!.admin();
      const databases = await admin.listDatabases();
      return databases.map((db: any) => db.name || db.id).filter(Boolean);
    } catch (error) {
      throw new Error(`Failed to list keyspaces: ${error instanceof Error ? error.message : String(error)}`);
    }
  }

  async listCollections(dbName?: string): Promise<string[]> {
    if (!this.db) {
      await this.connect();
    }

    try {
      const targetDbName = dbName || this.dbName;
      if (!targetDbName) {
        throw new Error('Database name is required');
      }

      // Get the database
      const targetDb = this.client!.db(this.endpoint!, { token: this.token });
      
      // List collections
      const collections = await targetDb.listCollections();
      return collections.map((col: any) => col.name || col);
    } catch (error) {
      throw new Error(`Failed to list collections: ${error instanceof Error ? error.message : String(error)}`);
    }
  }

  async listTables(dbName?: string): Promise<string[]> {
    if (!this.db) {
      await this.connect();
    }

    try {
      const targetDbName = dbName || this.dbName;
      if (!targetDbName) {
        throw new Error('Database name is required');
      }

      // Get the database
      const targetDb = this.client!.db(this.endpoint!, { token: this.token });
      
      // List tables
      const tables = await targetDb.listTables();
      return tables.map((table: any) => table.name || table);
    } catch (error) {
      throw new Error(`Failed to list tables: ${error instanceof Error ? error.message : String(error)}`);
    }
  }
}

// Utility function to extract all attributes from documents (including nested)
export function extractAttributes(documents: any[]): string[] {
  const attributes = new Set<string>();

  function extractFromObject(obj: any, prefix: string = '') {
    if (obj === null || obj === undefined) {
      return;
    }

    if (Array.isArray(obj)) {
      // For arrays, check the first element if it's an object
      if (obj.length > 0 && typeof obj[0] === 'object' && obj[0] !== null) {
        extractFromObject(obj[0], prefix);
      }
      return;
    }

    if (typeof obj === 'object') {
      for (const [key, value] of Object.entries(obj)) {
        // Transform keys starting with "$" to "_$"
        const transformedKey = key.startsWith('$') ? `_${key}` : key;
        const fullKey = prefix ? `${prefix}.${transformedKey}` : transformedKey;

        attributes.add(fullKey);

        // Recursively extract nested attributes
        if (value !== null && typeof value === 'object' && !Array.isArray(value)) {
          extractFromObject(value, fullKey);
        } else if (Array.isArray(value) && value.length > 0 && typeof value[0] === 'object') {
          extractFromObject(value[0], fullKey);
        }
      }
    }
  }

  // Extract attributes from all documents
  for (const doc of documents) {
    extractFromObject(doc);
  }

  return Array.from(attributes).sort();
}

// Singleton instance
let astraClientInstance: AstraClient | null = null;

export function getAstraClient(): AstraClient {
  if (!astraClientInstance) {
    astraClientInstance = new AstraClient();
  }
  return astraClientInstance;
}


