import type {BuildResult, Options, ViewStorage} from "@bokehjs/core/build_views"
import {build_views} from "@bokehjs/core/build_views"
import type {HasProps} from "@bokehjs/core/has_props"
import type {ViewOf} from "@bokehjs/core/view"
import type {StyleSheetLike} from "@bokehjs/core/dom"
import type {DOMView} from "@bokehjs/core/dom_view"
import {ClassList, InlineStyleSheet, ImportedStyleSheet} from "@bokehjs/core/dom"
import type {CSSStyles, CSSStyleSheetDecl} from "@bokehjs/core/css"
import type * as p from "@bokehjs/core/properties"
import {difference} from "@bokehjs/core/util/array"
import {assert} from "@bokehjs/core/util/assert"
import {isString} from "@bokehjs/core/util/types"
import type {UIElementView} from "@bokehjs/models/ui/ui_element"
import type {Transform} from "sucrase"

import {
  ReactiveESM, ReactiveESMView, model_getter, model_setter,
} from "./reactive_esm"

export class HostedStyleSheet extends InlineStyleSheet {
  host_id: string

  constructor(css?: string | CSSStyleSheetDecl, id?: string, override readonly persistent: boolean = false, host_id: string = "") {
    super(css, id, persistent)
    this.host_id = host_id
  }

  override replace(css: string, styles?: CSSStyles): void {
    css = css.replace(/:host\b/g, `#${this.host_id}`)
    super.replace(css, styles)
  }

  override prepend(css: string, styles?: CSSStyles): void {
    css = css.replace(/:host\b/g, `#${this.host_id}`)
    super.prepend(css, styles)
  }

  override append(css: string, styles?: CSSStyles): void {
    css = css.replace(/:host\b/g, `#${this.host_id}`)
    super.append(css, styles)
  }

}

async function _build_view<T extends HasProps>(view_cls: T["default_view"], model: T, options: Options<ViewOf<T>>): Promise<ViewOf<T>> {
  assert(view_cls != null, "model doesn't implement a view")
  const view = new view_cls({...options, model})
  view.initialize()
  await view.lazy_initialize()
  return view
}

// Custom build_views implementation which does not eagerly destroy old views
async function build_views_no_remove<T extends HasProps>(
  view_storage: ViewStorage<T>,
  models: T[],
  options: Options<ViewOf<T>> = {parent: null},
  cls: (model: T) => T["default_view"] = (model) => model.default_view,
): Promise<BuildResult<T>> {

  const to_remove = difference([...view_storage.keys()], models)

  const removed_views: ViewOf<T>[] = []
  for (const model of to_remove) {
    const view = view_storage.get(model)
    if (view != null) {
      view_storage.delete(model)
      removed_views.push(view)
    }
  }

  const created_views: ViewOf<T>[] = []
  const new_models = models.filter((model) => !view_storage.has(model))

  for (const model of new_models) {
    const view = await _build_view(cls(model), model, options)
    view_storage.set(model, view)
    created_views.push(view)
  }

  for (const view of created_views) {
    view.connect_signals()
  }

  return {
    created: created_views,
    removed: removed_views,
  }
}

export class ReactComponentView extends ReactiveESMView {
  declare model: ReactComponent
  declare style_cache: HTMLHeadElement
  model_getter = model_getter
  model_setter = model_setter
  react_root: any = null
  mounted: boolean = false

  _force_update_callbacks: (() => void)[] = []
  _mounted_resolve: (() => void) | null = null
  _scheduled_removals: DOMView[] = []

  override initialize(): void {
    super.initialize()
    if (!this.use_shadow_dom) {
      (this as any).display = new HostedStyleSheet("", "display", false, this.model.id);
      (this as any).style = new HostedStyleSheet("", "style", false, this.model.id);
      (this as any).parent_style = new HostedStyleSheet("", "parent", true, this.model.id)
    }
  }

  get use_shadow_dom(): boolean {
    return this.model.use_shadow_dom || !(this.parent instanceof ReactComponentView)
  }

