// ✅ On-device PII scrubber using BERT NER via Transformers.js
// file: apps/web/src/lib/ai/PiiScrubber.ts

// Package moved from @xenova/transformers to @huggingface/transformers (v3+)
// Model hub: Xenova/ namespace still valid for existing ONNX exports
import { pipeline, env } from '@huggingface/transformers'

// Use WebGPU when available, fallback to WASM
env.backends.onnx.wasm.numThreads = 4

type NerPipeline = Awaited<ReturnType<typeof pipeline>>

let nerPipeline: NerPipeline | null = null

async function getNerPipeline(): Promise<NerPipeline> {
  if (!nerPipeline) {
    // Detect WebGPU support — Safari <18, Firefox without flags, and older Android
    // WebViews do not support WebGPU. Transformers.js falls back to WASM silently.
    // WASM inference is 10-30× slower — prefer to avoid for real-time PII scrubbing.
    const supportsWebGpu = typeof navigator !== 'undefined' && !!navigator.gpu
    const device = supportsWebGpu ? 'webgpu' : 'wasm'

    nerPipeline = await pipeline(
      'token-classification',
      'Xenova/bert-base-NER',
      { device, dtype: 'q8' },  // INT8 quantised; 'q8' works on both webgpu and wasm
    )
  }
  return nerPipeline
}

export interface ScrubResult {
  scrubbed: string
  redactedEntities: { type: string; original: string; placeholder: string }[]
  sensitivityScore: number  // 0.0 – 1.0
}

const SENSITIVE_ENTITY_TYPES = new Set([
  'PER',   // Person names
  'ORG',   // Organisations (for internal-only docs)
  'MISC',  // Miscellaneous named entities
])

export async function scrubPii(text: string): Promise<ScrubResult> {
  const ner = await getNerPipeline()
  const entities = await ner(text, { ignore_labels: [] })

  let scrubbed = text
  const redactedEntities: ScrubResult['redactedEntities'] = []
  let offset = 0

  // Sort by start position to apply redactions in order
  const sensitive = (entities as any[])
    .filter((e: any) => SENSITIVE_ENTITY_TYPES.has(e.entity_group ?? e.entity?.split('-').pop()))
    .sort((a: any, b: any) => a.start - b.start)

  for (const entity of sensitive) {
    const placeholder = `[${entity.entity_group ?? 'ENTITY'}]`
    redactedEntities.push({
      type: entity.entity_group,
      original: entity.word,
      placeholder,
    })

    // Replace in scrubbed text (offset-adjusted)
    const start = entity.start + offset
    const end = entity.end + offset
    scrubbed = scrubbed.slice(0, start) + placeholder + scrubbed.slice(end)
    offset += placeholder.length - (entity.end - entity.start)
  }

  const sensitivityScore = Math.min(1.0, sensitive.length * 0.15)

  return { scrubbed, redactedEntities, sensitivityScore }
}

// ✅ On-device embedding for sensitive documents
// file: apps/web/src/lib/ai/LocalEmbedder.ts

import { pipeline, env } from '@xenova/transformers'

let embedder: any = null

async function getEmbedder() {
  if (!embedder) {
    embedder = await pipeline(
      'feature-extraction',
      'Xenova/all-MiniLM-L6-v2',
      { device: 'webgpu', dtype: 'q8' },
    )
  }
  return embedder
}

export async function embedLocally(text: string): Promise<number[]> {
  const model = await getEmbedder()
  const output = await model(text, { pooling: 'mean', normalize: true })
  return Array.from(output.data as Float32Array)
}

// Note: dimensions are 384 (MiniLM) vs 1536 (cloud text-embedding-3-large)
// For private documents: use local 384d embeddings stored in IndexedDB
// For cross-document search in vault: use cloud 1536d embeddings (after scrubbing)

// ✅ Hybrid dispatcher — selects on-device vs cloud based on sensitivity
// file: apps/web/src/lib/ai/HybridDispatcher.ts

