"""
Module for writing CHARMM CARD coordinate files (CRD).

This module provides functions to write coordinate data from crimm entities
to CHARMM CRD format files. The CRD format stores atomic coordinates along
with residue and segment information.

Format specification (from CHARMM source io/coorio.F90):
- Extended format: 2I10, 2X,A8, 2X,A8, 3F20.10, 2X,A8, 2X,A8, F20.10
- Standard format: 2I5, 1X,A4, 1X,A4, 3F10.5, 1X,A4, 1X,A4, F10.5
"""

from typing import Union, List
from crimm.StructEntities.Atom import Atom
from crimm.StructEntities.Residue import Residue
from crimm.StructEntities.Chain import Chain
from crimm.StructEntities.Model import Model
from crimm.StructEntities.Structure import Structure


# Extended format (EXT) - for systems with >99999 atoms
_EXT_HEADER_FORMAT = "{:>10d}  EXT\n"
_EXT_ATOM_FORMAT = (
    "{serial:>10d}{resno:>10d}  {resname:<8s}  {atomname:<8s}"
    "{x:>20.10f}{y:>20.10f}{z:>20.10f}  {segid:<8s}  {resid:<8s}{weight:>20.10f}\n"
)

# Standard format - for smaller systems
_STD_HEADER_FORMAT = "{:>5d}\n"
_STD_ATOM_FORMAT = (
    "{serial:>5d}{resno:>5d} {resname:<4s} {atomname:<4s}"
    "{x:>10.5f}{y:>10.5f}{z:>10.5f} {segid:<4s} {resid:<4s}{weight:>10.5f}\n"
)


def _get_atoms_from_entity(entity, include_lonepairs: bool = True) -> List[Atom]:
    """Get all atoms from an entity, optionally including lone pairs.

    Parameters
    ----------
    entity : Model, Structure, Chain, or Residue
        The entity to extract atoms from
    include_lonepairs : bool
        Whether to include lone pair pseudo-atoms for CGENFF ligands

    Returns
    -------
    List[Atom]
        List of atoms in order
    """
    atoms = []

    if hasattr(entity, 'get_atoms'):
        atoms.extend(list(entity.get_atoms()))
    elif isinstance(entity, Atom):
        atoms.append(entity)

    # Include lone pairs if present and requested
    if include_lonepairs:
        if hasattr(entity, 'get_residues'):
            for residue in entity.get_residues():
                if hasattr(residue, 'lone_pair_dict') and residue.lone_pair_dict:
                    atoms.extend(residue.lone_pair_dict.values())
        elif hasattr(entity, 'lone_pair_dict') and entity.lone_pair_dict:
            atoms.extend(entity.lone_pair_dict.values())

    return atoms


