// Types
export interface ParsedElement {
  symbol: string
  count: number
}

export interface FormulaParseResult {
  elements: Record<string, number>
  valid: boolean
  error?: string
}

/** `coefficient` is a full-size multiplier at the start of a hydrate segment (the 5 in CuSO4·5H2O). */
export type FormulaPartType = 'element' | 'subscript' | 'superscript' | 'paren' | 'dot' | 'charge' | 'coefficient'

export interface FormulaPart {
  type: FormulaPartType
  text: string
}

// Every element with its standard atomic weight (mass number of the most stable isotope when none is defined)
export const ATOMIC_WEIGHTS: Record<string, number> = {
  H: 1.008, He: 4.003, Li: 6.941, Be: 9.012, B: 10.81,
  C: 12.011, N: 14.007, O: 15.999, F: 18.998, Ne: 20.180,
  Na: 22.990, Mg: 24.305, Al: 26.982, Si: 28.086, P: 30.974,
  S: 32.065, Cl: 35.453, Ar: 39.948, K: 39.098, Ca: 40.078,
  Sc: 44.956, Ti: 47.867, V: 50.942, Cr: 51.996,
  Mn: 54.938, Fe: 55.845, Co: 58.933, Ni: 58.693, Cu: 63.546,
  Zn: 65.38, Ga: 69.723, Ge: 72.630, As: 74.922, Se: 78.971,
  Br: 79.904, Kr: 83.798, Rb: 85.468, Sr: 87.62, Y: 88.906,
  Zr: 91.224, Nb: 92.906, Mo: 95.95, Tc: 98, Ru: 101.07,
  Rh: 102.906, Pd: 106.42, Ag: 107.868, Cd: 112.414, In: 114.818,
  Sn: 118.710, Sb: 121.760, Te: 127.60, I: 126.904, Xe: 131.293,
  Cs: 132.905, Ba: 137.327, La: 138.905, Ce: 140.116, Pr: 140.908,
  Nd: 144.242, Pm: 145, Sm: 150.36, Eu: 151.964, Gd: 157.25,
  Tb: 158.925, Dy: 162.500, Ho: 164.930, Er: 167.259, Tm: 168.934,
  Yb: 173.045, Lu: 174.967, Hf: 178.49, Ta: 180.948, W: 183.84,
  Re: 186.207, Os: 190.23, Ir: 192.217, Pt: 195.084, Au: 196.967,
  Hg: 200.592, Tl: 204.38, Pb: 207.2, Bi: 208.980, Po: 209,
  At: 210, Rn: 222, Fr: 223, Ra: 226, Ac: 227,
  Th: 232.038, Pa: 231.036, U: 238.029, Np: 237, Pu: 244,
  Am: 243, Cm: 247, Bk: 247, Cf: 251, Es: 252,
  Fm: 257, Md: 258, No: 259, Lr: 266, Rf: 267,
  Db: 268, Sg: 269, Bh: 270, Hs: 277, Mt: 278,
  Ds: 281, Rg: 282, Cn: 285, Nh: 286, Fl: 289,
  Mc: 290, Lv: 293, Ts: 294, Og: 294,
}

function parseFormulaSegment(formula: string): FormulaParseResult {
  const elements: Record<string, number> = {}
  const stack: Record<string, number>[] = [{}]
  let i = 0

  while (i < formula.length) {
    const ch = formula[i]

    if (ch === '(') {
      stack.push({})
      i++
    } else if (ch === ')') {
      if (stack.length < 2) {
        return { elements: {}, valid: false, error: 'Unmatched closing parenthesis' }
      }
      i++
      // Read multiplier after closing paren
      let numStr = ''
      while (i < formula.length && /\d/.test(formula[i])) {
        numStr += formula[i]
        i++
      }
      const multiplier = numStr ? parseInt(numStr, 10) : 1
      const top = stack.pop()!
      const current = stack[stack.length - 1]
      for (const [el, count] of Object.entries(top)) {
        current[el] = (current[el] || 0) + count * multiplier
      }
    } else if (/[A-Z]/.test(ch)) {
      // Element symbol: uppercase letter optionally followed by lowercase
      let symbol = ch
      i++
      while (i < formula.length && /[a-z]/.test(formula[i])) {
        symbol += formula[i]
        i++
      }
      // Validate element
      if (!(symbol in ATOMIC_WEIGHTS)) {
        return { elements: {}, valid: false, error: `Unknown element: ${symbol}` }
      }
      // Read count
      let numStr = ''
      while (i < formula.length && /\d/.test(formula[i])) {
        numStr += formula[i]
        i++
      }
      const count = numStr ? parseInt(numStr, 10) : 1
      const current = stack[stack.length - 1]
      current[symbol] = (current[symbol] || 0) + count
    } else {
      return { elements: {}, valid: false, error: `Unexpected character: ${ch}` }
    }
  }

  if (stack.length !== 1) {
    return { elements: {}, valid: false, error: 'Unmatched opening parenthesis' }
  }

  Object.assign(elements, stack[0])
  return { elements, valid: true }
}

