forge/src/index.tsx

import { type TagDef, type HTMLTag, UnitlessProps, NonStyleKeys } from './types'

type ForgeComponent<Parts = {}> =
  ((props: Record<string, any>) => any) &
  { [K in keyof Parts]: (props: Record<string, any>) => any }

export const styles: Record<string, Record<string, string>> = {}
const themes: Record<string, Record<string, any>> = {}

export function clearStyles() {
  for (const key in styles) delete styles[key]
  for (const key in themes) delete themes[key]
}

export function createTheme<const T extends Record<string, string | number>>(name: string, values: T): T {
  themes[name] = values as Record<string, string | number>
  return values
}

// Generate CSS for all registered themes
export function themesToCSS(): string {
  let out: string[] = []

  for (const [name, vars] of Object.entries(themes)) {
    out.push(`[data-theme="${name}"] {`)
    for (const [key, value] of Object.entries(vars)) {
      out.push(`  --theme-${key}: ${value};`)
    }
    out.push(`}\n`)
  }

  return out.join('\n')
}

// Helper type to extract theme keys from multiple theme objects
type ThemeKeys<T> = T extends Record<string, any> ? keyof T : never

// Create a typed themeVar function from your themes
export function createThemedVar<T extends Record<string, any>>(_themes: T) {
  return function themeVar<K extends ThemeKeys<T[keyof T]>>(name: K): string {
    return `var(--theme-${name as string})`
  }
}

type Theme = Record<string, string | number>

export function createThemes<T extends Record<string, Theme>>(themeDefs: T) {
  for (const [name, values] of Object.entries(themeDefs))
    createTheme(name, values)

  return (name: keyof T[keyof T]) => `var(--theme-${name as string})`
}

export function extendThemes(overrides: Record<string, Theme>) {
  for (const [name, values] of Object.entries(overrides))
    themes[name] = { ...themes[name], ...values }

  return (name: string) => `var(--theme-${name})`
}

// Generic themeVar (untyped fallback)
export function themeVar(name: string): string {
  return `var(--theme-${name as string})`
}

// All CSS styles inside <style></style.
// Use w/ SSR: <Styles/>
export const Styles = () => <style dangerouslySetInnerHTML={{ __html: stylesToCSS() }} />

const isBrowser = typeof document !== 'undefined'
let styleElement: HTMLStyleElement | null = null

// automatically inject <style> tag into browser for SPA
function injectStylesInBrowser() {
  if (!isBrowser) return

  styleElement ??= document.getElementById('forge-styles') as HTMLStyleElement

  if (!styleElement) {
    styleElement = document.createElement('style')
    styleElement.id = 'forge-styles'
    document.head.appendChild(styleElement)
  }

  styleElement.textContent = stylesToCSS()
}

// turns style object into string CSS definition
export function stylesToCSS(): string {
  let out: string[] = []

  // Include theme CSS first
  const themeCSS = themesToCSS()
  if (themeCSS) {
    out.push(themeCSS)
    out.push('\n')
  }

  // Then component styles
  for (const [selector, style] of Object.entries(styles)) {
    if (Object.keys(style).length === 0) continue

    out.push(`${selector} { `)
    for (const [name, value] of Object.entries(style).sort(([a], [b]) => a.localeCompare(b)))
      out.push(`  ${name}: ${value};`)
    out.push(`}\n`)
  }

  return out.join('\n')
}

// creates a CSS class name
function makeClassName(baseName: string, partName?: string, variantName?: string, variantKey?: string): string {
  const cls = partName ? `${baseName}_${partName}` : baseName

  if (variantName && variantKey) {
    return cls + `.${variantName}-${variantKey}`
  } else {
    return cls
  }
}

// 'fontSize' => 'font-size'
function camelToDash(name: string): string {
  let out = ''

  for (const letter of name.split(''))
    out += letter.toUpperCase() === letter ? `-${letter.toLowerCase()}` : letter

  return out
}

// turns a TagDef into a JSX style object
function makeStyle(def: TagDef) {
  const style: Record<string, string> = {}

  for (const [name, value] of Object.entries(def)) {
    if (NonStyleKeys.has(name)) continue
    style[camelToDash(name)] = `${typeof value === 'number' && !UnitlessProps.has(name) ? `${value}px` : value}`
  }

  return style
}