  override render_esm(): void {
    if (this.model.compiled === null || this.model.render_module === null || this.container == null) {
      return
    }
    this._rendered = false
    if (this.model.usesMui) {
      if (this.model.root_node) {
        this.style_cache = document.head
      } else {
        this.style_cache = document.createElement("head")
        this.shadow_el.insertBefore(this.style_cache, this.container)
      }
    }
    this.accessed_properties = []
    for (const lf of this._lifecycle_handlers.keys()) {
      (this._lifecycle_handlers.get(lf) || []).splice(0)
    }
    this.model.disconnect_watchers(this)
    const mounted_promise = new Promise<void>((resolve) => {
      this._mounted_resolve = resolve
    })
    this.model.render_module.then((mod: any) => {
      if (this.container == null) {
        // Nothing will mount, so the ready promise has to be settled here or
        // the view never finishes and the document never goes idle.
        this._resolve_mounted()
        return
      }
      this.react_root = mod.default.render(this.model.id)
    }).catch((e: unknown) => {
      this._resolve_mounted()
      throw e
    })
    this._await_ready(mounted_promise)
  }

  override render_error(error: SyntaxError): void {
    // A component that errored will never mount and so never settles its
    // promise via `after_rendered`.
    this._resolve_mounted()
    super.render_error(error)
  }

  /**
   * Settles the promise handed to `_await_ready` by `render_esm`. Safe to call
   * more than once; only the first call has an effect.
   */
  _resolve_mounted(): void {
    const resolve = this._mounted_resolve
    if (resolve == null) {
      return
    }
    this._mounted_resolve = null
    resolve()
  }

  on_force_update(cb: () => void): void {
    this._force_update_callbacks.push(cb)
  }

  force_update(): void {
    for (const cb of this._force_update_callbacks) {
      cb()
    }
  }

  override remove(): void {
    this._force_update_callbacks = []
    this.mounted = false
    // A view removed before it mounted still owes a resolution to the promise
    // handed to `_await_ready`, otherwise the root never reaches idle.
    this._resolve_mounted()
    if (this.react_root && this.use_shadow_dom) {
      super.remove()
      this.react_root.then((root: any) => root && root.unmount())
      this.flush_scheduled_removals()
    } else {
      this._applied_stylesheets.forEach((stylesheet) => stylesheet.uninstall())
      for (const cb of (this._lifecycle_handlers.get("remove") || [])) {
        cb()
      }
      this._child_callbacks.clear()
      this._child_rendered.clear()
      this._mounted.clear()
    }
    this.react_root = null
  }

  get root_view(): ReactComponentView {
    let root: ReactComponentView = this
    if (this.use_shadow_dom) {
      return root
    }
    while (root.parent instanceof ReactComponentView) {
      root = root.parent
    }
    return root
  }

  protected override _apply_stylesheets(stylesheets: StyleSheetLike[]): void {
    const resolved_stylesheets = stylesheets.map((style) => isString(style) ? new InlineStyleSheet(style) : style)
    const root_view = this.root_view
    const target = root_view.shadow_el
    // When shadow DOM is disabled every component in the tree installs its
    // stylesheets into the same root, so the existing CSS is indexed once and
    // looked up by value. Scanning the root per stylesheet made this quadratic
    // in the number of components times the number of stylesheets each.
    const installed_css = new Set<string | null>()
    const installed_hrefs = new Set<string>()
    if (!this.use_shadow_dom) {
      for (const style of target.querySelectorAll("style")) {
        installed_css.add(style.textContent)
      }
      for (const link of target.querySelectorAll("link")) {
        installed_hrefs.add(link.href)
      }
    }
    resolved_stylesheets.forEach((stylesheet) => {
      if (!this.use_shadow_dom) {
        if (stylesheet instanceof InlineStyleSheet) {
          if (installed_css.has(stylesheet.css)) {
            return
          }
          installed_css.add(stylesheet.css)
        } else if (stylesheet instanceof ImportedStyleSheet) {
          const {href} = (stylesheet as any).el
          if (installed_hrefs.has(href)) {
            return
          }
          installed_hrefs.add(href)
        }
      }
      this._applied_stylesheets.push(stylesheet)
      stylesheet.install(target)
    })
  }

