"""
Utility functions for PDB file formatting and manipulation.

This module centralizes PDB header/footer generation and content manipulation
to avoid code duplication across the codebase.
"""

from datetime import datetime
from typing import Optional


def get_current_date_pdb_format() -> str:
    """
    Returns the current date formatted for PDB HEADER record (DD-MON-YY).

    Returns:
        str: Date string in PDB format (e.g., "21-JAN-26")
    """
    now = datetime.now()
    month_abbr = now.strftime("%b").upper()
    return f"{now.day:02d}-{month_abbr}-{now.year % 100:02d}"


def extract_atomic_content(full_pdb_content: str) -> str:
    """
    Extracts only ATOM, HETATM, and TER lines from a full PDB content string.

    Args:
        full_pdb_content: Complete PDB file content including headers

    Returns:
        str: PDB content containing only atomic coordinate lines
    """
    atomic_lines = []
    for line in full_pdb_content.splitlines():
        if line.startswith("ATOM") or line.startswith("HETATM") or line.startswith("TER"):
            atomic_lines.append(line)
    return "\n".join(atomic_lines)


def create_pdb_header(
    sequence_length: int,
    date: Optional[str] = None,
    command_args: Optional[str] = None,
) -> str:
    """
    Creates standardized PDB header records.

    Args:
        sequence_length: Number of residues in the peptide
        date: Date string in PDB format (DD-MON-YY). If None, uses current date.
        command_args: Command-line arguments used for generation (for reproducibility)

    Returns:
        str: PDB header records (HEADER through MODEL)
    """
    if date is None:
        date = get_current_date_pdb_format()

    header_lines = []
    header_lines.append(f"HEADER    PEPTIDE           {date}    ")
    header_lines.append(f"TITLE     GENERATED LINEAR PEPTIDE OF LENGTH {sequence_length}")
    header_lines.append("REMARK 1  This PDB file was generated by the CLI 'synth-pdb' tool.")
    header_lines.append("REMARK 2  It represents a simplified model of a linear peptide chain.")
    header_lines.append(
        "REMARK 2  Coordinates are idealized and do not reflect real-world physics."
    )

    # Add command-line arguments for reproducibility (REMARK 3)
    if command_args:
        header_lines.append("REMARK 3  ")
        header_lines.append("REMARK 3  GENERATION PARAMETERS:")

        # Split command into multiple lines if needed (PDB REMARK lines are max 80 chars)
        # Format: "REMARK 3  " = 10 chars, leaving 70 chars for content
        if len(command_args) <= 68:  # 10 + 68 = 78, leaving some margin
            header_lines.append(f"REMARK 3  Command: {command_args}")
        else:
            # Split into multiple lines
            header_lines.append("REMARK 3  Command:")
            # Wrap at 68 characters per line
            for i in range(0, len(command_args), 68):
                chunk = command_args[i : i + 68]
                header_lines.append(f"REMARK 3    {chunk}")

        header_lines.append("REMARK 3  ")

    header_lines.append(
        f"COMPND    MOL_ID: 1; MOLECULE: LINEAR PEPTIDE; CHAIN: A; LENGTH: {sequence_length};"
    )
    header_lines.append("SOURCE    ENGINEERED; SYNTHETIC CONSTRUCT;")
    header_lines.append("KEYWDS    PEPTIDE, LINEAR, GENERATED, THEORETICAL MODEL")
    header_lines.append("EXPDTA    THEORETICAL MODEL")
    header_lines.append("AUTHOR    SYNTH PDB")
    header_lines.append("MODEL        1")

    return "\n".join(header_lines)


def create_pdb_footer() -> str:
    """
    Creates standardized PDB footer records.

    Returns:
        str: PDB footer records (ENDMDL and END)
    """
    footer_lines = [
        "ENDMDL",
        "END         ",  # 10 spaces to match typical PDB END record length
    ]
    return "\n".join(footer_lines)


def extract_header_records(full_pdb_content: str, record_type: str = "SSBOND") -> str:
    """
    Extracts specific header records (e.g. SSBOND) from a full PDB content string.

    Args:
        full_pdb_content: Complete PDB file content
        record_type: Record type to extract (default: SSBOND)

    Returns:
        str: Extracted records joined by newlines
    """
    records = []
    for line in full_pdb_content.splitlines():
        if line.startswith(record_type):
            records.append(line)
    return "\n".join(records)


def assemble_pdb_content(
    atomic_content: str,
    sequence_length: int,
    date: Optional[str] = None,
    command_args: Optional[str] = None,
    extra_records: Optional[str] = None,
    conect_records: Optional[str] = None,
) -> str:
    """
    Assembles complete PDB file content from atomic coordinates and metadata.

    Args:
        atomic_content: PDB content containing only ATOM/HETATM/TER lines
        sequence_length: Number of residues in the peptide
        date: Date string in PDB format. If None, uses current date.
        command_args: Command-line arguments for reproducibility
        extra_records: Optional string of extra header records (e.g. SSBONDs) to insert
        conect_records: Optional string of CONECT records to insert at the end

    Returns:
        str: Complete PDB file content with headers and footers
    """
    header = create_pdb_header(sequence_length, date, command_args)
    footer = create_pdb_footer()

    content = f"{header}\n"
    if extra_records:
        content += f"{extra_records}\n"
    content += f"{atomic_content.strip()}\n"
    if conect_records:
        content += f"{conect_records.strip()}\n"
    content += footer

    return content