// turns a TagDef into a JSX component
function makeComponent(baseName: string, rootDef: TagDef, rootProps: Record<string, any>, partName?: string) {
  const def = partName ? rootDef.parts?.[partName]! : rootDef
  const base = def.base ?? 'div'

  // Extract element name from base (e.g., 'input[type=radio]' -> 'input')
  const tagName = base.split('[')[0]
  const Tag = tagName as HTMLTag

  // Extract attributes from base (e.g., 'input[type=radio]' -> { type: 'radio' })
  const baseAttrs: Record<string, string> = {}
  const attrMatch = base.match(/\[([^\]]+)\]/)
  if (attrMatch && attrMatch[1]) {
    const attrStr = attrMatch[1]
    const [attrName, attrValue] = attrStr.split('=')
    if (attrName && attrValue) {
      baseAttrs[attrName] = attrValue
    }
  }

  return ({ children, ...props }: { children: any, [key: string]: any }) => {
    const classNames = [makeClassName(baseName, partName)]

    const allProps = { ...rootProps, ...props }

    for (const [key, value] of Object.entries(allProps)) {
      const variantConfig = rootDef.variants?.[key]
      if (!variantConfig) continue

      // Remove variant prop from being passed to HTML element
      delete props[key]

      const variantName = key
      const variantKey = value

      let variantDef: TagDef | undefined
      // Distinguish boolean variants from keyed variants:
      // - Boolean variants: component({ variant: true }) → variantConfig is a TagDef
      // - Keyed variants: component({ variant: 'key' }) → variantConfig[key] is a TagDef
      if (value === true) {
        variantDef = variantConfig as TagDef
      } else if (typeof value === 'string') {
        variantDef = (variantConfig as Record<string, TagDef>)[value]
      }
      if (!variantDef) continue

      classNames.push(variantKey === true ? variantName : `${variantName}-${variantKey}`)
    }

    const finalProps = { class: classNames.join(' '), ...baseAttrs, ...props, children }

    if (finalProps.dangerouslySetInnerHTML) {
      const { children: _, ...rest } = finalProps
      return <Tag {...rest} />
    }

    const content = (partName && def.render) ? def.render(finalProps) : children

    return <Tag {...finalProps}>{content}</Tag>
  }
}

// ensures 'hover' is ':hover'
function stateName(state: string): string {
  return state.startsWith(':') ? state : `:${state}`
}

// Register base styles, selectors, and states for a class
function registerClassStyles(name: string, className: string, def: TagDef) {
  styles[`.${className}`] ??= makeStyle(def)

  for (let [selector, selectorDef] of Object.entries(def.selectors ?? {})) {
    selector = selector.replace(/@(\w+)/g, (_, partName) => `.${makeClassName(name, partName)}`)
    selector = selector.replace('&', `.${className}`)
    if (styles[selector]) throw `${selector} already defined!`
    styles[selector] = makeStyle(selectorDef)
  }

  for (const [state, style] of Object.entries(def.states ?? {}))
    styles[`.${className}${stateName(state)}`] = makeStyle(style)
}

