import json
import os

from constants import PROTON_MASS, OutputFile
from logging_utils import setup_logger
from utils import calculate_mass, parse_formula_to_dict

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


def gen_refmol(output_file: str = OutputFile.REFMOL) -> None:
    """Generate reference molecule data file from JSON"""

    logger.info("\n" + "=" * 60)
    logger.info("GENERATING REFERENCE MOLECULE DATA")
    logger.info("=" * 60)

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

    logger.info(f"  ✓ Parsed {len(data)} reference molecules")

    logger.info(f"\n  📝 Writing to: {output_file}")

    # Generate enum entries and dictionary entries
    enum_entries: list[str] = []
    literal_entries: list[str] = []
    dict_entries: list[str] = []

    for name, info in sorted(data.items()):
        # Convert name to enum format (replace hyphens with underscores, handle special chars)
        enum_name = name.replace("-", "_").replace("+", "_PLUS").upper()

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

        # Extract fields
        label_type = info.get("label_type", "")
        molecule_type = info.get("molecule_type", "")
        chemical_formula = info.get("chemical_formula", "")

        # Parse formula using lightweight parser and compute masses
        composition_dict = {}
        monoisotopic_mass = 0.0
        average_mass = 0.0

        if chemical_formula:
            try:
                composition_dict = parse_formula_to_dict(chemical_formula)
                monoisotopic_mass = calculate_mass(composition_dict, monoisotopic=True)
                average_mass = calculate_mass(composition_dict, monoisotopic=False)
            except Exception as e:
                logger.warning(f"Error parsing formula for {name}: {chemical_formula}, {e}")

        # Override with provided masses if available
        if "neutral_mass" in info:
            provided_mass = info["neutral_mass"]
            if monoisotopic_mass != 0.0 and abs(provided_mass - monoisotopic_mass) > 0.01:
                symbol = "🔴" if abs(provided_mass - monoisotopic_mass) > 1.0 else "⚠️"

                logger.warning(
                    f"\n  {symbol}  REFMOL MASS MISMATCH: {name} | NEUTRAL MASS\n"
                    f"      Calculated: {monoisotopic_mass:.6f}, Provided: {provided_mass:.6f}\n"
                    f"      Difference: {abs(provided_mass - monoisotopic_mass):.6f}\n"
                    f"      Composition: {composition_dict}"
                )

        if "ion_mz" in info:
            # For ions, subtract proton mass to get neutral mass
            provided_mz = info["ion_mz"]
            neutral_mass = provided_mz - PROTON_MASS
            if monoisotopic_mass != 0.0 and abs(neutral_mass - monoisotopic_mass) > 0.01:
                symbol = "🔴" if abs(neutral_mass - monoisotopic_mass) > 1.0 else "⚠️"

                logger.warning(
                    f"\n  {symbol}  REFMOL MASS MISMATCH: {name} | ION M/Z\n"
                    f"      Calculated: {monoisotopic_mass:.6f}, Ion m/z: {neutral_mass:.6f}\n"
                    f"      Difference: {abs(neutral_mass - monoisotopic_mass):.6f}\n"
                    f"      Composition: {composition_dict}"
                )

        dict_entry = f'''RefMolID.{enum_name}: RefMolInfo(
    name="{name}",
    label_type="{label_type}",
    molecule_type="{molecule_type}",
    chemical_formula='{chemical_formula}',
    monoisotopic_mass={monoisotopic_mass},
    average_mass={average_mass},
    dict_composition={composition_dict},
),'''
        dict_entries.append(dict_entry)

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

    # Write the complete file
    content = f'''"""Auto-generated reference molecule data"""
# DO NOT EDIT - generated by gen_refmol.py

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


class RefMolID(StrEnum):
    """Enum of reference molecule IDs"""
{enum_str}


RefMolLiteral = Literal[
{literal_str}
]


REFMOL_DICT: dict[RefMolID, RefMolInfo] = {{
{dict_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_refmol()
