import importlib.util
import logging
from pathlib import Path

import numpy as np
import numpy.typing as npt
import torch
import typer

from where_fi.application import CSIApplication
from where_fi.processing.aoa import AoA

from . import file, globals

cli = typer.Typer(callback=globals.main)
logger = logging.getLogger(__name__)
device = torch.device("cuda" if torch.cuda.is_available() else "cpu")


@cli.command()
def antennas() -> None:
    """Utility to help determine the order in which antennas are plugged in

    Once the script is running, unplug and replug antennas from left to right, to get
    the correct order. Every time an antenna is unplugged and replugged, the script will
    print the antenna identifier. When you are done, press Ctrl+C to stop the script and
    get the final order.
    """
    from ..utils import antenna_order

    if not globals.csi_producer.is_live:
        raise NotImplementedError("This command only works with live data")

    antenna_order.main()


@cli.command()
def heatmap() -> None:
    app = CSIApplication(globals.csi_producer, visualise_raw=True)
    aoa = AoA()

    @app.on_sample
    def _(sample: npt.NDArray[np.complex64]) -> None:
        processed_tensor = torch.tensor(sample, device=device)
        aoa.update(processed_tensor)
        aoa.heatmap(visualiser=app.visualise_data)

    app.start()
    logger.info("Finished processing CSI data")


@cli.command()
def run(file: Path, app_name: str = "app") -> None:
    """
    Run the application with the given file.

    Can be used for running arbitrary CSI applications with non-default data streams
    (e.g. from file or environment simulation).
    """
    print(file.absolute())
    if not file.exists():
        print(f"Error: File '{file}' does not exist.")
        raise typer.Exit(1)
    module_name = file.stem

    spec = importlib.util.spec_from_file_location(module_name, str(file))
    if spec is None:
        print(f"Could not load spec from {file}")
        raise typer.Exit(1)

    module = importlib.util.module_from_spec(spec)
    if spec.loader is None:
        print(f"Could not load module from {file}")
        raise typer.Exit(1)
    try:
        spec.loader.exec_module(module)
    except Exception as e:
        print(f"Failed to execute {file}: {e}")
        raise typer.Exit(1) from e

    if not hasattr(module, app_name):
        print(f"Error: '{app_name}' not defined in the module.")
        raise typer.Exit(1)
    if not isinstance(module.app, CSIApplication):
        print(f"Error: '{app_name}' is not a CSIApplication.")
        raise typer.Exit(1)

    module.app.producer = globals.csi_producer
    module.app.start()


cli.add_typer(file.app, name="file", help="Commands for working with CSI files")