"""xTB input generation for workflow engine."""

import logging
from typing import Any, Optional


logger = logging.getLogger(__name__)

__all__ = [
    "generate_xtb_input_files",
    "generate_xtb_inputs",
]


def generate_xtb_input_files(
    node_type: str,
    params: dict[str, Any],
    structure_str: Optional[str],
) -> dict[str, str]:
    """Pure function: return {filename: content} for xTB calculation."""
    method = params.get("method", "GFN2-xTB")
    accuracy = params.get("accuracy", 1.0)
    etemp = params.get("electronic_temperature", 300)

    if node_type == "xtb_relax":
        fmax = params.get("fmax", 0.01)
        max_steps = params.get("max_steps", 500)
        script = f'''#!/usr/bin/env python3
"""xTB relaxation generated by CatGo workflow engine."""
from ase.io import read, write
from ase.optimize import BFGS

atoms = read("POSCAR", format="vasp")

# Try tblite first, fallback to xtb-python
method = "{method}"
try:
    from tblite.ase import TBLite
    calc = TBLite(method=method, accuracy={accuracy}, electronic_temperature={etemp})
except ImportError:
    from xtb.ase.calculator import XTB
    calc = XTB(method=method)

atoms.calc = calc
opt = BFGS(atoms, trajectory="opt.traj", logfile="opt.log")
opt.run(fmax={fmax}, steps={max_steps})

write("CONTCAR", atoms, format="vasp")
print(f"Final energy: {{atoms.get_potential_energy():.6f}} eV")
'''
    elif node_type == "xtb_static":
        script = f'''#!/usr/bin/env python3
"""xTB single-point calculation generated by CatGo workflow engine."""
from ase.io import read, write

atoms = read("POSCAR", format="vasp")

# Try tblite first, fallback to xtb-python
method = "{method}"
try:
    from tblite.ase import TBLite
    calc = TBLite(method=method, accuracy={accuracy}, electronic_temperature={etemp})
except ImportError:
    from xtb.ase.calculator import XTB
    calc = XTB(method=method)

atoms.calc = calc
energy = atoms.get_potential_energy()
forces = atoms.get_forces()

write("CONTCAR", atoms, format="vasp")
print(f"Final energy: {{energy:.6f}} eV")
print(f"Max force: {{forces.max():.6f}} eV/A")
'''
    else:
        raise RuntimeError(f"Unknown xTB node type: {node_type}")

    files = {"run_xtb.py": script}
    if structure_str:
        from workflow.engines import ensure_poscar
        try:
            files["POSCAR"] = ensure_poscar(structure_str)
        except Exception as e:
            logger.warning("Could not convert structure to POSCAR: %s", e)
            files["POSCAR"] = structure_str
    return files


async def generate_xtb_inputs(
    hpc: Any,
    work_dir: str,
    node_type: str,
    params: dict[str, Any],
    structure_str: Optional[str],
):
    """Generate xTB input files and upload to HPC."""
    from catgo.utils.job_parser import write_remote_files
    files = generate_xtb_input_files(node_type, params, structure_str)
    await write_remote_files(hpc.conn, {f"{work_dir}/{k}": v for k, v in files.items()})