/** Parses and renders chemical formulas, computes molecular weight, and generates rich-text formula parts. */
export function useChemicalFormula() {
  function parseFormula(formula: string): FormulaParseResult {
    const trimmed = formula.trim()
    if (!trimmed) {
      return { elements: {}, valid: false, error: 'Empty formula' }
    }

    // A trailing charge (SO4^2-, Fe2+, NH4+) does not change the composition. Like renderFormulaParts,
    // digits before a bare sign are the charge only on a lone element (Fe2+); otherwise they are a
    // subscript and the sign a unit charge (NH4+). Write Hg2^2+ for any other multiple charge.
    const body = trimmed.replace(/^([A-Z][a-z]?)\d*[+-]$/, '$1').replace(/(\^\d*)?[+-]$/, '')
    if (!body) {
      return { elements: {}, valid: false, error: 'Empty formula' }
    }
    // Split on hydrate separator (·, . or *)
    const parts = body.split(/[·.*]/)
    const combined: Record<string, number> = {}

    for (const part of parts) {
      const segment = part.trim()
      if (!segment) continue

      // Check for leading coefficient (e.g., 5H2O)
      let coefficient = 1
      let formulaStr = segment
      const coeffMatch = segment.match(/^(\d+)([A-Z].*)$/)
      if (coeffMatch) {
        coefficient = parseInt(coeffMatch[1], 10)
        formulaStr = coeffMatch[2]
      }

      const result = parseFormulaSegment(formulaStr)
      if (!result.valid) {
        return result
      }

      for (const [el, count] of Object.entries(result.elements)) {
        combined[el] = (combined[el] || 0) + count * coefficient
      }
    }

    return { elements: combined, valid: true }
  }

  function calculateMW(elements: Record<string, number>): number {
    let mw = 0
    for (const [el, count] of Object.entries(elements)) {
      const weight = ATOMIC_WEIGHTS[el]
      if (weight !== undefined) {
        mw += weight * count
      }
    }
    return Math.round(mw * 1000) / 1000
  }

  function renderFormulaParts(formula: string): FormulaPart[] {
    const parts: FormulaPart[] = []
    const trimmed = formula.trim()
    let i = 0

    while (i < trimmed.length) {
      const ch = trimmed[i]

      // Hydrate dot separator
      if (ch === '·' || (ch === '.' && i > 0 && /\d/.test(trimmed[i - 1]) && i + 1 < trimmed.length && /\d/.test(trimmed[i + 1]))) {
        // Disambiguate: '.' as hydrate separator only if followed by a digit then uppercase
        // For simplicity, treat '·' always as dot
        if (ch === '·') {
          parts.push({ type: 'dot', text: '·' })
          i++
          continue
        }
      }

      // Parentheses
      if (ch === '(' || ch === ')') {
        parts.push({ type: 'paren', text: ch })
        i++
        // After ')' check for subscript number
        if (ch === ')' && i < trimmed.length && /\d/.test(trimmed[i])) {
          let numStr = ''
          while (i < trimmed.length && /\d/.test(trimmed[i])) {
            numStr += trimmed[i]
            i++
          }
          parts.push({ type: 'subscript', text: numStr })
        }
        continue
      }

      // Charge notation: ^2+, ^2-, 2+, 2- at end, or +, -
      if (ch === '^') {
        i++
        let chargeStr = ''
        while (i < trimmed.length && /[\d+\-]/.test(trimmed[i])) {
          chargeStr += trimmed[i]
          i++
        }
        if (chargeStr) {
          parts.push({ type: 'charge', text: chargeStr })
        }
        continue
      }

      // Element symbol: uppercase followed by optional lowercase
      if (/[A-Z]/.test(ch)) {
        let symbol = ch
        i++
        while (i < trimmed.length && /[a-z]/.test(trimmed[i])) {
          symbol += trimmed[i]
          i++
        }
        parts.push({ type: 'element', text: symbol })

        // Check for subscript number after element
        if (i < trimmed.length && /\d/.test(trimmed[i])) {
          let numStr = ''
          while (i < trimmed.length && /\d/.test(trimmed[i])) {
            numStr += trimmed[i]
            i++
          }
          parts.push({ type: 'subscript', text: numStr })

          // On a lone element the digits before a trailing sign are the charge (Fe2+);
          // elsewhere they stay a subscript and the sign renders as a unit charge (NH4+).
          if (parts.length === 2 && i === trimmed.length - 1 && (trimmed[i] === '+' || trimmed[i] === '-')) {
            parts.push({ type: 'charge', text: numStr + trimmed[i] })
            parts.splice(parts.length - 2, 1) // Remove the subscript
            i++
            continue
          }
        }
        continue
      }

      // Leading digit (hydrate coefficient like 5H2O)
      if (/\d/.test(ch)) {
        let numStr = ''
        while (i < trimmed.length && /\d/.test(trimmed[i])) {
          numStr += trimmed[i]
          i++
        }
        // Could be a subscript or coefficient depending on context
        // If next char is uppercase, it's a coefficient (render as element-like text)
        // If at end or followed by +/-, it's a charge component
        const prev = parts[parts.length - 1]
        if (i < trimmed.length && (trimmed[i] === '+' || trimmed[i] === '-')) {
          parts.push({ type: 'charge', text: numStr + trimmed[i] })
          i++
        } else if (!prev || prev.type === 'dot') {
          parts.push({ type: 'coefficient', text: numStr })
        } else {
          parts.push({ type: 'subscript', text: numStr })
        }
        continue
      }

      // Standalone + or - (charge)
      if (ch === '+' || ch === '-') {
        parts.push({ type: 'charge', text: ch })
        i++
        continue
      }

      // Dot or asterisk as hydrate separator
      if (ch === '.' || ch === '*') {
        parts.push({ type: 'dot', text: '·' })
        i++
        continue
      }

      // Skip whitespace
      if (/\s/.test(ch)) {
        i++
        continue
      }

      // Unknown character, skip
      i++
    }

    return parts
  }

  return {
    parseFormula,
    calculateMW,
    renderFormulaParts,
  }
}