# ============================
# exporters.py — Structure Export Utilities (ASE-Enhanced)
# ============================
from __future__ import annotations
import numpy as np
from typing import List, Optional, Tuple, Sequence
try:
    from ase import Atoms
    from ase.io import write
    HAS_ASE = True
except ImportError:
    HAS_ASE = False

def export_structure_xyz(filename: str, positions: np.ndarray, lattice: np.ndarray, 
                         elements: List[str], energy_total: float = 0.0, 
                         energy_formation: float = 0.0, append: bool = False):
    """Exports the current structure to an Extended XYZ file (Custom implementation)."""
    if not filename.endswith(".extxyz") and not filename.endswith(".xyz"): 
        filename += ".xyz"
    
    try:
        N = positions.shape[0]
        mode = "a" if append else "w"
        with open(filename, mode) as f:
            f.write(f"{N}\n")
            
            # Construct Lattice String: ASE expects vectors as rows in cell matrix
            L = lattice.flatten()
            lat_str = " ".join([f"{v:.8f}" for v in L])
            
            # Header Line with Properties and Energy metadata
            header = f'Lattice="{lat_str}" Properties=species:S:1:pos:R:3 energy={energy_total:.6f} formation_energy={energy_formation:.6f} pbc="T T T"'
            f.write(f"{header}\n")
            
            species = elements if elements else ["X"] * N
            for i in range(N):
                e = species[i] if i < len(species) else "X"
                x, y, z = positions[i]
                f.write(f"{e} {x:.8f} {y:.8f} {z:.8f}\n")
        
        return True, f"Exported: {filename}"
        
    except Exception as e:
        return False, f"Export Failed: {e}"

def export_structure(filename: str, fmt: str, positions: np.ndarray, lattice: np.ndarray,
                     elements: Sequence[str], energy: float = 0.0, append: bool = False) -> Tuple[bool, str]:
    """
    Exports structural data to various formats using ASE if available.
    Supports formats like 'cif', 'pdb', 'vasp' (POSCAR), 'aims', etc.
    """
    if fmt.lower() in ["xyz", "extxyz"]:
        return export_structure_xyz(filename, positions, lattice, list(elements), energy_total=energy, append=append)

    if not HAS_ASE:
        return False, "ASE is required for this export format."
    
    try:
        # Map generic format names to ASE format names if needed
        ase_fmt = fmt.lower()
        ext = ase_fmt
        if ase_fmt == "poscar" or ase_fmt == "vasp":
            ase_fmt = "vasp"
            ext = "poscar"
        elif ase_fmt == "aims":
            ext = "in"
        
        if not filename.lower().endswith(f".{ext}"):
            filename += f".{ext}"
        
        # Build ASE Atoms object
        # Note: ASE cell takes 3x3 lattice vectors
        atoms = Atoms(symbols=elements, positions=positions, cell=lattice, pbc=True)
        # Store metadata if applicable
        atoms.info['energy'] = energy
        
        # Write using ASE
        write(filename, atoms, format=ase_fmt, append=append)
        return True, f"Exported: {filename}"
        
    except Exception as e:
        return False, f"Export Failed: {e}"
