from __future__ import annotations

import argparse
import os
from pathlib import Path

import mntools
from mntools.load.loader_sbml import load_from_sbml

from aam_annotator.generate_mapped_network import generate_mapped_network


def parse_args() -> argparse.Namespace:
    parser = argparse.ArgumentParser(
        description="Generate an atom-level annotated metabolic network from an input model file."
    )
    parser.add_argument(
        "input_filename",
        type=Path,
        help="Input model file (SBML .xml).",
    )
    parser.add_argument(
        "-o",
        "--output",
        type=Path,
        default=None,
        help="Optional output filename for the generated annotated JSON.",
    )
    parser.add_argument(
        "-c",
        "--curations",
        type=Path,
        default=None,
        help=("Optional curations JSON filename."),
    )
    parser.add_argument(
        "-db",
        "--database",
        type=str,
        default="public",
        help=(
            "Optionally specify the ChemRecon data source. Set to 'public' to use the public database,"
            "or 'local' to use a local docker image (see ChemRecon documentation)."
        ),
    )
    return parser.parse_args()


def load_network(input_filename: Path) -> mntools.MetabolicNetwork:

    suffix = input_filename.suffix.lower()
    if suffix not in {".xml", ".sbml"}:
        raise ValueError(
            f"Unsupported input format '{suffix}'. "
            "Expected an SBML file with extension .xml or .sbml."
        )

    return load_from_sbml(str(input_filename))


def main() -> None:
    args = parse_args()

    input_filename = args.input_filename.resolve()
    output_filename = args.output.resolve() if args.output is not None else None
    curations_file = args.curations.resolve() if args.curations is not None else None
    db_selection = args.database if args.database is not None else "public"

    if not input_filename.exists():
        raise FileNotFoundError(f"Input file not found: {input_filename}")

    # Import
    import chemrecon

    # Connect to local ChemRecon backend, like demo_run.py.
    match db_selection:
        case "public":
            chemrecon.connect_public()
        case "local":
            chemrecon.connect_local_docker()
        case _:
            raise ValueError("Invalid argument for the database connection.")

    print(f"\nCWD: {os.getcwd()}")
    print("\nParsing CRN")
    network = load_network(input_filename)

    print("\nGenerating ACRN")

    generate_mapped_network(
        network,
        output_filename=output_filename,
        curations_path=curations_file,
    )


if __name__ == "__main__":
    main()
