import { useEffect, useMemo, useRef, useState } from 'react'

type Message = { role: 'user' | 'assistant' | 'system'; content: string }

type OIDCConfig = {
  issuer: string
  authorization_endpoint: string
  token_endpoint: string
}

const OIDC_CLIENT_ID = import.meta.env.VITE_OIDC_CLIENT_ID
const OIDC_REDIRECT_URI = import.meta.env.VITE_OIDC_REDIRECT_URI || window.location.origin
const OIDC_SCOPE = 'openid profile email'
const OIDC_RESPONSE_TYPE = 'code'
const OIDC_CODE_CHALLENGE_METHOD = 'S256'
const API_BASE = import.meta.env.VITE_API_BASE || 'http://localhost:8000'

function base64UrlEncode(buffer: ArrayBuffer) {
  const bytes = new Uint8Array(buffer)
  let str = ''
  for (const byte of bytes) {
    str += String.fromCharCode(byte)
  }
  return btoa(str)
    .replace(/\+/g, '-')
    .replace(/\//g, '_')
    .replace(/=+$/, '')
}

async function sha256(text: string) {
  // Prefer SubtleCrypto when available (secure contexts). Fallback to a
  // pure-JS implementation for environments where `crypto.subtle` is undefined.
  const data = new TextEncoder().encode(text)
  if (typeof crypto !== 'undefined' && crypto.subtle && typeof crypto.subtle.digest === 'function') {
    const hash = await crypto.subtle.digest('SHA-256', data)
    return base64UrlEncode(hash)
  }

  // JS fallback SHA-256 (returns ArrayBuffer)
  function sha256Fallback(message: Uint8Array) {
    const H = new Uint32Array([0x6a09e667, 0xbb67ae85, 0x3c6ef372, 0xa54ff53a, 0x510e527f, 0x9b05688c, 0x1f83d9ab, 0x5be0cd19])
    const K = new Uint32Array([
      0x428a2f98,0x71374491,0xb5c0fbcf,0xe9b5dba5,0x3956c25b,0x59f111f1,0x923f82a4,0xab1c5ed5,
      0xd807aa98,0x12835b01,0x243185be,0x550c7dc3,0x72be5d74,0x80deb1fe,0x9bdc06a7,0xc19bf174,
      0xe49b69c1,0xefbe4786,0x0fc19dc6,0x240ca1cc,0x2de92c6f,0x4a7484aa,0x5cb0a9dc,0x76f988da,
      0x983e5152,0xa831c66d,0xb00327c8,0xbf597fc7,0xc6e00bf3,0xd5a79147,0x06ca6351,0x14292967,
      0x27b70a85,0x2e1b2138,0x4d2c6dfc,0x53380d13,0x650a7354,0x766a0abb,0x81c2c92e,0x92722c85,
      0xa2bfe8a1,0xa81a664b,0xc24b8b70,0xc76c51a3,0xd192e819,0xd6990624,0xf40e3585,0x106aa070,
      0x19a4c116,0x1e376c08,0x2748774c,0x34b0bcb5,0x391c0cb3,0x4ed8aa4a,0x5b9cca4f,0x682e6ff3,
      0x748f82ee,0x78a5636f,0x84c87814,0x8cc70208,0x90befffa,0xa4506ceb,0xbef9a3f7,0xc67178f2
    ])

    const ml = message.length * 8
    const withOne = new Uint8Array(((message.length + 9 + 63) >> 6) << 6)
    withOne.set(message)
    withOne[message.length] = 0x80
    const view = new DataView(withOne.buffer)
    view.setUint32(withOne.length - 4, ml >>> 0)
    view.setUint32(withOne.length - 8, Math.floor(ml / 0x100000000))

    const w = new Uint32Array(64)
    for (let i = 0; i < withOne.length; i += 64) {
      for (let t = 0; t < 16; ++t) w[t] = view.getUint32(i + t * 4)
      for (let t = 16; t < 64; ++t) {
        const s0 = (rightRotate(w[t - 15], 7) ^ rightRotate(w[t - 15], 18) ^ (w[t - 15] >>> 3)) >>> 0
        const s1 = (rightRotate(w[t - 2], 17) ^ rightRotate(w[t - 2], 19) ^ (w[t - 2] >>> 10)) >>> 0
        w[t] = (w[t - 16] + s0 + w[t - 7] + s1) >>> 0
      }

      let a = H[0], b = H[1], c = H[2], d = H[3], e = H[4], f = H[5], g = H[6], h = H[7]
      for (let t = 0; t < 64; ++t) {
        const S1 = (rightRotate(e, 6) ^ rightRotate(e, 11) ^ rightRotate(e, 25)) >>> 0
        const ch = ((e & f) ^ (~e & g)) >>> 0
        const temp1 = (h + S1 + ch + K[t] + w[t]) >>> 0
        const S0 = (rightRotate(a, 2) ^ rightRotate(a, 13) ^ rightRotate(a, 22)) >>> 0
        const maj = ((a & b) ^ (a & c) ^ (b & c)) >>> 0
        const temp2 = (S0 + maj) >>> 0

        h = g; g = f; f = e; e = (d + temp1) >>> 0
        d = c; c = b; b = a; a = (temp1 + temp2) >>> 0
      }

      H[0] = (H[0] + a) >>> 0
      H[1] = (H[1] + b) >>> 0
      H[2] = (H[2] + c) >>> 0
      H[3] = (H[3] + d) >>> 0
      H[4] = (H[4] + e) >>> 0
      H[5] = (H[5] + f) >>> 0
      H[6] = (H[6] + g) >>> 0
      H[7] = (H[7] + h) >>> 0
    }

    const out = new Uint8Array(32)
    const outView = new DataView(out.buffer)
    for (let i = 0; i < 8; ++i) outView.setUint32(i * 4, H[i])
    return out.buffer
  }

  function rightRotate(value: number, amount: number) {
    return (value >>> amount) | (value << (32 - amount))
  }

  const hashBuf = sha256Fallback(data)
  return base64UrlEncode(hashBuf)
}

async function createCodeVerifier() {
  const array = new Uint8Array(32)
  window.crypto.getRandomValues(array)
  return base64UrlEncode(array)
}

function generateUUID() {
  // Use native crypto.randomUUID when available, otherwise fallback to RFC4122 v4
  // Compatible with older browsers that lack crypto.randomUUID
  if (typeof crypto !== 'undefined' && 'randomUUID' in crypto) {
    // @ts-ignore
    return (crypto as any).randomUUID()
  }

  const bytes = crypto.getRandomValues(new Uint8Array(16))
  bytes[6] = (bytes[6] & 0x0f) | 0x40
  bytes[8] = (bytes[8] & 0x3f) | 0x80
  const hex = Array.from(bytes, (b) => b.toString(16).padStart(2, '0')).join('')
  return `${hex.substr(0, 8)}-${hex.substr(8, 4)}-${hex.substr(12, 4)}-${hex.substr(16, 4)}-${hex.substr(20, 12)}`
}

function getLocalToken() {
  return window.localStorage.getItem('oidc_access_token')
}

function setLocalToken(token: string) {
  window.localStorage.setItem('oidc_access_token', token)
}

function clearLocalToken() {
  window.localStorage.removeItem('oidc_access_token')
}

async function fetchOIDCConfig(): Promise<OIDCConfig> {
  const response = await fetch(`${API_BASE}/auth/config`)
  if (!response.ok) {
    throw new Error('Unable to fetch OIDC configuration')
  }
  return response.json()
}

async function handleRedirectCallback(): Promise<string | null> {
  const params = new URLSearchParams(window.location.search)
  const code = params.get('code')
  const state = params.get('state')
  const storedState = window.sessionStorage.getItem('oidc_state')
  const codeVerifier = window.sessionStorage.getItem('oidc_code_verifier')

  if (!code || !state || state !== storedState || !codeVerifier) {
    return null
  }

  window.history.replaceState({}, document.title, window.location.pathname)

  const response = await fetch(`${API_BASE}/auth/callback`, {
    method: 'POST',
    headers: { 'Content-Type': 'application/json' },
    body: JSON.stringify({ code, code_verifier: codeVerifier, redirect_uri: OIDC_REDIRECT_URI }),
  })

  const result = await response.json()
  if (response.ok && result.access_token) {
    setLocalToken(result.access_token)
    return result.access_token
  }

  throw new Error(result.error_description || 'OIDC callback failed')
}

function parseSSE(text: string, onData: (chunk: string) => void) {
  const events = text.split('\n\n')
  for (const event of events) {
    const trimmed = event.trim()
    if (!trimmed) continue
    const lines = trimmed.split('\n')
    const dataLines = lines
      .filter((line) => line.startsWith('data:'))
      .map((line) => line.replace(/^data:\s?/, ''))
    if (dataLines.length) {
      onData(dataLines.join('\n'))
    }
  }
}

function App() {
  const [messages, setMessages] = useState<Message[]>([])
  const [input, setInput] = useState('')
  const [token, setToken] = useState<string | null>(getLocalToken())
  const [loading, setLoading] = useState(false)
  const [error, setError] = useState<string | null>(null)
  const controllerRef = useRef<AbortController | null>(null)

  const isAuthenticated = !!token

  useEffect(() => {
    handleRedirectCallback()
      .then((newToken) => {
        if (newToken) {
          setToken(newToken)
        }
      })
      .catch((err) => {
        console.error(err)
        setError(err instanceof Error ? err.message : 'Authentication failed')
      })

    return () => {
      controllerRef.current?.abort()
    }
  }, [])

  const authHeaders = useMemo(
    () => ({ Authorization: token ? `Bearer ${token}` : '' }),
    [token],
  )

  async function login() {
    try {
      const config = await fetchOIDCConfig()
      const state = generateUUID()
      const codeVerifier = await createCodeVerifier()
      const codeChallenge = await sha256(codeVerifier)

      window.sessionStorage.setItem('oidc_state', state)
      window.sessionStorage.setItem('oidc_code_verifier', codeVerifier)

      const params = new URLSearchParams({
        response_type: OIDC_RESPONSE_TYPE,
        client_id: OIDC_CLIENT_ID,
        redirect_uri: OIDC_REDIRECT_URI,
        scope: OIDC_SCOPE,
        state,
        code_challenge: codeChallenge,
        code_challenge_method: OIDC_CODE_CHALLENGE_METHOD,
      })

      window.location.href = `${config.authorization_endpoint}?${params.toString()}`
    } catch (err) {
      setError(err instanceof Error ? err.message : 'Login failed')
    }
  }

  function logout() {
    clearLocalToken()
    setToken(null)
    setMessages([])
  }

  async function submitMessage() {
    if (!input.trim() || !token) return

    const userMessage: Message = { role: 'user', content: input.trim() }
    const messageBatch = [...messages, userMessage]
    setMessages(messageBatch)
    setInput('')
    setLoading(true)
    setError(null)

    const controller = new AbortController()
    controllerRef.current = controller

    try {
      const response = await fetch(`${API_BASE}/chat`, {
        method: 'POST',
        headers: {
          'Content-Type': 'application/json',
          ...authHeaders,
        },
        body: JSON.stringify({ model: 'llama2', messages: messageBatch }),
        signal: controller.signal,
      })

      if (!response.ok) {
        const payload = await response.text()
        throw new Error(payload || 'Chat request failed')
      }

      const reader = response.body?.getReader()
      if (!reader) {
        throw new Error('No streaming body available')
      }

      const decoder = new TextDecoder()
      let buffer = ''

      const appendAssistantText = (chunk: string) => {
        setMessages((prev) => {
          const last = prev[prev.length - 1]
          if (last?.role === 'assistant') {
            return [...prev.slice(0, -1), { ...last, content: last.content + chunk }]
          }
          return [...prev, { role: 'assistant', content: chunk }]
        })
      }

      while (true) {
        const { done, value } = await reader.read()
        if (done) break
        buffer += decoder.decode(value, { stream: true })
        const boundary = buffer.indexOf('\n\n')
        if (boundary === -1) continue

        const chunk = buffer.slice(0, boundary + 2)
        buffer = buffer.slice(boundary + 2)
        parseSSE(chunk, appendAssistantText)
      }
    } catch (err) {
      if (err instanceof Error) {
        setError(err.message)
      } else {
        setError('Streaming chat failed')
      }
    } finally {
      setLoading(false)
      controllerRef.current = null
    }
  }

  return (
    <div style={{ maxWidth: 800, margin: '0 auto', padding: 24, fontFamily: 'system-ui, sans-serif' }}>
      <header style={{ display: 'flex', justifyContent: 'space-between', alignItems: 'center' }}>
        <h1>AI Chat</h1>
        <div>
          {isAuthenticated ? (
            <button onClick={logout}>Logout</button>
          ) : (
            <button onClick={login}>Login with OIDC</button>
          )}
        </div>
      </header>

      <main>
        <section>
          <h2>Chat</h2>
          {!isAuthenticated && <p>Please log in to start chatting.</p>}
          {error && <div style={{ color: 'red', marginBottom: 12 }}>{error}</div>}
          <div style={{ minHeight: 300, border: '1px solid #ddd', padding: 16, borderRadius: 8, marginBottom: 16 }}>
            {messages.map((message, index) => (
              <div key={index} style={{ marginBottom: 12 }}>
                <strong>{message.role}:</strong>
                <p style={{ margin: '4px 0 0', whiteSpace: 'pre-wrap' }}>{message.content}</p>
              </div>
            ))}
          </div>
          <div style={{ display: 'flex', gap: 8, flexWrap: 'wrap' }}>
            <input
              value={input}
              onChange={(event) => setInput(event.target.value)}
              disabled={!isAuthenticated || loading}
              style={{ flex: 1, minWidth: 0, padding: '8px 12px' }}
              placeholder="Type your message..."
            />
            <button onClick={submitMessage} disabled={!isAuthenticated || loading || !input.trim()}>
              Send
            </button>
          </div>
          {loading && <p>Waiting for response...</p>}
        </section>
      </main>
    </div>
  )
}

export default App