def _generate_system_info(entity) -> list:
    """Generate informative title lines from entity metadata.

    Extracts useful information about the system:
    - Structure/PDB ID
    - Chain IDs and types
    - Disulfide bonds
    - CGenFF usage for ligands
    - Total atom count

    Parameters
    ----------
    entity : Model, Structure, Chain, or Residue
        The entity to extract information from

    Returns
    -------
    list
        List of title lines describing the system
    """
    import datetime
    import getpass

    lines = []

    # Get timestamp and user
    timestamp = datetime.datetime.now().strftime("%Y-%m-%d %H:%M:%S")
    try:
        user = getpass.getuser()
    except Exception:
        user = "unknown"

    # Get structure ID (PDB ID)
    structure_id = None
    model = None

    if hasattr(entity, 'level'):
        if entity.level == 'M':  # Model
            model = entity
            if hasattr(entity, 'parent') and entity.parent is not None:
                structure_id = entity.parent.id
        elif entity.level == 'S':  # Structure
            structure_id = entity.id
            models = list(entity.get_models())
            if models:
                model = models[0]
        elif entity.level == 'C':  # Chain
            model = entity.parent
            if model and hasattr(model, 'parent') and model.parent:
                structure_id = model.parent.id

    # Build first line
    if structure_id:
        lines.append(f"Structure: {structure_id} | Generated by crimm")
    else:
        lines.append("Generated by crimm")

    # Get chain info
    chain_info = []
    disu_count = 0
    has_cgenff = False
    has_tip3 = False
    ion_types = set()

    if model is not None:
        for chain in model:
            chain_id = chain.id
            chain_type = getattr(chain, 'chain_type', 'Unknown')

            # Count disulfide bonds from patches
            if hasattr(chain, 'topology') and chain.topology is not None:
                topo = chain.topology
                if hasattr(topo, 'inter_res_patches'):
                    for patch in topo.inter_res_patches:
                        if hasattr(patch, 'name') and 'DISU' in patch.name.upper():
                            disu_count += 1
                # Check for CGenFF
                if hasattr(topo, 'source') and topo.source:
                    if 'cgenff' in str(topo.source).lower():
                        has_cgenff = True

            # Categorize chains
            if chain_type in ('Polypeptide(L)', 'Protein'):
                n_res = len(list(chain.get_residues()))
                chain_info.append(f"{chain_id}(protein,{n_res}res)")
            elif chain_type == 'Solvent':
                has_tip3 = True
            elif chain_type == 'Ion':
                for res in chain.get_residues():
                    ion_types.add(res.resname.strip())
            elif chain_type == 'Ligand':
                chain_info.append(f"{chain_id}(ligand)")
            elif chain_type in ('Polyribonucleotide', 'RNA'):
                chain_info.append(f"{chain_id}(RNA)")
            elif chain_type in ('Polydeoxyribonucleotide', 'DNA'):
                chain_info.append(f"{chain_id}(DNA)")

    # Add chain summary
    if chain_info:
        lines.append(f"Chains: {', '.join(chain_info)}")

    # Count waters and ions
    n_waters = 0
    n_ions = 0
    if model is not None:
        for chain in model:
            chain_type = getattr(chain, 'chain_type', 'Unknown')
            if chain_type == 'Solvent':
                n_waters = len(list(chain.get_residues()))
            elif chain_type == 'Ion':
                n_ions = len(list(chain.get_residues()))

    # Check for solvation metadata (salt concentration, preserved waters)
    solvation_info = []
    if model is not None and hasattr(model, '_solvation_info'):
        sol_info = model._solvation_info
        if 'concentration' in sol_info:
            solvation_info.append(f"{sol_info['concentration']*1000:.0f} mM salt")
        if sol_info.get('preserved_waters', 0) > 0:
            solvation_info.append(f"{sol_info['preserved_waters']} crystal waters kept")
        if sol_info.get('preserved_ions', 0) > 0:
            solvation_info.append(f"{sol_info['preserved_ions']} crystal ions kept")

    # Add solvation summary
    if n_waters > 0 or n_ions > 0:
        solv_parts = []
        if n_waters > 0:
            solv_parts.append(f"{n_waters} waters")
        if n_ions > 0:
            solv_parts.append(f"{n_ions} ions ({','.join(sorted(ion_types))})")
        if solvation_info:
            solv_parts.extend(solvation_info)
        lines.append(f"Solvation: {', '.join(solv_parts)}")

    # Add crystal/box information
    if model is not None and hasattr(model, '_solvation_info'):
        sol_info = model._solvation_info
        box_type = sol_info.get('box_type')
        box_dim = sol_info.get('box_dim')
        if box_type and box_dim:
            box_type_names = {
                'cube': 'Cubic',
                'octa': 'Truncated Octahedron',
                'rhdo': 'Rhombic Dodecahedron',
                'ortho': 'Orthorhombic',
                'tetra': 'Tetragonal',
                'hexa': 'Hexagonal',
                'mono': 'Monoclinic',
                'tric': 'Triclinic',
                'rhomb': 'Rhombohedral',
            }
            box_name = box_type_names.get(box_type, box_type.capitalize())
            # Handle orthorhombic with different a, b, c dimensions
            box_dims = sol_info.get('box_dims')
            if box_type == 'ortho' and box_dims is not None:
                a, b, c = box_dims
                lines.append(f"Crystal: {box_name} box, {a:.2f} x {b:.2f} x {c:.2f} A")
            else:
                lines.append(f"Crystal: {box_name} box, {box_dim:.2f} A")

    # Add special features
    features = []
    if disu_count > 0:
        features.append(f"{disu_count} DISU")
    if has_cgenff:
        features.append("CGenFF")

    if features:
        lines.append(f"Features: {' | '.join(features)}")

    # Add timestamp
    lines.append(f"Created: {timestamp} by {user}")

    return lines


def _format_title_lines(title: str = "", entity=None) -> str:
    """Format title lines with asterisk prefix.

    CHARMM CRD format requires title section to have at least two lines:
    one or more content lines starting with "* " followed by a terminating "*".

    Parameters
    ----------
    title : str, optional
        Title text (can be multiline). If empty and entity is provided,
        system info is extracted from the entity. Otherwise a basic
        default title is used.
    entity : Model, Structure, Chain, or Residue, optional
        Entity to extract system information from for auto-generated title.

    Returns
    -------
    str
        Formatted title lines with asterisk prefix
    """
    import datetime
    import getpass

    lines = []
    if title:
        # User-provided title
        for line in title.strip().split('\n'):
            lines.append(f"* {line}\n")
    elif entity is not None:
        # Auto-generate informative title from entity
        for line in _generate_system_info(entity):
            lines.append(f"* {line}\n")
    else:
        # Basic default title
        timestamp = datetime.datetime.now().strftime("%Y-%m-%d %H:%M:%S")
        try:
            user = getpass.getuser()
        except Exception:
            user = "unknown"
        lines.append(f"* CRD file generated by crimm\n")
        lines.append(f"* Created: {timestamp} by {user}\n")
    # Terminating asterisk (required by CHARMM format)
    lines.append("*\n")
    return "".join(lines)


