import csv
import os
import re
from pathlib import Path
from typing import Any, 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 _normalize_input_entry(raw: Any, fallback_line: int) -> Dict[str, Any]:
    if isinstance(raw, dict):
        entry = dict(raw)
        entry.setdefault("original_line", str(fallback_line))
        entry.setdefault("input_text", "")
        entry["input_text"] = str(entry.get("input_text", ""))
        return entry

    return {
        "original_line": str(fallback_line),
        "input_text": str(raw),
        "source_type": "txt",
    }


def preprocess_lines(
    lines: List[Any],
    input_format: str,
    model_input_type: str,
) -> Tuple[List[Dict[str, Any]], List[Dict[str, Any]]]:
    """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, Any]] = []
    failures: List[Dict[str, Any]] = []

    for idx, raw in enumerate(lines, start=1):
        entry = _normalize_input_entry(raw, idx)
        text = str(entry.get("input_text", "")).strip()
        if not text:
            continue

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

        shared_fields: Dict[str, Any] = {
            "original_line": str(entry.get("original_line", idx)),
            "input_text": text,
        }
        if "source_type" in entry:
            shared_fields["source_type"] = entry.get("source_type")
        if "source_header" in entry:
            shared_fields["source_header"] = entry.get("source_header")
        if "source_row" in entry:
            shared_fields["source_row"] = entry.get("source_row")

        if smiles_canon is None or selfies_norm is None:
            failed = dict(shared_fields)
            failed["reason"] = "circularity_failed"
            failures.append(failed)
            continue

        model_input = selfies_norm if model_fmt == "selfies" else smiles_canon
        record = dict(shared_fields)
        record.update(
            {
                "smiles": smiles_canon,
                "selfies": selfies_norm,
                "model_input": model_input,
            }
        )
        records.append(record)

    return records, failures


def load_input_lines(path: Path) -> List[Dict[str, Any]]:
    """Load molecule inputs from .txt or .csv.

    For .csv, assumes first row is a header and first column stores molecules.
    The returned entries preserve full source rows to support filtered CSV outputs.
    """
    suffix = path.suffix.lower()

    if suffix == ".csv":
        entries: List[Dict[str, Any]] = []
        with path.open("r", encoding="utf-8-sig", newline="") as handle:
            reader = csv.reader(handle)
            try:
                header = next(reader)
            except StopIteration:
                return entries

            header = [str(col) for col in header]
            if not header:
                header = ["molecule"]

            for csv_line_no, row in enumerate(reader, start=2):
                if row is None:
                    continue
                source_row = [str(cell) for cell in row]
                if not source_row or not any(cell.strip() for cell in source_row):
                    continue

                molecule_text = source_row[0].strip() if source_row else ""
                if len(source_row) < len(header):
                    source_row = source_row + [""] * (len(header) - len(source_row))

                entries.append(
                    {
                        "original_line": str(csv_line_no),
                        "input_text": molecule_text,
                        "source_type": "csv",
                        "source_header": list(header),
                        "source_row": source_row,
                    }
                )
        return entries

    entries = []
    with path.open("r", encoding="utf-8") as handle:
        for line_no, line in enumerate(handle, start=1):
            entries.append(
                {
                    "original_line": str(line_no),
                    "input_text": line.rstrip("\n"),
                    "source_type": "txt",
                }
            )
    return entries


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 _base_stem_for_sidecar(path: Path) -> str:
    stem = path.stem
    if stem.endswith("_map"):
        return stem[: -len("_map")]
    if stem.endswith("_failures"):
        return stem[: -len("_failures")]
    return stem


def _normalize_csv_row(row: Any, width: int) -> List[str]:
    if isinstance(row, list):
        cells = [str(cell) for cell in row]
    else:
        cells = [str(row)] if row is not None else []

    if len(cells) < width:
        cells.extend([""] * (width - len(cells)))
    if len(cells) > width:
        cells = cells[:width]
    return cells


def _extract_csv_header(rows: List[Dict[str, Any]]) -> List[str]:
    for row in rows:
        header = row.get("source_header")
        if isinstance(header, list) and header:
            return [str(col) for col in header]
    return []


def _write_filtered_csv(mapping_path: Path, records: List[Dict[str, Any]]) -> None:
    filtered_path = mapping_path.with_name(f"{_base_stem_for_sidecar(mapping_path)}_filtered.csv")
    header = _extract_csv_header(records)

    if header:
        with filtered_path.open("w", newline="", encoding="utf-8") as csvfile:
            writer = csv.writer(csvfile)
            writer.writerow(header)
            for record in records:
                writer.writerow(_normalize_csv_row(record.get("source_row", []), len(header)))
        return

    # Text fallback: keep a simple one-column filtered view aligned with npy rows.
    with filtered_path.open("w", newline="", encoding="utf-8") as csvfile:
        writer = csv.writer(csvfile)
        writer.writerow(["molecule"])
        for record in records:
            writer.writerow([record.get("input_text", "")])


def _write_map_csv_enabled() -> bool:
    raw = os.environ.get("MOREHUB_WRITE_MAP_CSV", "0").strip().lower()
    return raw in {"1", "true", "yes", "on"}


def write_mapping_csv(path: Path, records: List[Dict[str, Any]], write_map: Optional[bool] = None) -> None:
    if write_map is None:
        write_map = _write_map_csv_enabled()

    if write_map:
        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"],
                    ]
                )
    elif path.exists():
        path.unlink()

    _write_filtered_csv(path, records)


def write_failures_csv(path: Path, failures: List[Dict[str, Any]]) -> None:
    header = _extract_csv_header(failures)

    with path.open("w", newline="", encoding="utf-8") as csvfile:
        writer = csv.writer(csvfile)
        if header:
            writer.writerow(["original_line", "input_text", "reason", *header])
            for row in failures:
                writer.writerow(
                    [
                        row.get("original_line", ""),
                        row.get("input_text", ""),
                        row.get("reason", ""),
                        *_normalize_csv_row(row.get("source_row", []), len(header)),
                    ]
                )
            return

        writer.writerow(["original_line", "input_text", "reason"])
        for row in failures:
            writer.writerow([row["original_line"], row["input_text"], row["reason"]])