  override render(): void {
    if (this.react_root) {
      this.react_root.then((root: any) => root.unmount())
    }
    this._force_update_callbacks = []
    this.mounted = false
    // `super.render()` calls `render_esm`, which installs a fresh promise, so
    // settle any promise the previous render left outstanding first.
    this._resolve_mounted()
    super.render()
  }

  override r_after_render(): void {
    // If the DOM node was re-inserted, e.g. due to the parent
    // children changing order we must force an update in the
    // React component to ensure anything depending on the DOM
    // structure (e.g. emotion caches) is updated
    if (!this.model.use_shadow_dom) { this._apply_visible() }
    super.r_after_render()
    if (this.use_shadow_dom) {
      this.force_update()
    }
  }

  override _update_layout(): void {
    super._update_layout()
    const handlers = (this._lifecycle_handlers.get("update_layout") || [])
    for (const cb of handlers) {
      cb()
    }
  }

  override async build_child_views(): Promise<UIElementView[]> { // TODO BuildResult<UIElement>
    const build_fn = this.model.use_shadow_dom ? build_views_no_remove : build_views
    const {created, removed} = await build_fn(this._child_views, this.child_models, {parent: this})

    for (const view of removed) {
      this._resize_observer.unobserve(view.el)
      if (this.model.use_shadow_dom) {
        this._child_rendered.delete(view)
        if (created.length) {
          this._scheduled_removals.push(view)
        } else {
          view.remove()
        }
      }
    }

    for (const view of created) {
      this._resize_observer.observe(view.el, {box: "border-box"})
    }

    return created
  }

  protected override async _update_children_pass(): Promise<void> {
    const created_children = new Set(await this.build_child_views())

    const new_views = new Map()
    for (const child_view of this.child_views) {
      if (!created_children.has(child_view)) {
        continue
      }
      const child = this._lookup_child(child_view)
      if (!child) {
        continue
      }

      if (new_views.has(child)) {
        new_views.get(child).push(child_view)
      } else {
        new_views.set(child, [child_view])
      }
    }

    for (const child of this.model.children) {
      const callbacks = this._child_callbacks.get(child) || []
      const new_children = new_views.get(child) || []
      if (new_children) {
        for (const callback of callbacks) {
          callback(new_children)
        }
      }
    }
    this._update_children()
    // Removals are normally drained by the replacement child once it mounts,
    // but a child that never mounts would leak the old views, so flush any
    // that are still pending after React has had a chance to commit.
    setTimeout(() => this.flush_scheduled_removals(), 0)
  }

  flush_scheduled_removals(): void {
    const removals = this._scheduled_removals
    this._scheduled_removals = []
    for (const view of removals) {
      if (!view.is_destroyed) {
        view.remove()
      }
    }
  }

  override _on_mounted(): void {
    this.invalidate_layout()
    this.mounted = true
  }

  override has_finished(): boolean {
    return super.has_finished() && this._rendered
  }

  patch_container(container: HTMLDivElement): void {
    this.el = this.container = container
    this._update_stylesheets()
    this.class_list = new ClassList(this.container.classList)
    this._apply_html_attributes()
  }

  override after_rendered(): void {
    const handlers = (this._lifecycle_handlers.get("after_render") || [])
    for (const cb of handlers) {
      cb()
    }
    if (!this._rendered) {
      for (const cb of (this._lifecycle_handlers.get("after_layout") || [])) {
        cb()
      }
    }
    this._rendered = true
    if (this._mounted_resolve) {
      const child_ready: Promise<void>[] = []
      for (const child_view of this.child_views) {
        child_ready.push(child_view.ready)
      }
      if (child_ready.length > 0) {
        Promise.all(child_ready).then(() => this._resolve_mounted())
      } else {
        this._resolve_mounted()
      }
    }
    this.finish()
  }
}

