import json
import os

from constants import OutputFile
from logging_utils import setup_logger

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


# Map fragment ion IDs to IonType enum values when they differ
ID_TO_ION_TYPE = {
    "z_radical": "z.",
    "z_plus_h": "z+H",
    "c_minus_h": "c-H",
    "d_valine": "d-valine",
    "da_threonine": "da-threonine",
    "da_isoleucine": "da-isoleucine",
    "db_threonine": "db-threonine",
    "db_isoleucine": "db-isoleucine",
    "w_valine": "w-valine",
    "wa_threonine": "wa-threonine",
    "wa_isoleucine": "wa-isoleucine",
    "wb_threonine": "wb-threonine",
    "wb_isoleucine": "wb-isoleucine",
}


def gen_frag(output_file: str = OutputFile.FRAGMENT_IONS) -> None:
    """Generate fragment ion data file from JSON"""

    logger.info("\n" + "=" * 60)
    logger.info("GENERATING FRAGMENT ION DATA")
    logger.info("=" * 60)

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

    fragment_ions = data["fragment_ions"]
    logger.info(f"  ✓ Parsed {len(fragment_ions)} fragment ions")

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

    # Generate IonType enum entries
    iontype_entries: list[str] = []
    iontype_literals: list[str] = []

    for ion in fragment_ions:
        ion_id = ion["id"]
        capitalized_id = ion["capitalized_id"]
        # Get the IonType value (either from mapping or use the id directly)
        ion_type_value = ID_TO_ION_TYPE.get(ion_id, ion_id)

        iontype_entries.append(f'    {capitalized_id} = "{ion_type_value}"')
        iontype_literals.append(f'    "{ion_type_value}",')

    iontype_enum_str = "\n".join(iontype_entries)
    iontype_literal_str = "\n".join(iontype_literals)

    # Generate the fragment ion dictionary entries
    entries: list[str] = []
    for ion in fragment_ions:
        ion_id = ion["id"]
        capitalized_id = ion["capitalized_id"]

        # Build IonTypeProperty flags based on the boolean fields
        properties_flags = []

        if ion.get("is_forward", False):
            properties_flags.append("IonTypeProperty.FORWARD")
        if ion.get("is_backward", False):
            properties_flags.append("IonTypeProperty.BACKWARD")
        if ion.get("is_internal", False):
            properties_flags.append("IonTypeProperty.INTERNAL")
        if ion.get("is_intact", False):
            properties_flags.append("IonTypeProperty.INTACT")
        if ion.get("is_aa_specific_fwd", False):
            properties_flags.append("IonTypeProperty.AA_SPECIFIC_FWD")
        if ion.get("is_aa_specific_bwd", False):
            properties_flags.append("IonTypeProperty.AA_SPECIFIC_BWD")

        # Combine flags with bitwise OR
        properties_str = " | ".join(properties_flags) if properties_flags else "IonTypeProperty.NONE"

        # Parse formula to get element counts as a dict
        element_dict = {}
        formula = ion.get("chemical_formula", "")
        if formula and formula.strip():
            # Simple formula parser for element composition
            import re

            # 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
            formula = f"'{str(formula)}'"
        elif formula is None:
            formula = None

        else:
            formula = "''"

        entry = f'''IonType.{capitalized_id}: FragmentIonInfo(
    id=IonType.{capitalized_id},
    name="{ion["name"]}",
    formula={formula},
    monoisotopic_mass={ion["monoisotopic_mass"]},
    average_mass={ion["average_mass"]},
    dict_composition={element_dict},
    properties={properties_str},
),'''
        entries.append(entry)

    entries_str = "\n".join(entries)

    # Write the complete file
    content = f'''"""Auto-generated fragment ion data"""
# DO NOT EDIT - generated by gen_fragment_ions.py

from enum import StrEnum
from typing import Literal
from .dclass import FragmentIonInfo, IonTypeProperty


class IonType(StrEnum):
    """Fragment ion types following ProForma notation"""
{iontype_enum_str}


IonTypeLiteral = Literal[
{iontype_literal_str}
]

ION_TYPE_DICT: dict[IonType, FragmentIonInfo] = {{
{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(fragment_ions)}")


if __name__ == "__main__":
    gen_frag()