// adds CSS styles for tag definition
function registerStyles(name: string, def: TagDef) {
  const rootClassName = makeClassName(name)
  registerClassStyles(name, rootClassName, def)

  for (const [partName, partDef] of Object.entries(def.parts ?? {})) {
    const partClassName = makeClassName(name, partName)
    registerClassStyles(name, partClassName, partDef)
  }

  for (const [variantName, variantConfig] of Object.entries(def.variants ?? {})) {
    // Detect boolean vs keyed variants by checking if config has structural keys or looks like a TagDef
    const isBooleanVariant = 'parts' in variantConfig || 'styles' in variantConfig || 'states' in variantConfig ||
      // If first key is camelCase or contains CSS-like properties, treat as boolean variant
      Object.keys(variantConfig).some(k => k !== k.toLowerCase() || typeof (variantConfig as any)[k] !== 'object')

    if (isBooleanVariant) {
      // Boolean variant - variantConfig is a TagDef
      const variantDef = variantConfig as TagDef
      const baseClassName = makeClassName(name)
      const className = `${baseClassName}.${variantName}`
      registerClassStyles(name, className, variantDef)

      for (const [partName, partDef] of Object.entries(variantDef.parts ?? {})) {
        const basePartClassName = makeClassName(name, partName)
        // Direct class on part (for render({ parts }) path where parts inherit variant props)
        const partClassName = `${basePartClassName}.${variantName}`
        registerClassStyles(name, partClassName, partDef)
        // Descendant selector (for Component.Part path where root has variant class)
        const descendantClassName = `${baseClassName}.${variantName} .${basePartClassName}`
        registerClassStyles(name, descendantClassName, partDef)
      }
    } else {
      // Keyed variant - iterate over the keys
      for (const [variantKey, variantDef] of Object.entries(variantConfig as Record<string, TagDef>)) {
        const className = makeClassName(name, undefined, variantName, variantKey)
        registerClassStyles(name, className, variantDef)

        for (const [partName, partDef] of Object.entries(variantDef.parts ?? {})) {
          const basePartClassName = makeClassName(name, partName)
          // Direct class on part (for render({ parts }) path)
          const partClassName = makeClassName(name, partName, variantName, variantKey)
          registerClassStyles(name, partClassName, partDef)
          // Descendant selector (for Component.Part path)
          const descendantClassName = `${className} .${basePartClassName}`
          registerClassStyles(name, descendantClassName, partDef)
        }
      }
    }
  }

  // In browser, inject styles into DOM immediately
  injectStylesInBrowser()
}

// module-level scoping
export function createScope(scope: string) {
  return {
    define: (nameOrDef: string | TagDef, defIfNamed?: TagDef) => {
      if (typeof nameOrDef === 'string')
        return define(`${scope}${nameOrDef === 'Root' ? '' : nameOrDef}`, defIfNamed)
      else
        return define(`${scope}${anonName(nameOrDef)}`, nameOrDef as TagDef)
    }
  }
}

// the main event
export function define<const T extends TagDef>(name: string, def: T): ForgeComponent<T['parts']>
export function define<const T extends TagDef>(def: T): ForgeComponent<T['parts']>
export function define(nameOrDef: string | TagDef, defIfNamed?: TagDef) {
  const def = defIfNamed ?? nameOrDef as TagDef
  const name = defIfNamed ? (nameOrDef as string) : anonName(def)

  if (styles[`.${name}`]) throw `${name} is already defined! Must use unique names.`
  registerStyles(name, def)

  const currentProps: Record<string, any> = {}
  const Root = makeComponent(name, def, currentProps)
  const parts: Record<string, Function> = { Root }

  for (const [part] of Object.entries(def.parts ?? {}))
    parts[part] = makeComponent(name, def, currentProps, part)

  const component = (props: Record<string, any>) => {
    for (const key in currentProps) delete currentProps[key]
    Object.assign(currentProps, props)

    return def.render?.({ props, parts }) ?? <Root {...props}>{props.children}</Root>
  }

  // Attach parts as static properties for Component.Part syntax
  for (const [part] of Object.entries(def.parts ?? {}))
    (component as any)[part] = makeComponent(name, def, {}, part)

  return component
}

// automatic names
const anonComponents: Record<string, number> = {}

// div tag -> Div1
function anonName(def: TagDef): string {
  const base = (def.base ?? 'div')
  const count = (anonComponents[base] ??= 1)
  anonComponents[base] += 1
  return tagName(base) + (count === 1 ? '' : String(count))
}

// a => Anchor, nav => Nav
function tagName(base: string): string {
  const capitalized = base.slice(0, 1).toUpperCase() + base.slice(1)
  return capitalized === 'A' ? 'Anchor' : capitalized
}

// global CSS for resets, element defaults, pseudo-elements
export function css(rules: Record<string, TagDef>) {
  for (const [selector, def] of Object.entries(rules)) {
    styles[selector] = makeStyle(def)

    for (const [state, stateDef] of Object.entries(def.states ?? {}))
      styles[`${selector}${stateName(state)}`] = makeStyle(stateDef)
  }

  injectStylesInBrowser()
}

// shortcut so you only have to import one thing, if you want
define.Styles = Styles
define.css = css