/**
 * RAG Data Safety Test Fixtures
 * Tests for detecting data exfiltration risks in RAG systems
 */

import type { TestGroup } from '../../types'

export const aiRagSafetyTests: TestGroup = {
  name: 'RAG Data Safety',
  tier: 'A',
  layer: 2,
  description: 'Detection of cross-tenant data access and context exposure in RAG systems',

  truePositives: [
    {
      name: 'RAG Safety - True Positives',
      expectFindings: true,
      expectedCategories: ['ai_rag_exfiltration'],
      description: 'RAG safety issues that MUST be detected',
      file: {
        path: 'src/api/rag/unsafe-retrieval.ts',
        content: `
import { Pinecone } from '@pinecone-database/pinecone'
import { OpenAI } from 'openai'

const pinecone = new Pinecone()
const openai = new OpenAI()

// Unscoped vector query - no user/tenant filtering - HIGH
export async function unscopedQuery(query: string) {
  const index = pinecone.index('documents')
  const embedding = await openai.embeddings.create({
    model: 'text-embedding-3-small',
    input: query
  })

  // No filter! Returns documents from ALL users
  const results = await index.query({
    vector: embedding.data[0].embedding,
    topK: 10,
  })

  return results.matches
}

// Raw context returned in response - MEDIUM
export async function exposeRawContext(req: Request) {
  const { query } = await req.json()
  const docs = await retrieveDocs(query)

  // Exposing full document content to client
  return Response.json({
    answer: 'Here is the answer',
    sourceDocuments: docs,  // RAW CONTEXT EXPOSURE
    chunks: docs.map(d => d.pageContent)
  })
}

// LangChain retriever without filter - HIGH
import { VectorStoreRetriever } from 'langchain/vectorstores'

export async function langchainUnscopedRetrieval(query: string) {
  const retriever = new VectorStoreRetriever({ vectorStore })
  // No metadata filter - gets all documents
  const docs = await retriever.invoke(query)
  return docs
}

// Logging retrieved documents - INFO (but still flag)
export async function loggingContext(query: string) {
  const docs = await retrieveDocs(query)
  console.log('Retrieved documents:', docs)  // Logging sensitive content
  console.debug('Context for RAG:', docs.map(d => d.content))
  return processWithLLM(docs, query)
}

// Chroma query without where clause - MEDIUM
import { Chroma } from 'chromadb'

export async function chromaUnscopedQuery(query: string) {
  const collection = await chroma.getCollection('documents')
  // No where filter - queries all documents
  const results = await collection.query({
    query_texts: [query],
    n_results: 10
  })
  return results
}
`,
        language: 'typescript',
        size: 1800,
      },
    },
    // AI Detection Roadmap Phase 1: Enhanced RAG Detection
    {
      name: 'RAG Safety - Corpus Poisoning True Positives',
      expectFindings: true,
      expectedCategories: ['ai_rag_corpus_poisoning'],
      description: 'User uploads directly embedded without sanitization',
      file: {
        path: 'src/api/rag/upload-handler.ts',
        content: `
import { Pinecone } from '@pinecone-database/pinecone'
import { OpenAI } from 'openai'

const pinecone = new Pinecone()
const openai = new OpenAI()

// Dangerous: User upload embedded directly without sanitization
export async function handleDocumentUpload(req: Request) {
  const { userId, document } = await req.json()

  // Direct embedding of user content - no sanitization!
  const embedding = await openai.embeddings.create({
    model: 'text-embedding-3-small',
    input: document.content  // Raw user content embedded
  })

  const index = pinecone.index('knowledge-base')
  await index.upsert([{
    id: document.id,
    values: embedding.data[0].embedding,
    metadata: { userId, content: document.content }
  }])

  return { success: true }
}

// Dangerous: External content fetched and embedded without validation
export async function indexExternalContent(url: string) {
  const response = await fetch(url)
  const content = await response.text()

  // External content embedded without sanitization
  const embedding = await createEmbedding(content)
  await vectorStore.addDocument({ content, embedding })
}

// Dangerous: User-provided PDF indexed without scanning
export async function indexUploadedPdf(pdfBuffer: Buffer, userId: string) {
  const pdfText = await pdfParser.parse(pdfBuffer)

  // PDF content could contain injection instructions
  await addToCorpus(pdfText, { userId })
}

// Dangerous: Slack messages indexed without filtering
export async function indexSlackChannel(channelId: string) {
  const messages = await slack.getMessages(channelId)

  // User messages embedded - could contain injections
  for (const msg of messages) {
    await embedDocument(msg.text, { author: msg.user })
  }
}
`,
        language: 'typescript',
        size: 1500,
      },
    },
    {
      name: 'RAG Safety - PII Leakage True Positives',
      expectFindings: true,
      expectedCategories: ['ai_rag_pii_leakage'],
      description: 'PII fields exposed in RAG retrieval responses',
      file: {
        path: 'src/api/rag/search-handler.ts',
        content: `
import { VectorStore } from 'langchain/vectorstores'

// Dangerous: PII fields in embedded documents
export async function indexUserProfile(user: User) {
  const doc = {
    content: user.bio,
    metadata: {
      userId: user.id,
      email: user.email,           // PII in metadata!
      ssn: user.ssn,               // Highly sensitive PII!
      phoneNumber: user.phone,     // PII in metadata!
      fullName: user.fullName,     // PII in metadata!
      dateOfBirth: user.dob,       // PII in metadata!
      address: user.homeAddress,   // PII in metadata!
    }
  }

  await vectorStore.addDocument(doc)
}

// Dangerous: Returning PII in retrieval response
export async function searchDocuments(query: string) {
  const results = await vectorStore.similaritySearch(query)

  // Returning full metadata including PII
  return results.map(doc => ({
    content: doc.pageContent,
    email: doc.metadata.email,     // PII in response!
    phone: doc.metadata.phone,     // PII in response!
    ssn: doc.metadata.ssn,         // PII in response!
    author: doc.metadata.fullName, // PII in response!
  }))
}

// Dangerous: Medical records embedded with patient info
export async function indexMedicalRecords(records: MedicalRecord[]) {
  for (const record of records) {
    await embedDocument(record.diagnosis, {
      patientName: record.patientName,    // PHI!
      patientDob: record.dateOfBirth,     // PHI!
      patientSsn: record.ssn,             // PHI!
      insuranceId: record.insuranceNumber // PHI!
    })
  }
}

// Dangerous: Credit card info in indexed data
export async function indexTransactions(transactions: Transaction[]) {
  for (const tx of transactions) {
    await vectorStore.addDocument({
      content: tx.description,
      metadata: {
        cardNumber: tx.cardLastFour,  // PCI data!
        cvv: tx.cvv,                   // PCI data - CRITICAL!
        accountNumber: tx.accountNum,  // Financial PII!
      }
    })
  }
}
`,
        language: 'typescript',
        size: 1800,
      },
    },
  ],

  falseNegatives: [
    {
      name: 'RAG Safety - False Negatives',
      expectFindings: false,
      description: 'Safe RAG patterns that should NOT be flagged',
      allowedInfoFindings: [
        {
          category: 'ai_rag_exfiltration',
          maxCount: 2,
          reason: 'Some patterns may generate info-level findings for hygiene notes',
        },
      ],
      file: {
        path: 'src/api/rag/safe-retrieval.ts',
        content: `
import { Pinecone } from '@pinecone-database/pinecone'

const pinecone = new Pinecone()

// Properly scoped query with user filter - SAFE
export async function scopedQuery(query: string, userId: string) {
  const index = pinecone.index('documents')
  const embedding = await getEmbedding(query)

  // Properly filtered by userId
  const results = await index.query({
    vector: embedding,
    topK: 10,
    filter: {
      userId: { $eq: userId }
    }
  })

  return results.matches
}

// Filtered response - only returning safe fields - SAFE
export async function filteredResponse(req: Request) {
  const { query } = await req.json()
  const docs = await retrieveDocs(query)

  // Only returning IDs and titles, not full content
  return Response.json({
    answer: 'Here is the answer',
    sources: docs.map(d => ({
      id: d.id,
      title: d.metadata.title,
      // NOT including pageContent or full document
    }))
  })
}

// LangChain with metadata filter - SAFE
export async function langchainScopedRetrieval(query: string, tenantId: string) {
  const retriever = vectorStore.asRetriever({
    filter: { tenantId },  // Properly scoped
    k: 5
  })
  const docs = await retriever.invoke(query)
  return docs
}

// Supabase with RLS - SAFE (RLS handles authorization)
export async function supabaseRagQuery(query: string) {
  // RLS policies enforce user scoping
  const { data } = await supabase.rpc('match_documents', {
    query_embedding: embedding,
    match_count: 10
  })
  return data
}

// Chroma with where filter - SAFE
export async function chromaScopedQuery(query: string, userId: string) {
  const collection = await chroma.getCollection('documents')
  const results = await collection.query({
    query_texts: [query],
    where: { userId: userId },  // Properly filtered
    n_results: 10
  })
  return results
}
`,
        language: 'typescript',
        size: 1700,
      },
    },
  ],
}
