import * as React from 'react'
import {__, _x} from '@wordpress/i18n'
import cx from 'classnames'

import {
	useState,
} from '@wordpress/element'

import {
	Button,
	type ButtonProps,
} from '@ska/components'

import {
	AIModal,
	type AIModalProps,
	type ChatProps,
} from '../..'

import type {
	AISession,
} from '../../types'

import './style.scss'

export type AIButtonProps = Omit<ButtonProps, 'onClick'> & {

	/** AI session context. */
	session: AISession

	/** Button component to render. */
	as?: React.ElementType

	/** Click handler. */
	onClick?: (open: () => void) => void
	/** Close modal after response is received. */
	closeOnResponse?: boolean

	initialActive?: boolean

	/** Props for Modal component. */
	modalProps?: Partial<AIModalProps>
	/** Props for Chat component. */
	chatProps?: Partial<ChatProps>

	/** `@tanstack/react-query` query client provider. */
	QueryClientProvider?: React.FC<React.PropsWithChildren>
}

const DefaultQueryClientProvider: React.FC<React.PropsWithChildren> = ({children}) => children

const AIButton: React.FC<AIButtonProps> = ({
	session = {},
	as: CustomComponent,
	onClick,
	closeOnResponse = false,
	initialActive = false,
	modalProps = {},
	chatProps = {},
	QueryClientProvider = DefaultQueryClientProvider,
	isBusy = false,
	...props
}) => {

	const [prompt, setPrompt] = useState('')
	const [model, setModel] = useState('')
	const [isActive, setIsActive] = useState(initialActive)
	const [isOpen, setIsOpen] = useState(isActive)
	const [isLoading, setIsLoading] = useState(false)
	const Component = CustomComponent || Button

	return <>
		<Component
			children={__('AI', 'ska-blocks')}
			aria-pressed={isActive}
			isBusy={isLoading || isBusy}
			{...props}
			className={cx('ska-ai-button', props?.className, {
				'ska-ai-button--loading': isLoading,
			})}
			onClick={() => {
				const open = () => {
					if(!isActive) {
						setIsActive(true)
					}
					setIsOpen(true)
				}
				onClick ? onClick(open) : open()
			}}
		/>
		{isActive && <>
			<QueryClientProvider>
				<AIModal
					session={{
						...session,
						onResponse: response => {
							setPrompt('')
							if('onResponse' in session && typeof session.onResponse === 'function') {
								session.onResponse(response)
							}
							if(closeOnResponse) {
								setIsOpen(false)
							}
						},
					}}
					isOpen={isOpen}
					prompt={prompt}
					setPrompt={setPrompt}
					model={model}
					setModel={setModel}
					onRequestBegin={() => {
						setIsLoading(true)
						const {
							closeOnSend = false,
						} = modalProps
						if(closeOnSend) {
							setIsOpen(false)
						}
					}}
					onRequestFinish={() => {
						setIsLoading(false)
					}}
					onRequestAbort={() => {
						setIsLoading(false)
					}}
					{...modalProps}
					onRequestClose={() => {
						setIsOpen(false)
						modalProps?.onRequestClose && modalProps.onRequestClose()
					}}
					chatProps={{
						...modalProps?.chatProps,
						...chatProps,
						onEndSession: () => {
							setIsOpen(false)
							setIsLoading(false)
							setIsActive(false)
							chatProps?.onEndSession && chatProps.onEndSession()
						},
					}}
				/>
			</QueryClientProvider>
		</>}
	</>
}

export default AIButton