import csv
from pathlib import Path
from typing import Dict, List, Optional, Tuple

import selfies as sf
from rdkit import Chem


def canonicalize_smiles(smiles: str) -> Optional[str]:
    try:
        mol = Chem.MolFromSmiles(smiles)
        if mol is None:
            return None
        return Chem.MolToSmiles(mol, canonical=True)
    except Exception:
        return None


def normalize_selfies(selfies_text: str) -> Tuple[Optional[str], Optional[str]]:
    try:
        smiles = sf.decoder(selfies_text.strip())
    except Exception:
        return None, None
    smiles_canon = canonicalize_smiles(smiles)
    if smiles_canon is None:
        return None, None
    try:
        selfies_roundtrip = sf.encoder(smiles_canon)
        smiles_roundtrip = sf.decoder(selfies_roundtrip)
    except Exception:
        return None, None
    smiles_roundtrip_canon = canonicalize_smiles(smiles_roundtrip)
    if smiles_roundtrip_canon != smiles_canon:
        return None, None
    return smiles_canon, selfies_roundtrip


def normalize_smiles(smiles_text: str) -> Tuple[Optional[str], Optional[str]]:
    smiles_canon = canonicalize_smiles(smiles_text.strip())
    if smiles_canon is None:
        return None, None
    try:
        selfies_roundtrip = sf.encoder(smiles_canon)
        smiles_roundtrip = sf.decoder(selfies_roundtrip)
    except Exception:
        return None, None
    smiles_roundtrip_canon = canonicalize_smiles(smiles_roundtrip)
    if smiles_roundtrip_canon != smiles_canon:
        return None, None
    return smiles_canon, selfies_roundtrip


def preprocess_lines(
    lines: List[str],
    input_format: str,
    model_input_type: str,
) -> Tuple[List[Dict[str, str]], List[Dict[str, str]]]:
    """Apply circularity validation and return valid records and failures."""
    fmt = input_format.strip().lower()
    model_fmt = model_input_type.strip().lower()
    if fmt not in {"selfies", "smiles"}:
        raise ValueError(f"Unsupported input format: {input_format}")
    if model_fmt not in {"selfies", "smiles"}:
        raise ValueError(f"Unsupported model input type: {model_input_type}")

    records: List[Dict[str, str]] = []
    failures: List[Dict[str, str]] = []

    for idx, raw in enumerate(lines, start=1):
        text = raw.strip()
        if not text:
            continue

        if fmt == "selfies":
            smiles_canon, selfies_norm = normalize_selfies(text)
        else:
            smiles_canon, selfies_norm = normalize_smiles(text)

        if smiles_canon is None or selfies_norm is None:
            failures.append(
                {
                    "original_line": str(idx),
                    "input_text": text,
                    "reason": "circularity_failed",
                }
            )
            continue

        model_input = selfies_norm if model_fmt == "selfies" else smiles_canon
        records.append(
            {
                "original_line": str(idx),
                "input_text": text,
                "smiles": smiles_canon,
                "selfies": selfies_norm,
                "model_input": model_input,
            }
        )

    return records, failures


def load_input_lines(path: Path) -> List[str]:
    with path.open("r", encoding="utf-8") as handle:
        return [line.rstrip("\n") for line in handle]


def resolve_output_paths(output_dir: Path, output_stem: str) -> Tuple[Path, Path, Path]:
    output_dir.mkdir(parents=True, exist_ok=True)
    npy_path = output_dir / f"{output_stem}.npy"
    map_path = output_dir / f"{output_stem}_map.csv"
    failures_path = output_dir / f"{output_stem}_failures.csv"
    return npy_path, map_path, failures_path


def write_mapping_csv(path: Path, records: List[Dict[str, str]]) -> None:
    with path.open("w", newline="", encoding="utf-8") as csvfile:
        writer = csv.writer(csvfile)
        writer.writerow(
            [
                "row_index",
                "original_line",
                "input_text",
                "smiles",
                "selfies",
                "model_input",
            ]
        )
        for row_index, record in enumerate(records):
            writer.writerow(
                [
                    row_index,
                    record["original_line"],
                    record["input_text"],
                    record["smiles"],
                    record["selfies"],
                    record["model_input"],
                ]
            )


def write_failures_csv(path: Path, failures: List[Dict[str, str]]) -> None:
    with path.open("w", newline="", encoding="utf-8") as csvfile:
        writer = csv.writer(csvfile)
        writer.writerow(["original_line", "input_text", "reason"])
        for row in failures:
            writer.writerow([row["original_line"], row["input_text"], row["reason"]])