export namespace ReactComponent {
  export type Attrs = p.AttrsOf<Props>

  export type Props = ReactiveESM.Props & {
    root_node: p.Property<string | null>
    use_shadow_dom: p.Property<boolean>
  }
}

export interface ReactComponent extends ReactComponent.Attrs {}

export class ReactComponent extends ReactiveESM {
  declare properties: ReactComponent.Props
  override sucrase_transforms: Transform[] = ["typescript", "jsx"]

  constructor(attrs?: Partial<ReactComponent.Attrs>) {
    super(attrs)
  }

  get usesMui(): boolean {
    if (this.importmap?.imports) {
      return Object.keys(this.importmap?.imports).some(k => k.startsWith("@mui"))
    }
    return false
  }

  protected override _render_code(): string {
    let [import_code, bundle_code] = ["", ""]
    const cache_key = (this.bundle === "url") ? this.esm : (this.bundle || `${this.class_name}-${this.esm.length}`)
    if (this.bundle) {
      bundle_code = `
  const ns = await view._module_cache.get("${cache_key}")
  const {React, createRoot} = ns.default`
    } else {
      import_code = `
import * as React from "react"
import { createRoot } from "react-dom/client"`
    }
    let init_code = ""
    let render_code = ""
    if (this.usesMui) {
      if (this.bundle) {
        bundle_code = `
  const ns = await view._module_cache.get("${cache_key}")
  const {CacheProvider, React, createCache, createRoot} = ns.default`
      } else {
        import_code = `
${import_code}
import createCache from "@emotion/cache"
import { CacheProvider } from "@emotion/react"`
      }
      init_code = `
  if (view.use_shadow_dom) {
    const css_key = id.replace("-", "").replace(/\d/g, (digit) => String.fromCharCode(digit.charCodeAt(0) + 49)).toLowerCase()
    view.mui_cache = createCache({
      key: 'css-'+css_key,
      prepend: true,
      container: view.style_cache,
    })
  }`
      render_code = `
  if (rendered && ((view.parent?.react_root === undefined) || view.model.use_shadow_dom)) {
    rendered = React.createElement(CacheProvider, {value: view.mui_cache}, rendered)
  }`
    }
    return `
${import_code}

async function render(id) {
  const view = Bokeh.index.find_one_by_id(id)
  if (view == null) {
    return null
  }

  ${bundle_code}

  class Child extends React.PureComponent {

    state = {rendered: null}

    constructor(props) {
      super(props)
      this.render_callback = null
      this.containerRef = React.createRef()
      // Registers the child as tracked but not yet rendered. React's render
      // phase has to stay free of side effects, since a render may be
      // discarded without ever committing, so the flag is only ever flipped
      // here and from getSnapshotBeforeUpdate.
      this._mark_stale()
    }

    _mark_stale() {
      const view = this.view
      if (view) {
        this.props.parent._child_rendered.set(view, false)
      }
    }

    updateElement() {
      const childView = this.view
      const el = childView?.el
      if (el && this.containerRef.current && !this.containerRef.current.contains(el)) {
        this.containerRef.current.innerHTML = ""
        this.containerRef.current.appendChild(el)
      }
    }

    get view() {
      const child = this.props.parent.model.data[this.props.name]
      const model = this.props.id == null ? child : child.find(item => item.id === this.props.id)
      return this.props.parent.get_child_view(model)
    }

    get element() {
      const view = this.view
      return view == null ? null : view.el
    }

    get use_shadow_dom() {
      return this.view?.model.use_shadow_dom || (this.view?.react_root === undefined)
    }

    componentDidMount() {
      const view = this.view
      this.render_callback = (new_views) => {
        const view = this.view
        if (!view || !new_views.includes(view)) {
          return
        }
        this.updateElement()
        if (this.use_shadow_dom) {
          this.props.parent.flush_scheduled_removals()
          this.props.parent.rerender_(view)
          this.props.parent._child_rendered.set(view, true)
        } else {
          view.patch_container(this.containerRef.current)
          view.model.render_module.then(async (mod) => {
            this.props.parent.flush_scheduled_removals()
            this.setState(
              {rendered: await mod.default.render(view.model.id)},
              () => {
                this.props.parent.notify_mount(this.props.name, view.model.id)
                this.view.r_after_render()
                this.view.after_rendered()
              }
            )
          })
        }
      }
      this.props.parent.on_child_render(this.props.name, this.render_callback)
      if (view == null) { return }
      this.props.parent.flush_scheduled_removals()
      if (this.use_shadow_dom) {
        this.updateElement()
        this.props.parent.rerender_(view)
        this.props.parent._child_rendered.set(view, true)
        this.props.parent.notify_mount(this.props.name, view.model.id)
      } else {
        view.patch_container(this.containerRef.current)
        view.model.render_module.then(async (mod) => {
          this.setState(
            {rendered: await mod.default.render(view.model.id)},
            () => {
              this.props.parent.notify_mount(this.props.name, view.model.id)
              this.view.r_after_render()
              this.view.after_rendered()
            }
          )
        })
      }
    }

    componentWillUnmount() {
      if (this.render_callback) {
        this.props.parent.remove_on_child_render(this.props.name, this.render_callback)
      }
      // The mount bookkeeping lives on the parent, and the view may already be
      // gone by the time React unmounts us, so fall back to the model id prop.
      const id = this.view?.model.id ?? this.props.id
      if (id != null) {
        this.props.parent.notify_mount(this.props.name, id, true)
      }
    }

    getSnapshotBeforeUpdate() {
      // Commit-phase equivalent of the constructor's registration: the view
      // this Child renders may have been swapped out, so the incoming one is
      // registered as stale here rather than during render.
      this._mark_stale()
      return null
    }

    componentDidUpdate() {
      if (this.use_shadow_dom) {
        this.updateElement()
      }
    }

    render() {
      const child = this.state.rendered
      const class_name = (this.use_shadow_dom ?
        "child-wrapper" : this.view.model.class_name.replace(/([a-z])([A-Z])/g, "$1-$2").toLowerCase()
      )
      return React.createElement('div', {id: this.view?.model.id, className: class_name, ref: this.containerRef}, child)
    }
  }

  function react_getter(target, name) {
    if (name === "useMount") {
      return (callback) => React.useEffect(() => {
        target.model_proxy.on('lifecycle:mounted', callback)
        return () => target.model_proxy.off('lifecycle:mounted', callback)
      }, [])
    } if (name == "useState") {
      return (prop) => {
        const data_model = target.model.data
        const propPath = prop.split(".")
        let targetModel = data_model
        let resolvedProp = null

        for (let i = 0; i < propPath.length - 1; i++) {
          if (targetModel && targetModel.properties && propPath[i] in targetModel.properties) {
            targetModel = targetModel[propPath[i]]
          } else {
            // Stop if any part of the path is missing
            targetModel = null
            break
          }
        }
        if (targetModel && targetModel.attributes && propPath[propPath.length - 1] in targetModel.attributes) {
          resolvedProp = propPath[propPath.length - 1]
        }
        if (resolvedProp && targetModel) {
          const [value, setValue] = React.useState(targetModel.attributes[resolvedProp])

          React.useEffect(() => {
            const cb = () => {
              if (target.model.events.includes(resolvedProp)) {
                targetModel.attributes[resolvedProp] && (setValue((v) => v+1) || targetModel.setv({[resolvedProp]: false}))
              } else {
                setValue(targetModel.attributes[resolvedProp])
              }
            }
            react_proxy.on(prop, cb, true)
            return () => react_proxy.off(prop, cb)
          }, [])

          let initialized = React.useRef(false)
          React.useEffect(() => {
            if (!target.model.events.includes(resolvedProp) && initialized.current) {
              targetModel.setv({ [resolvedProp]: value })
            } else { initialized.current = true }
          }, [value])

          return [value, setValue]
        }
        throw ReferenceError("Could not resolve " + prop + " on " + target.model.class_name)
      }
    } else if (name === "get_child") {
      return (child) => {
        const data_model = target.model.data
        const value = data_model.attributes[child]
        if (Array.isArray(value)) {
          const [children_state, set_children] = React.useState(value.map((model) =>
            React.createElement(Child, { parent: target, name: child, key: model.id, id: model.id })
          ))
          React.useEffect(() => {
            target.on_child_render(child, () => {
              const current_models = data_model.attributes[child]
              const previous_models = children_state.map(child => child.props.id)
              if (current_models.some((model, i) => model.id !== previous_models[i])) {
                set_children(current_models.map((model, i) => (
                  React.createElement(Child, { parent: target, name: child, key: model.id, id: model.id })
                )))
              }
            })
          }, [])
          return children_state
        } else {
          return React.createElement(Child, {parent: target, name: child})
        }
      }
    }
    return target.model_getter(target, name)
  }

  const react_proxy = new Proxy(view, {
    get: react_getter,
    set: view.model_setter
  })

  class ErrorBoundary extends React.Component {
    constructor(props) {
      super(props)
      // initialize the error state
      this.state = { hasError: false }
    }

    // if an error happened, set the state to true
    static getDerivedStateFromError(error) {
      return { hasError: true }
    }

    componentDidCatch(error) {
      this.props.view.render_error(error)
    }

    render() {
      if (this.state.hasError) {
        return React.createElement('div')
      }
      return React.createElement('div', {className: "error-wrapper"}, this.props.children)
    }
  }

  class Component extends React.Component {

    constructor(props) {
      super(props)
      ${init_code}
    }

    componentDidMount() {
      if (!this.props.view.use_shadow_dom) { return }
      this.props.view.on_force_update(() => {
        ${init_code}
        this.forceUpdate()
      })
      this.props.view.after_rendered()
    }

    render() {
      let rendered = React.createElement(this.props.view.render_fn, this.props)
      if (this.props.view.model.dev) {
        rendered = React.createElement(ErrorBoundary, {view}, rendered)
      }
      ${render_code}
      return rendered
    }
  }

  const props = {view, model: react_proxy, data: view.model.data, el: view.container}
  const rendered = React.createElement(Component, props)
  if (!view.model.use_shadow_dom && (view.parent?.react_root !== undefined)) {
    return rendered
  }
  if (rendered) {
    let container
    if (view.model.root_node) {
      container = document.querySelector(view.model.root_node)
      if (container == null) {
        container = document.createElement("div")
        container.id = view.model.root_node.replace("#", "")
        document.body.append(container)
      }
    } else if (view.container == null) {
      view._resolve_mounted()
      return null
    } else {
      container = view.container
    }
    const root = createRoot(container)
    try {
      root.render(rendered)
    } catch(e) {
      view._resolve_mounted()
      view.render_error(e)
    }
    return root
  }
}

export default {render}`
  }

  override compile(): string | null {
    const compiled = super.compile()
    if (this.bundle) {
      return compiled
    } else if (compiled === null || !compiled.includes("React")) {
      return compiled
    }
    return `
import * as React from "react"

${compiled}`
  }

  protected override get _render_cache_key() {
    const cache_key = (this.bundle === "url") ? this.esm : (this.bundle || `${this.class_name}-${this.esm.length}`)
    return `react-${this.usesMui}-${cache_key}`
  }

  static override __module__ = "panel.models.esm"

  static {
    this.prototype.default_view = ReactComponentView
    this.define<ReactComponent.Props>(({Bool, Nullable, Str}) => ({
      root_node:  [ Nullable(Str), null ],
      use_shadow_dom:   [ Bool,    true ],
    }))

    this.override<ReactComponent.Props>({
      render_policy: "manual",
    })
  }
}