from __future__ import annotations

import json
from pathlib import Path

import pandas as pd

from .analyzer import MacroLactoneAnalyzer
from .errors import MacrolactoneError
from .fragmenter import MacrolactoneFragmenter
from .models import FragmentationResult
from .visualization import save_numbered_molecule_png


def fragment_csv(
    input_csv: str | Path,
    smiles_column: str = "smiles",
    id_column: str = "id",
    ring_size: int | None = None,
    max_rows: int | None = None,
) -> list[FragmentationResult]:
    results, errors = _fragment_csv_with_errors(
        input_csv=input_csv,
        smiles_column=smiles_column,
        id_column=id_column,
        ring_size=ring_size,
        max_rows=max_rows,
    )
    if errors:
        first_error = errors[0]
        raise first_error["exception"]
    return results


def results_to_dataframe(results: list[FragmentationResult]) -> pd.DataFrame:
    rows: list[dict] = []
    for result in results:
        for fragment in result.fragments:
            rows.append(
                {
                    "parent_id": result.parent_id,
                    "parent_smiles": result.parent_smiles,
                    "ring_size": result.ring_size,
                    **fragment.to_dict(),
                }
            )
    return pd.DataFrame(rows)


def write_result_json(result: FragmentationResult, output_path: str | Path) -> Path:
    path = Path(output_path)
    path.parent.mkdir(parents=True, exist_ok=True)
    path.write_text(json.dumps(result.to_dict(), indent=2, ensure_ascii=False) + "\n", encoding="utf-8")
    return path


def export_numbered_macrolactone_csv(
    input_csv: str | Path,
    output_dir: str | Path,
    smiles_column: str = "smiles",
    id_column: str = "id",
    output_csv_name: str = "numbered_macrolactones.csv",
    ring_size: int | None = None,
    allowed_ring_atom_types: list[str] | None = None,
    image_size: tuple[int, int] = (800, 800),
    dpi: int = 600,
) -> Path:
    analyzer = MacroLactoneAnalyzer()
    output_dir = Path(output_dir)
    images_dir = output_dir / "images"
    output_dir.mkdir(parents=True, exist_ok=True)
    images_dir.mkdir(parents=True, exist_ok=True)

    rows = _read_csv_rows(
        input_csv=input_csv,
        smiles_column=smiles_column,
        id_column=id_column,
    )
    export_rows: list[dict] = []
    for row in rows:
        record = {
            "parent_id": row["parent_id"],
            "smiles": row["smiles"],
            "status": "success",
            "image_path": "",
            "classification": None,
            "primary_reason_code": None,
            "ring_size": None,
            "candidate_ring_sizes": [],
            "error_type": None,
            "error_message": None,
        }
        try:
            classification = analyzer.classify_macrocycle(row["smiles"], ring_size=ring_size)
            record.update(
                {
                    "classification": classification.classification,
                    "primary_reason_code": classification.primary_reason_code,
                    "ring_size": classification.ring_size,
                    "candidate_ring_sizes": classification.candidate_ring_sizes,
                }
            )
            image_path = images_dir / f"{row['parent_id']}.png"
            save_numbered_molecule_png(
                row["smiles"],
                image_path,
                ring_size=ring_size,
                size=image_size,
                allowed_ring_atom_types=allowed_ring_atom_types,
                dpi=dpi,
            )
            record["image_path"] = str(image_path.relative_to(output_dir.parent))
        except Exception as exc:  # pragma: no cover - surfaced in CSV
            record.update(
                {
                    "status": "error",
                    "error_type": type(exc).__name__,
                    "error_message": str(exc),
                }
            )
        export_rows.append(record)

    output_path = output_dir / output_csv_name
    pd.DataFrame(export_rows).to_csv(output_path, index=False)
    return output_path


def _fragment_csv_with_errors(
    input_csv: str | Path,
    smiles_column: str = "smiles",
    id_column: str = "id",
    ring_size: int | None = None,
    max_rows: int | None = None,
) -> tuple[list[FragmentationResult], list[dict]]:
    fragmenter = MacrolactoneFragmenter(ring_size=ring_size)
    rows = _read_csv_rows(
        input_csv=input_csv,
        smiles_column=smiles_column,
        id_column=id_column,
        max_rows=max_rows,
    )

    results: list[FragmentationResult] = []
    errors: list[dict] = []
    for row in rows:
        try:
            results.append(fragmenter.fragment_molecule(row["smiles"], parent_id=row["parent_id"]))
        except MacrolactoneError as exc:
            errors.append(
                {
                    "parent_id": row["parent_id"],
                    "smiles": row["smiles"],
                    "error_type": type(exc).__name__,
                    "error_message": str(exc),
                    "exception": exc,
                }
            )
    return results, errors


def _read_csv_rows(
    input_csv: str | Path,
    smiles_column: str = "smiles",
    id_column: str = "id",
    max_rows: int | None = None,
) -> list[dict]:
    dataframe = pd.read_csv(input_csv)
    if max_rows is not None:
        dataframe = dataframe.head(max_rows)

    rows = []
    for index, row in dataframe.iterrows():
        parent_id = row[id_column] if id_column in dataframe.columns else f"row_{index}"
        rows.append(
            {
                "parent_id": str(parent_id),
                "smiles": row[smiles_column],
            }
        )
    return rows
