"""
Fragment dataclass for storing side chain information.
"""
from dataclasses import dataclass, field
from typing import List, Optional
from dataclasses_json import dataclass_json
import json


@dataclass_json
@dataclass
class Fragment:
    """
    Data class to store fragment information from macrocycle side chain cleavage.
    
    Attributes:
        fragment_smiles: SMILES string of the cleaved fragment (contains dummy atom '*' 
                        at the attachment point for molzip compatibility)
        parent_smiles: SMILES string of the parent molecule
        cleavage_position: Position on the ring where cleavage occurred (1-N, where N is ring size)
        fragment_id: Unique identifier for this fragment
        parent_id: Identifier of the parent molecule (e.g., "macro_0")
        atom_count: Number of atoms in the fragment
        molecular_weight: Molecular weight of the fragment
    
    Note:
        The fragment_smiles contains a dummy atom (*) at the attachment point, which allows
        for easy reconstruction using RDKit's molzip functionality. This format is compatible
        with SQLModel and PostgreSQL storage while maintaining chemical information.
    """
    fragment_smiles: str
    parent_smiles: str
    cleavage_position: int
    fragment_id: str
    parent_id: str
    atom_count: Optional[int] = None
    molecular_weight: Optional[float] = None
    
    def to_json_file(self, filepath: str):
        """Save fragment to JSON file."""
        with open(filepath, 'w') as f:
            json.dump(self.to_dict(), f, indent=2, ensure_ascii=False)
    
    @classmethod
    def from_json_file(cls, filepath: str):
        """Load fragment from JSON file."""
        with open(filepath, 'r') as f:
            data = json.load(f)
        return cls.from_dict(data)


@dataclass_json
@dataclass
class MoleculeFragments:
    """
    Collection of all fragments from a single macrocycle molecule.
    
    Attributes:
        parent_id: Identifier of the parent molecule
        parent_smiles: SMILES string of the parent molecule
        ring_numbering: Dictionary mapping atom indices to ring positions (1-16)
        fragments: List of Fragment objects
    """
    parent_id: str
    parent_smiles: str
    ring_numbering: dict = field(default_factory=dict)
    fragments: List[Fragment] = field(default_factory=list)
    
    def add_fragment(self, fragment: Fragment):
        """Add a fragment to the collection."""
        self.fragments.append(fragment)
    
    def to_json_file(self, filepath: str):
        """Save all fragments to JSON file."""
        with open(filepath, 'w') as f:
            json.dump(self.to_dict(), f, indent=2, ensure_ascii=False)
    
    @classmethod
    def from_json_file(cls, filepath: str):
        """Load fragments from JSON file."""
        with open(filepath, 'r') as f:
            data = json.load(f)
        return cls.from_dict(data)
    
    def get_cleavage_positions(self) -> List[int]:
        """Get list of all cleavage positions."""
        return [frag.cleavage_position for frag in self.fragments]
    
    def get_fragment_by_position(self, position: int) -> Optional[Fragment]:
        """Get fragment by cleavage position."""
        for frag in self.fragments:
            if frag.cleavage_position == position:
                return frag
        return None

