import code
import json
import logging
import sys
from pathlib import Path
from typing import Optional

import click
import pandas as pd
from click import Path as ClickPath

from rpmeta.cli.ctx import Context
from rpmeta.constants import ALL_FEATURES, TARGET, ModelEnum
from rpmeta.dataset import InputRecord
from rpmeta.predictor import Predictor

logger = logging.getLogger(__name__)


@click.group("run")
@click.option(
    "-m",
    "--model-dir",
    type=ClickPath(exists=True, dir_okay=True, resolve_path=True, file_okay=False, path_type=Path),
    required=True,
    help="Path to the model directory",
)
@click.option(
    "-n",
    "--model-name",
    type=click.Choice(ModelEnum, case_sensitive=False),
    required=True,
    help="Type of the model to use",
)
@click.option(
    "-c",
    "--categories",
    type=ClickPath(exists=True, dir_okay=False, resolve_path=True, file_okay=True, path_type=Path),
    required=True,
    help="Path to the categories file",
)
@click.pass_context
def run(ctx: click.Context, model_dir: Path, model_name: ModelEnum, categories: Path):
    """
    Subcommand to collect model-related commands and run operations on them.

    The model dir is expected to be a directory with model files, and the categories are
    expected to be in JSON format. The categories file should contain a mapping of categorical
    features to their possible values.

    The response of the model is a single integer representing the predicted build time duration
    in minutes by default. If the model does not recognize the package, it will return -1 as
    a prediction and log an error message.
    """
    ctx.ensure_object(Context)
    ctx.obj.predictor = Predictor.load(model_dir, model_name, categories, ctx.obj.config)


@run.command("serve")
@click.option("--host", type=str, default=None, help="Host to serve the API on")
@click.option(
    "-p",
    "--port",
    type=int,
    default=None,
    help="Port to serve the API on",
)
@click.option(
    "--debug",
    is_flag=True,
    default=False,
    help="Enable debug mode",
)
@click.pass_context
def serve(ctx: click.Context, host: Optional[str], port: Optional[int], debug: bool):
    """
    Start the API server on specified host and port.

    \b
    The server will accept HTTP POST requests with JSON payloads containing the input data in
    format:
        {
            "package_name": "example-package",
            "epoch": 0,
            "version": "1.0.0",
            "mock_chroot": "fedora-42-x86_64",
            "hw_info": {
                "cpu_model_name": "Intel Xeon Processor (Cascadelake)",
                "cpu_arch": "x86_64",
                "cpu_model": "85",
                "cpu_cores": 6,
                "ram": 15324520,
                "swap": 8388604
            }
        }

    The server will return a JSON response with the predicted build time duration in minutes
    (by default) in the following format:
        {
            "prediction": 21,
            ...
        }
    or -1 as the prediction if the package name is not recognized.
    """
    from uvicorn.config import Config as UvicornConfig
    from uvicorn.server import Server

    from rpmeta.server.api import app, reload_predictor

    reload_predictor(ctx.obj.predictor)

    config = ctx.obj.config

    if host:
        config.api.host = host
    if port:
        config.api.port = port
    if debug:
        config.api.debug = debug

    logger.info("Serving on: %s:%s", config.api.host, config.api.port)

    uvicorn_config = UvicornConfig(
        app=app,
        host=config.api.host,
        port=config.api.port,
        log_level=ctx.obj.log_level.lower() if ctx.obj.log_level else "info",
    )
    server = Server(config=uvicorn_config)
    server.run()


