import json
import os
import re

from constants import OutputFile
from logging_utils import setup_logger
from utils import calculate_mass

logger = setup_logger(__name__, os.path.splitext(os.path.basename(__file__))[0])


def parse_formula_to_dict(formula: str) -> dict[str, int]:
    """Parse a chemical formula string to element composition dict"""
    element_dict = {}
    if formula and formula.strip():
        # Match element symbol followed by optional number or negative number
        pattern = r"([A-Z][a-z]?)(-?\d*)"
        matches = re.findall(pattern, formula)
        for elem, count_str in matches:
            if elem:
                count = int(count_str) if count_str and count_str != "-" else (1 if not count_str else -1)
                element_dict[elem] = count
    return element_dict


def gen_deltas(output_file: str = OutputFile.NEUTRAL_DELTAS) -> None:
    """Generate neutral delta data file from JSON"""

    logger.info("\n" + "=" * 60)
    logger.info("GENERATING NEUTRAL DELTA DATA")
    logger.info("=" * 60)

    logger.info("  📖 Reading from: ./data/neutral_losses.json")
    with open("./data/neutral_losses.json") as f:
        data = json.load(f)

    logger.info(f"  ✓ Parsed {len(data)} neutral deltas")
    logger.info(f"\n  📝 Writing to: {output_file}")

    # Generate NeutralDelta enum entries
    enum_entries: list[str] = []
    literal_entries: list[str] = []

    for delta in data:
        # Create enum name from the name field (e.g., "Ammonia" -> "AMMONIA")
        enum_name = delta["name"].upper().replace(" ", "_").replace("-", "_")
        formula = delta["formula"]

        enum_entries.append(f'    {enum_name} = "{formula}"')
        literal_entries.append(f'    "{formula}",')

    enum_str = "\n".join(enum_entries)
    literal_str = "\n".join(literal_entries)

    # Generate the neutral delta dictionary entries
    entries: list[str] = []
    for delta in data:
        enum_name = delta["name"].upper().replace(" ", "_").replace("-", "_")
        formula = delta["formula"]
        name = delta["name"]
        description = delta["description"]
        amino_acids = delta.get("amino_acids", [])

        # Parse formula to get element counts as a dict
        element_dict = parse_formula_to_dict(formula)

        # make all values negative for neutral losses
        for elem in element_dict:
            element_dict[elem] = -abs(element_dict[elem])

        # Calculate masses
        monoisotopic_mass = calculate_mass(element_dict, monoisotopic=True)
        average_mass = calculate_mass(element_dict, monoisotopic=False)

        # Format amino acids as frozenset
        aa_set = f"frozenset({amino_acids})" if amino_acids else "frozenset()"

        entry = f'''NeutralDelta.{enum_name}: NeutralDeltaInfo(
    formula="{formula}",
    name="{name}",
    description="{description}",
    amino_acids={aa_set},
    monoisotopic_mass={monoisotopic_mass},
    average_mass={average_mass},
    dict_composition={element_dict},
),'''
        entries.append(entry)

    entries_str = "\n".join(entries)

    # Write the complete file
    content = f'''"""Auto-generated neutral delta data"""
# DO NOT EDIT - generated by gen_neutral_deltas.py

from enum import StrEnum
from typing import Literal
from .dclass import NeutralDeltaInfo


class NeutralDelta(StrEnum):
    """Neutral loss deltas for fragment ions"""
{enum_str}


NeutralDeltaLiteral = Literal[
{literal_str}
]

NEUTRAL_DELTA_DICT: dict[NeutralDelta, NeutralDeltaInfo] = {{
{entries_str}
}}
'''

    with open(output_file, "w") as f:
        f.write(content)

    logger.info(f"✅ Successfully generated {output_file}")
    logger.info(f"   Total entries: {len(data)}")


if __name__ == "__main__":
    gen_deltas()