import { scrubPii, ScrubResult } from './PiiScrubber'
import { fetchSSE } from '../sseClient'

export type ExecutionMode = 'local' | 'cloud' | 'cloud_scrubbed'

export interface DispatchConfig {
  piiThreshold: number        // Above this → scrub before cloud (default: 0.3)
  cloudRequiredThreshold: number // Above this → always cloud, never local (default: 0.8)
  maxLocalTokens: number      // Tasks longer than this → cloud (default: 500 tokens)
}

const DEFAULT_CONFIG: DispatchConfig = {
  piiThreshold: 0.3,
  cloudRequiredThreshold: 0.8,
  maxLocalTokens: 500,
}

export class HybridDispatcher {
  constructor(private config: DispatchConfig = DEFAULT_CONFIG) {}

  async dispatch(
    userInput: string,
    taskComplexity: 'simple' | 'complex',
  ): Promise<{ mode: ExecutionMode; scrubResult?: ScrubResult; response: string }> {
    const scrubResult = await scrubPii(userInput)
    const tokenEstimate = Math.ceil(userInput.length / 4)

    // Determine routing
    let mode: ExecutionMode
    let inputToSend: string

    if (
      taskComplexity === 'simple' &&
      tokenEstimate <= this.config.maxLocalTokens &&
      scrubResult.sensitivityScore < this.config.piiThreshold
    ) {
      // Low sensitivity, simple task → run locally
      mode = 'local'
      inputToSend = userInput
    } else if (scrubResult.sensitivityScore >= this.config.piiThreshold) {
      // High sensitivity → scrub, then send to cloud
      mode = 'cloud_scrubbed'
      inputToSend = scrubResult.scrubbed
    } else {
      // Complex task, low sensitivity → cloud
      mode = 'cloud'
      inputToSend = userInput
    }

    const response = mode === 'local'
      ? await this.runLocally(inputToSend)
      : await this.runOnCloud(inputToSend)

    return { mode, scrubResult, response }
  }

  private async runLocally(text: string): Promise<string> {
    // Qwen 2.5 0.5B on WebGPU — handles simple summarisation, classification
    const model = await pipeline('text-generation', 'Xenova/Qwen2.5-0.5B-Instruct', {
      device: 'webgpu', dtype: 'q4',
    })
    const result = await model(text, { max_new_tokens: 200 })
    return (result as any)[0].generated_text
  }

  private async runOnCloud(text: string): Promise<string> {
    // Route to gateway SSE — text has already been scrubbed if needed
    const chunks: string[] = []
    for await (const event of fetchSSE('/api/chat/stream', { message: text, history: [] }, new AbortController().signal)) {
      const action = JSON.parse(event.data)
      if (action.type === 'text_reply') chunks.push(action.payload.content)
    }
    return chunks.join('')
  }
}

// ✅ Model loading progress with React Suspense-compatible state
// file: apps/web/src/hooks/useLocalModel.ts

import { useState, useEffect } from 'react'
import { env } from '@huggingface/transformers'

export interface ModelLoadState {
  status: 'idle' | 'downloading' | 'loading' | 'ready' | 'error'
  progress: number  // 0–100
  error?: string
}

export function useLocalModel(modelId: string): ModelLoadState {
  const [state, setState] = useState<ModelLoadState>({ status: 'idle', progress: 0 })

  useEffect(() => {
    setState({ status: 'downloading', progress: 0 })

    env.backends.onnx.wasm.numThreads = 4

    pipeline('feature-extraction', modelId, {
      device: 'webgpu',
      dtype: 'q8',
      progress_callback: (progress: { status: string; progress?: number }) => {
        if (progress.status === 'downloading') {
          setState({ status: 'downloading', progress: progress.progress ?? 0 })
        } else if (progress.status === 'loading') {
          setState({ status: 'loading', progress: 100 })
        }
      },
    }).then(() => {
      setState({ status: 'ready', progress: 100 })
    }).catch(error => {
      setState({ status: 'error', progress: 0, error: String(error) })
    })
  }, [modelId])

  return state
}