import { extname, relative } from 'node:path'
import BabelGenerate from '@babel/generator'
import { parse } from '@babel/parser'
import BabelTraverse from '@babel/traverse'
import type * as t from '@babel/types'
import {
  deadCodeElimination,
  findReferencedIdentifiers,
} from 'babel-dead-code-elimination'
import type { Plugin } from 'vite'
import { EMPTY_LOADER_STRING, makeLoaderRouteIdStub } from '../constants'

const traverse = (BabelTraverse['default'] || BabelTraverse) as typeof BabelTraverse
const generate = (BabelGenerate['default'] ||
  BabelGenerate) as any as typeof BabelGenerate

// Collect type-only imports before dead code elimination runs
// These should never be removed since TypeScript erases them at compile time
function collectTypeImports(ast: t.File): t.ImportDeclaration[] {
  const typeImports: t.ImportDeclaration[] = []
  traverse(ast, {
    ImportDeclaration(path) {
      // Check if the entire import is type-only: `import type { X } from '...'`
      if (path.node.importKind === 'type') {
        typeImports.push(path.node)
      }
    },
  })
  return typeImports
}

// Restore type-only imports that may have been removed by dead code elimination
function restoreTypeImports(ast: t.File, typeImports: t.ImportDeclaration[]) {
  if (typeImports.length === 0) return

  // Get existing import sources to avoid duplicates
  const existingSources = new Set<string>()
  traverse(ast, {
    ImportDeclaration(path) {
      if (path.node.importKind === 'type') {
        existingSources.add(path.node.source.value)
      }
    },
  })

  // Add back any type imports that were removed
  for (const typeImport of typeImports) {
    if (!existingSources.has(typeImport.source.value)) {
      ast.program.body.unshift(typeImport)
    }
  }
}

export const clientTreeShakePlugin = (opts?: {
  // 'rolldown' when used in the native rolldown DevEngine (no vite environment context)
  runtime?: 'vite' | 'rolldown'
}): Plugin => {
  const runtime = opts?.runtime ?? 'vite'

  return {
    name: 'one-client-tree-shake',

    enforce: 'pre',

    ...(runtime === 'vite' && {
      applyToEnvironment(env: { name: string }) {
        return env.name === 'client' || env.name === 'ios' || env.name === 'android'
      },
    }),

    transform: {
      order: 'pre',
      async handler(code, id, settings) {
        if (runtime === 'vite' && this.environment?.name === 'ssr') {
          return
        }
        if (!/\.(js|jsx|ts|tsx)/.test(extname(id))) {
          return
        }
        if (/node_modules/.test(id)) {
          return
        }

        const out = await transformTreeShakeClient(code, id, process.cwd())

        return out
      },
    },
  } satisfies Plugin
}

export async function transformTreeShakeClient(code: string, id: string, root?: string) {
  if (!/generateStaticParams|loader/.test(code)) {
    return
  }

  let ast: any
  try {
    // `as any` because babel-dead-code-elimination using @types and it conflicts :/
    ast = parse(code, {
      sourceType: 'module',
      plugins: ['typescript', 'jsx'],
    }) as any
  } catch (error) {
    // If there's a syntax error, skip transformation and let Vite handle the error
    // This prevents the dev server from crashing on syntax errors
    const errorMessage = error instanceof Error ? error.message : String(error)
    console.warn(
      `[one] Skipping tree shaking for ${id} due to syntax error:`,
      errorMessage
    )
    return
  }

  let referenced: any
  try {
    referenced = findReferencedIdentifiers(ast)
  } catch (error) {
    // If finding referenced identifiers fails, skip transformation
    const errorMessage = error instanceof Error ? error.message : String(error)
    console.warn(
      `[one] Skipping tree shaking for ${id} due to identifier analysis error:`,
      errorMessage
    )
    return
  }

  const removed = {
    loader: false,
    generateStaticParams: false,
  }

  // note: only handles inline declarations (export function loader / export const loader)
  // importing a loader from another file (export { loader } from './other') is not supported
  // and will break client-side tree shaking. loaders must be defined in the route file.
  try {
    traverse(ast, {
      ExportNamedDeclaration(path) {
        if (
          path.node.declaration &&
          path.node.declaration.type === 'FunctionDeclaration'
        ) {
          if (!path.node.declaration.id) return
          const functionName = path.node.declaration.id.name
          if (functionName === 'loader' || functionName === 'generateStaticParams') {
            path.remove()
            removed[functionName] = true
          }
        } else if (
          path.node.declaration &&
          path.node.declaration.type === 'VariableDeclaration'
        ) {
          path.node.declaration.declarations.forEach((declarator, index) => {
            if (
              declarator.id.type === 'Identifier' &&
              (declarator.id.name === 'loader' ||
                declarator.id.name === 'generateStaticParams')
            ) {
              const declaration = path.get('declaration.declarations.' + index)
              if (!Array.isArray(declaration) && declaration) {
                ;(declaration as any).remove()
                removed[declarator.id.name] = true
              }
            }
          })
        }
      },
    })
  } catch (error) {
    // If traversal fails, skip transformation
    const errorMessage = error instanceof Error ? error.message : String(error)
    console.warn(
      `[one] Skipping tree shaking for ${id} due to traversal error:`,
      errorMessage
    )
    return
  }

  const removedFunctions = Object.keys(removed).filter((key) => removed[key])

  if (removedFunctions.length) {
    try {
      // Collect type-only imports before dead code elimination
      // These should be preserved since TypeScript erases them at compile time
      const typeImports = collectTypeImports(ast)

      deadCodeElimination(ast, referenced)

      // Restore any type imports that were incorrectly removed
      restoreTypeImports(ast, typeImports)

      const out = generate(ast, { retainLines: true })

      // add back in empty or filled loader and genparams
      const codeOut =
        out.code +
        '\n\n' +
        removedFunctions
          .map((key) => {
            if (key === 'loader') {
              if (root) {
                // compute routeId relative to the app/ directory to match route contextKey format
                // contextKeys are like "./_layout.tsx", "./matches-test/page1+ssg.tsx"
                const fromRoot = relative(root, id).replace(/\\/g, '/')
                const routeId = './' + fromRoot.replace(/^app\//, '')
                return makeLoaderRouteIdStub(routeId)
              }
              return EMPTY_LOADER_STRING
            }

            return `export function generateStaticParams() {};`
          })
          .join('\n')

      console.info(
        ` 🧹 [one]      ${relative(process.cwd(), id)} removed ${removedFunctions.length} server-only exports`
      )

      return {
        code: codeOut,
        map: out.map,
      }
    } catch (error) {
      // If code generation fails, skip transformation
      const errorMessage = error instanceof Error ? error.message : String(error)
      console.warn(
        `[one] Skipping tree shaking for ${id} due to code generation error:`,
        errorMessage
      )
      return
    }
  }
}