@run.command("predict")
@click.option(
    "-d",
    "--data",
    type=str,
    required=True,
    help="Input data to make prediction on (file path or JSON string)",
)
@click.option(
    "--output-type",
    type=click.Choice(["json", "text"], case_sensitive=False),
    default="text",
    show_default=True,
    help="Output type for the prediction",
)
@click.pass_context
def predict(ctx: click.Context, data: str, output_type: str):
    """
    Make single prediction on the input data.

    \b
    WARNING: The model must be loaded into memory on each query. This mode is extremely
    inefficient for frequent real-time queries.

    The command accepts raw string in JSON format or Linux path to the JSON file in format:
        {
            "package_name": "example-package",
            "epoch": 0,
            "version": "1.0.0",
            "mock_chroot": "fedora-42-x86_64",
            "hw_info": {
                "cpu_model_name": "AMD Ryzen 7 PRO 7840HS",
                "cpu_arch": "x86_64",
                "cpu_model": "116",
                "cpu_cores": 8,
                "ram": 123456789,
                "swap": 123456789
            }
        }

    Command response is in minutes by default.
    """
    if Path(data).exists():
        with open(data) as f:
            input_data = json.load(f)
    else:
        input_data = json.loads(data)

    logger.debug("Input data received: %s", input_data)

    prediction = ctx.obj.predictor.predict(InputRecord(**input_data), ctx.obj.config.model.behavior)

    if output_type == "json":
        print(json.dumps({"prediction": prediction}))
    else:
        print(f"Prediction: {prediction}")


@run.command("visualize")
@click.option(
    "-d",
    "--dataset",
    type=ClickPath(exists=True, dir_okay=False, resolve_path=True, file_okay=True, path_type=Path),
    required=True,
    help="Path to the dataset file used for testing the model",
)
@click.pass_context
def visualize(ctx: click.Context, dataset: Path):
    """
    Visualize the model performance.
    This command will generate various plots and save them to the model directory.
    """
    from sklearn.metrics import mean_absolute_error, r2_score, root_mean_squared_error

    from rpmeta.trainer.base import BestModelResult
    from rpmeta.trainer.visualizer import ResultsHandler

    predictor = ctx.obj.predictor

    with open(dataset, encoding="utf-8") as f:
        data = json.load(f)

    for record in data:
        record["mock_chroot"] = f"{record['os']}-{record['os_version']}-{record['os_arch']}"

    df = pd.json_normalize(data)

    for col, cat_list in predictor.category_maps.items():
        dtype = pd.CategoricalDtype(categories=cat_list, ordered=False)
        df[col] = df[col].astype(dtype)

    df_actuals = df[TARGET]
    df_predictions = predictor.model.predict(df[ALL_FEATURES])

    X_test = df[ALL_FEATURES]  # noqa: N806
    y_test = df[TARGET]

    r2 = r2_score(df_actuals, df_predictions)
    mae = mean_absolute_error(df_actuals, df_predictions)
    rmse = root_mean_squared_error(df_actuals, df_predictions)

    best_model_result = BestModelResult(
        model_name="testing_model",
        model_path=Path("."),  # Not applicable for loaded models
        neg_rmse=-rmse,
        neg_mae=-mae,
        r2=r2,
        params={},  # Not available when loading a model
        model=predictor.model,
    )

    best_models = {"testing_model": best_model_result}
    result_handler = ResultsHandler(
        all_trials={},
        best_models=best_models,
        studies={},
        X_test=X_test,
        y_test=y_test,
        config=ctx.obj.config,
    )

    result_handler.plot_predictions()
    result_handler.plot_distribution()
    result_handler.save_best_json()


@run.command("shell")
@click.pass_context
def shell(ctx: click.Context):
    """
    Start interactive Python shell with the model loaded.
    """
    predictor = ctx.obj.predictor
    config = ctx.obj.config

    local_vars = {
        "ctx": ctx,
        "predictor": predictor,
        "config": config,
        "InputRecord": InputRecord,
        "pd": pd,
    }

    var_list = "\n".join(f" - {name}: {var!s}" for name, var in local_vars.items())
    banner = (
        "Interactive RPMeta Shell\n"
        f"Python {sys.version}\n"
        "The following variables are available:\n"
        f"{var_list}"
    )

    code.interact(banner=banner, local=local_vars)