def get_crd_str(
    entity: Union[Model, Structure, Chain, Residue],
    extended: bool = True,
    title: str = "",
    include_lonepairs: bool = True,
    reset_serial: bool = True
) -> str:
    """Get CHARMM CRD format string from a crimm entity.

    Parameters
    ----------
    entity : Model, Structure, Chain, or Residue
        The entity to write coordinates from
    extended : bool, default True
        Use extended format (I10/A8/F20.10) for large systems
    title : str, default ""
        Title line(s) to include in the header
    include_lonepairs : bool, default True
        Include lone pair pseudo-atoms for CGENFF ligands
    reset_serial : bool, default True
        Reset atom serial numbers starting from 1

    Returns
    -------
    str
        CRD format string
    """
    # Get atoms
    atoms = _get_atoms_from_entity(entity, include_lonepairs)
    natom = len(atoms)

    if natom == 0:
        raise ValueError("No atoms found in entity")

    # Select format based on extended flag and atom count
    if extended or natom >= 100000:
        header_fmt = _EXT_HEADER_FORMAT
        atom_fmt = _EXT_ATOM_FORMAT
        str_width = 8
    else:
        header_fmt = _STD_HEADER_FORMAT
        atom_fmt = _STD_ATOM_FORMAT
        str_width = 4

    # Build output
    lines = []

    # Title section (pass entity for auto-generated system info)
    lines.append(_format_title_lines(title, entity=entity))

    # Atom count header
    lines.append(header_fmt.format(natom))

    # Build global residue numbering map
    # RESNO in CRD format must be sequential across the entire system (1, 2, 3...)
    # while RESID is the segment-local residue identifier
    resno_map = {}  # Maps residue object to global resno
    global_resno = 1
    last_residue = None

    for atom in atoms:
        residue = atom.parent
        if residue is not last_residue and residue is not None:
            if residue not in resno_map:
                resno_map[residue] = global_resno
                global_resno += 1
            last_residue = residue

    # Atom lines
    for idx, atom in enumerate(atoms, start=1):
        serial = idx if reset_serial else (atom.serial_number or idx)

        # Get parent residue info
        residue = atom.parent
        if residue is not None:
            resname = residue.resname[:str_width]
            if resname == 'HIS' and residue.topo_definition is not None:
                # Use residue definition resname for HIS variants
                resname = residue.topo_definition.resname
            segid = (residue.segid or "")[:str_width]
            # Use global sequential RESNO for CHARMM compatibility
            resno = resno_map.get(residue, 1)
            # RESID is the segment-local residue ID as string
            resid = str(residue.id[1])[:str_width]
        else:
            resname = "UNK"[:str_width]
            segid = ""[:str_width]
            resno = 1
            resid = "1"[:str_width]

        atomname = atom.name[:str_width]

        # Get coordinates
        if atom.coord is None:
            raise ValueError(f"Atom {atom.name} (serial {serial}) has no coordinates")
        x, y, z = atom.coord

        # Validate coordinates are finite numbers
        import numpy as np
        if not all(np.isfinite([x, y, z])):
            raise ValueError(
                f"Atom {atom.name} (serial {serial}) has invalid coordinates: "
                f"({x}, {y}, {z}). Coordinates must be finite numbers."
            )

        # Weight/temperature factor
        weight = atom.bfactor if atom.bfactor is not None else 0.0

        lines.append(atom_fmt.format(
            serial=serial,
            resno=resno,
            resname=resname,
            atomname=atomname,
            x=x,
            y=y,
            z=z,
            segid=segid,
            resid=resid,
            weight=weight
        ))

    return "".join(lines)


def write_crd(
    entity: Union[Model, Structure, Chain, Residue],
    filepath: str,
    extended: bool = True,
    title: str = "",
    include_lonepairs: bool = True,
    reset_serial: bool = True
) -> None:
    """Write CHARMM CRD format file from a crimm entity.

    Parameters
    ----------
    entity : Model, Structure, Chain, or Residue
        The entity to write coordinates from
    filepath : str
        Path to output file
    extended : bool, default True
        Use extended format (I10/A8/F20.10) for large systems
    title : str, default ""
        Title line(s) to include in the header
    include_lonepairs : bool, default True
        Include lone pair pseudo-atoms for CGENFF ligands
    reset_serial : bool, default True
        Reset atom serial numbers starting from 1
    """
    crd_str = get_crd_str(
        entity,
        extended=extended,
        title=title,
        include_lonepairs=include_lonepairs,
        reset_serial=reset_serial
    )

    with open(filepath, 'w', encoding='utf-8') as f:
        f.write(crd_str)
