"""
Tool I/O Type System for DAG workflow data passing.

Provides:
- FileType metadata via typing.Annotated
- Pre-defined file type aliases (PDBFile, SDFFile, etc.)
- ToolFile / ToolOutput dataclasses for structured runtime output
- @tool_io decorator for static I/O declaration on Tool classes
- Helper to extract contracts at planning time without instantiation
"""
from __future__ import annotations

import os
from dataclasses import dataclass, field
from typing import Annotated, Any, get_args, get_origin


# ---------------------------------------------------------------------------
# 1. File type metadata (attached to Annotated aliases)
# ---------------------------------------------------------------------------

@dataclass(frozen=True)
class FileType:
    """Compile-time file type metadata, extractable via get_type_hints()."""
    ext: str
    description: str


# Pre-defined type aliases - runtime value is just ``str`` (a file path).
# Grouped by domain for readability.

# --- Structure files ---
PDBFile    = Annotated[str, FileType(".pdb",    "Protein Data Bank structure file")]
CIFFile    = Annotated[str, FileType(".cif",    "mmCIF crystallographic structure file")]
MMCIFFile  = Annotated[str, FileType(".mmcif",  "Macromolecular CIF structure file")]
MMTFFile   = Annotated[str, FileType(".mmtf",   "Macromolecular Transmission Format file")]
PDBQTFile  = Annotated[str, FileType(".pdbqt",  "AutoDock PDBQT structure file")]
GROFile    = Annotated[str, FileType(".gro",    "GROMACS structure file")]
SDFFile    = Annotated[str, FileType(".sdf",    "Structure Data File (molecules)")]
MOLFile    = Annotated[str, FileType(".mol",    "MDL MOL molecular structure file")]
MOL2File   = Annotated[str, FileType(".mol2",   "Sybyl MOL2 molecular structure file")]
XYZFile    = Annotated[str, FileType(".xyz",    "Atomic coordinate file")]
PQRFile    = Annotated[str, FileType(".pqr",    "PQR charge/radius structure file")]
STRUFile   = Annotated[str, FileType(".stru",   "ABACUS atomic structure file")]
VASPFile   = Annotated[str, FileType(".vasp",   "VASP POSCAR/CONTCAR structure file")]

# --- Sequence / alignment files ---
FASTAFile  = Annotated[str, FileType(".fasta",  "Protein/nucleotide sequence file")]
A3MFile    = Annotated[str, FileType(".a3m",    "Multiple sequence alignment file")]

# --- Topology files ---
PRMTOPFile = Annotated[str, FileType(".prmtop", "AMBER topology file")]
PSFFile    = Annotated[str, FileType(".psf",    "Protein Structure File (CHARMM/NAMD topology)")]
TOPFile    = Annotated[str, FileType(".top",    "GROMACS topology file")]

# --- Trajectory files ---
DCDFile    = Annotated[str, FileType(".dcd",    "DCD binary trajectory file")]
XTCFile    = Annotated[str, FileType(".xtc",    "GROMACS compressed trajectory file")]
TRRFile    = Annotated[str, FileType(".trr",    "GROMACS full-precision trajectory file")]
NetCDFFile = Annotated[str, FileType(".netcdf", "NetCDF trajectory file")]
NCTRAJFile = Annotated[str, FileType(".nctraj", "NetCDF trajectory file (Amber)")]
LAMMPSFile = Annotated[str, FileType(".lammps", "LAMMPS trajectory/data file")]

# --- Electron density / volumetric ---
MRCFile    = Annotated[str, FileType(".mrc",    "MRC/CCP4 electron density map")]
MAPFile    = Annotated[str, FileType(".map",    "CCP4 electron density map")]
CCP4File   = Annotated[str, FileType(".ccp4",   "CCP4 electron density map")]
DSN6File   = Annotated[str, FileType(".dsn6",   "DSN6 electron density map")]
DXFile     = Annotated[str, FileType(".dx",     "OpenDX volumetric data file")]
CubeFile   = Annotated[str, FileType(".cube",   "Gaussian Cube volumetric file")]

# --- Tabular / data files ---
CSVFile    = Annotated[str, FileType(".csv",    "Tabular data file")]
JSONFile   = Annotated[str, FileType(".json",   "Structured data file")]
TXTFile    = Annotated[str, FileType(".txt",    "Plain text file")]
NPYFile    = Annotated[str, FileType(".npy",    "NumPy array file")]

# --- Image / mesh files ---
PNGFile    = Annotated[str, FileType(".png",    "Image file")]
OBJFile    = Annotated[str, FileType(".obj",    "Wavefront OBJ 3D mesh file")]
PLYFile    = Annotated[str, FileType(".ply",    "PLY polygon mesh file")]


# ---------------------------------------------------------------------------
# 2. Runtime output structures
# ---------------------------------------------------------------------------

@dataclass
class ToolFile:
    """One output port — single file or a list of files of the same type."""
    type: str                            # e.g. "pdb", "sdf"
    description: str                     # human-readable
    path: str | None = None              # single file
    paths: list[str] | None = None       # multiple files of the same type

    @property
    def is_multi(self) -> bool:
        return bool(self.paths)

    @property
    def all_paths(self) -> list[str]:
        """Return all file paths regardless of single/multi mode."""
        if self.paths:
            return [p for p in self.paths if p]
        return [self.path] if self.path else []


@dataclass
class ToolOutput:
    """Unified output returned by Tool.output()."""
    files: dict[str, ToolFile] = field(default_factory=dict)
    text: str = ""           # user-facing display text (replaces outputShow 1st return)
    summary: str = ""        # LLM-facing summary (replaces summary() return)
    metrics: dict[str, Any] = field(default_factory=dict)


# ---------------------------------------------------------------------------
# 3. @tool_io decorator
# ---------------------------------------------------------------------------

def tool_io(inputs: dict | None = None, outputs: dict | None = None):
    """Attach static I/O type declarations to a Tool subclass."""
    def decorator(cls):
        cls.__tool_io__ = {
            "inputs": inputs or {},
            "outputs": outputs or {},
        }
        return cls
    return decorator


# ---------------------------------------------------------------------------
# 4. Extraction helpers (used at planning time)
# ---------------------------------------------------------------------------

def extract_file_type(annotated_type) -> FileType | None:
    """Pull FileType metadata out of an Annotated alias."""
    if get_origin(annotated_type) is Annotated:
        for meta in get_args(annotated_type)[1:]:
            if isinstance(meta, FileType):
                return meta
    return None


def get_tool_contract(tool_cls) -> dict | None:
    """
    Read a Tool class's I/O contract without instantiation.

    Returns ``None`` for legacy tools that lack ``@tool_io``.
    """
    io = getattr(tool_cls, "__tool_io__", None)
    if not io:
        return None

    def _resolve(mapping: dict) -> dict:
        result = {}
        for key, typ in mapping.items():
            ft = extract_file_type(typ)
            if ft:
                result[key] = {"ext": ft.ext, "description": ft.description}
            else:
                result[key] = {"raw_type": str(typ)}
        return result

    return {
        "inputs": _resolve(io["inputs"]),
        "outputs": _resolve(io["outputs"]),
    }
