import React          from "https://esm.sh/react@19/";
import { createRoot } from "https://esm.sh/react-dom@19/client/";
import { pipeline, cos_sim }   from 'https://cdn.jsdelivr.net/npm/@huggingface/transformers';

import {initState,reducer} from './reducer.js';
import rawData from './rawData.js';
import BusySignal from './busy.js';

var corpus;

function App () {
  const [state, dispatch] = React.useReducer(reducer, initState());
  const setData = (key,data) => dispatch({type:'setData', key:key, data:data});
  async function text2corpus() {
    dispatch({type:'vectoring'}); // hide button; show busy signal
    let extractor = await pipeline('feature-extraction');
    corpus = await extractor(state.sentences, {
      pooling: 'mean', // Average the results to get a single vector
      normalize: true, // Normalize the vector (useful for cosine similarity)
    });
    dispatch({type:'showCorpusStats'}); // hide busy signal; show stats
  }

  async function ask() {
    // convert query to vector
    // find 5 most similar vectors
    let extractor = await pipeline('feature-extraction');
    let queryVector = await extractor(state.query, { pooling: 'mean', normalize: true });

    let results = state.sentences.map((text, indx) => {
      let score = cos_sim(queryVector.data, corpus[indx].data);
      return { text, score };
    });

    // sort results by score and grab first 5
    setData('top5', results
	    .sort((a, b) => b.score - a.score)
            .slice(0, 5));

    dispatch({type:'thinking'});
  }

  async function getAnswer(prompt) {
    let qaGen = await pipeline('text-generation', 'HuggingFaceTB/SmolLM2-135M-Instruct', {device:'webgpu'});
    let qaOutput = await qaGen(prompt, { max_new_tokens: 300 });
    qaOutput = qaOutput[0].generated_text.replace(/^.*ANSWER:/s, ''); 
    qaOutput = qaOutput.replace(/Answer:.*/s, ''); 
    //dispatch({type:'answerQ', prompt:prompt, answer:qaOutput[0].generated_text});
    dispatch({type:'answerQ', prompt:prompt, answer:qaOutput});

  }

  
  React.useEffect(()=>dispatch({type:'updateKB', data:rawData.chatbot}), []);

  React.useEffect(()=>{
    if (state.thinking) {
      let context= state.top5.map(x=>x.text).join('');
      //console.log(context);
      let prompt = `
You are a helpful assistant.

Use only the following context.

Context:
${context}

Question: ${state.query}

ANSWER:
`;
      getAnswer(prompt);
      setData('prompt', prompt);
    }
  }, [state.thinking]);
  
  return (
    <>

      {(state.sentences.length>0) &&
       <>
	 {state.error && <h3 className='error'>ERROR</h3>}

	 <p>
	   Below is the text for <a href="https://en.wikipedia.org/wiki/Chatbot">Wikipedia's chatbot article</a>.
	   Click on the button below to convert the text to vector embeddings.
	 </p>

	 <div id='KBtextDIV'>{state.KBtext}</div>

	 <div>
	   {state.wordCount} words; {state.sentences.length} sentences
	 </div>

	 {state.showConvertButton && <p><button onClick={text2corpus}>Convert to vector embeddings.</button></p>}

	 {state.busyConverting && <p id='busyDIV'>Converting... <BusySignal /></p>}

	 {state.showCorpustats &&
	  <>
	    <p>{corpus.data.length} vectors</p>
	    
	    <p className='flexColumn'>Ask a question.
	      <div className='flexColumn'>
		<textarea value={state.query}
			  onChange={e=>setData('query', e.target.value)} />
		{state.showAskButton && <div><button onClick={ask}>Ask</button></div>}
	      </div>
	    </p>
          </>
	 }

	 {(state.top5.length>0) &&
	  <>
	    <table>
	      <caption>Top 5 related sentences to your query</caption>
	      <thead>
		<tr><th>Score</th><th>Text</th></tr>
	      </thead>
	      <tbody>
		{state.top5.map((x,indx)=>
		  <tr key={'tr'+indx}><td>{x.score.toFixed(4)}</td><td>{x.text}</td></tr>
		)}
	      </tbody>
	    </table>
	    <p><b>Prompting the text generator with the following:</b> <i>{state.prompt}</i></p>

	 {state.thinking && <p id='busyDIV'>Thinking... <BusySignal /></p>}

	    {state.answers.map((x,indx)=>
	      <>
		<hr />
		<p key={'answer'+indx}><b>{state.queries[indx]}</b> - {x}</p>
	      </>
	    )}
	    
	  </>
	  
	 }
	 
       </>
      }
    </>
  )
}

createRoot(document.getElementById('root')).render( <React.StrictMode> <App/> </React.StrictMode> );