import fs from 'node:fs'
import path from 'node:path'

const MAX_FILE_BYTES = 200_000  // 200KB per file

/** Image extensions that should be sent as multimodal attachments, not text. */
const IMAGE_EXTENSIONS = new Set(['.png', '.jpg', '.jpeg', '.gif', '.webp', '.bmp', '.tiff', '.tif'])

const IMAGE_MEDIA_TYPES: Record<string, string> = {
  '.png':  'image/png',
  '.jpg':  'image/jpeg',
  '.jpeg': 'image/jpeg',
  '.gif':  'image/gif',
  '.webp': 'image/webp',
  '.bmp':  'image/bmp',
  '.tiff': 'image/tiff',
  '.tif':  'image/tiff',
}

const MAX_IMAGE_BYTES = 5_000_000  // 5MB per image

export interface ExpansionResult {
  /** Prompt text with file contents inlined. Image @mentions are replaced with [image: name]. */
  prompt: string
  filesIncluded: string[]
  filesNotFound: string[]
  /** Images extracted from @path.png mentions — passed as multimodal attachments. */
  imageAttachments: Array<{ base64: string; mediaType: string; filePath: string }>
}

export function expandFileMentions(rawPrompt: string, cwd: string): ExpansionResult {
  const filesIncluded: string[] = []
  const filesNotFound: string[] = []
  const imageAttachments: Array<{ base64: string; mediaType: string; filePath: string }> = []

  // Match `@<path>` where path contains at least one `/`, `.`, `~` or `\`
  // so it doesn't clash with `@modelAlias` (which has no path separators).
  const re = /@([~/\\][^\s]+|[^\s]*[/\\.][^\s]+)/g

  const expanded = rawPrompt.replace(re, (match, p1: string) => {
    const cleaned = p1.replace(/[,;:.!?)]+$/, '')  // strip trailing punctuation
    const resolvedPath = resolvePath(cleaned, cwd)
    if (!resolvedPath || !fs.existsSync(resolvedPath)) {
      filesNotFound.push(cleaned)
      return match
    }
    try {
      const stat = fs.statSync(resolvedPath)

      // ── Directory listing ──────────────────────────────────────────────
      if (stat.isDirectory()) {
        const entries = fs.readdirSync(resolvedPath).slice(0, 50)
        filesIncluded.push(cleaned + '/')
        return `\n\n--- Contents of ${cleaned}/ ---\n${entries.map(e => '  ' + e).join('\n')}\n--- end ${cleaned}/ ---\n\n`
      }

      // ── Image file → multimodal attachment ────────────────────────────
      const ext = path.extname(resolvedPath).toLowerCase()
      if (IMAGE_EXTENSIONS.has(ext)) {
        if (stat.size > MAX_IMAGE_BYTES) {
          filesNotFound.push(`${cleaned} (image too large: >${Math.round(stat.size / 1024)}KB, max 5MB)`)
          return match
        }
        const base64 = fs.readFileSync(resolvedPath).toString('base64')
        const mediaType = IMAGE_MEDIA_TYPES[ext] ?? 'image/png'
        imageAttachments.push({ base64, mediaType, filePath: cleaned })
        filesIncluded.push(cleaned)
        // Replace with a description so the text prompt makes sense
        return `[image: ${path.basename(cleaned)}]`
      }

      // ── Text/code file → inline as code block ─────────────────────────
      if (stat.size > MAX_FILE_BYTES) {
        filesNotFound.push(`${cleaned} (>${Math.round(MAX_FILE_BYTES / 1024)}KB)`)
        return match
      }
      const content = fs.readFileSync(resolvedPath, 'utf-8')
      filesIncluded.push(cleaned)
      const langExt = ext.slice(1) || ''
      return `\n\n--- ${cleaned} ---\n\`\`\`${langExt}\n${content}\n\`\`\`\n--- end ${cleaned} ---\n\n`
    } catch {
      filesNotFound.push(cleaned)
      return match
    }
  })

  return { prompt: expanded, filesIncluded, filesNotFound, imageAttachments }
}

function resolvePath(p: string, cwd: string): string | null {
  if (p.startsWith('~')) {
    const home = process.env.HOME || process.env.USERPROFILE || ''
    return path.join(home, p.slice(1))
  }
  if (path.isAbsolute(p)) return p
  return path.join(cwd, p)
}