/*
 * Copyright Amazon.com, Inc. or its affiliates. All Rights Reserved.
 * SPDX-License-Identifier: Apache-2.0
 */
// Code generated by software.amazon.smithy.rust.codegen.smithy-rs. DO NOT EDIT.

use std::sync::Arc;

use aws_smithy_http_server::Router;
use pyo3::prelude::*;

use crate::{PyHandler, PyHandlers, PyState, SharedSocket};

/// Python compatible wrapper for the [aws_smithy_http_server::Router] type.
#[pyclass(text_signature = "(router)")]
#[derive(Debug, Clone)]
pub struct PyRouter(pub Router);

/// Python application definition, holding the handlers map, the optional Python context object
/// and the asyncio task locals with the running event loop.
#[pyclass(subclass, text_signature = "()")]
#[derive(Debug, Clone)]
pub struct PyApp {
    pub handlers: PyHandlers,
    pub context: Option<Arc<PyObject>>,
    pub locals: pyo3_asyncio::TaskLocals,
}

#[pymethods]
impl PyApp {
    /// Create a new instance of [PyApp].
    #[new]
    pub fn new(py: Python) -> PyResult<Self> {
        let asyncio = py.import("asyncio")?;
        let event_loop = asyncio.call_method0("get_event_loop")?;
        let locals = pyo3_asyncio::TaskLocals::new(event_loop);
        Ok(Self {
            handlers: PyHandlers::new(),
            context: None,
            locals,
        })
    }

    /// Start a single worker with its own Tokio and Python async runtime and provided shared socket.
    ///
    /// Python asynchronous loop needs to be started and handled during the lifetime of the process.
    /// First of all we install [uvloop] as the main Python event loop. Thanks to libuv, uvloop
    /// performs ~20% better than Python standard event loop in most benchmarks, while being 100%
    /// compatible.
    /// We retrieve the Python context object, if setup by the user calling [PyApp::context] method,
    /// generate the [PyState] structure and build the [aws_smithy_http_server::Router], filling
    /// it with the functions generated by `PythonServerOperationHandlerGenerator.kt`.
    /// At last we get a cloned reference to the underlying [socket2::Socket].
    ///
    /// Now that all the setup is done, we can start the two runtimes and run the [hyper] server.
    /// We spawn a thread with a new [tokio::runtime], setup the middlewares and finally block the
    /// thread on `hyper::serve`.
    /// The main process continues and at the end it is blocked on Python `loop.run_forever()`.
    ///
    /// [uvloop]: https://github.com/MagicStack/uvloop
    #[pyo3(text_signature = "($self, socket, worker_number)")]
    fn start_hyper_thread(
        &mut self,
        py: Python,
        socket: &PyCell<SharedSocket>,
        worker_number: isize,
    ) -> PyResult<()> {
        // Setup the Python asyncio loop to use `uvloop`.
        let asyncio = py.import("asyncio")?;
        let uvloop = py.import("uvloop")?;
        uvloop.call_method0("install")?;
        tracing::debug!("Setting up uvloop for current process");
        let event_loop = asyncio.call_method0("new_event_loop")?;
        asyncio.call_method1("set_event_loop", (event_loop,))?;
        // Create the `PyState` object from the Python context object.
        let context = self.context.clone().unwrap_or_else(|| Arc::new(py.None()));
        let state = PyState::new(context);
        // Build the router.
        let router: PyRouter = self.router(py).expect("`start_hyper_thread()` is meant to be called only by subclasses implementing the `router()` method").extract(py)?;
        // Clone the socket.
        let borrow = socket.try_borrow_mut()?;
        let held_socket: &SharedSocket = &*borrow;
        let raw_socket = held_socket.get_socket()?;
        // Store Python event loop locals.
        self.locals = pyo3_asyncio::TaskLocals::new(event_loop);

        // Spawn a new background [std::thread] to run the application.
        tracing::debug!("Start the Tokio runtime in a background task");
        std::thread::spawn(move || {
            // The thread needs a new [tokio] runtime.
            let rt = tokio::runtime::Builder::new_current_thread()
                .enable_all()
                .thread_name(format!("smithy-rs[{}]", worker_number))
                .build()
                .expect("Unable to start a new tokio runtime for this process");
            // Register operations into a Router, add middleware and start the `hyper` server,
            // all inside a [tokio] blocking function.
            tracing::debug!("Add middlewares to Rust Python router");
            let app = router.0.layer(
                tower::ServiceBuilder::new()
                    .layer(aws_smithy_http_server::AddExtensionLayer::new(state)),
            );
            let server = hyper::Server::from_tcp(
                raw_socket
                    .try_into()
                    .expect("Unable to convert `socket2::Socket` into `std::net::TcpListener`"),
            )
            .expect("Unable to create hyper server from shared socket")
            .serve(app.into_make_service());

            tracing::debug!("Starting hyper server from shared socket");
            rt.block_on(async move {
                // Run forever-ish...
                if let Err(err) = server.await {
                    tracing::error!("server error: {}", err);
                }
            });
        });
        // Block on the event loop forever.
        tracing::debug!("Run and block on the Python event loop");
        let event_loop = (*event_loop).call_method0("run_forever");
        tracing::info!("Rust Python server started successfully");
        if event_loop.is_err() {
            tracing::warn!("Ctrl-c handler, quitting");
        }
        Ok(())
    }

    /// Register a new operation in the handlers map.
    ///
    /// The operation registered in the map are used inside the code-generated `router()` method
    /// and passed to the [aws_smithy_http_server::Router] as part of the operation handlers call.
    #[pyo3(text_signature = "($self, name, func)")]
    pub fn register_operation(&mut self, py: Python, name: &str, func: PyObject) -> PyResult<()> {
        let inspect = py.import("inspect")?;
        // Check if the function is a coroutine.
        // NOTE: that `asyncio.iscoroutine()` doesn't work here.
        let is_coroutine = inspect
            .call_method1("iscoroutinefunction", (&func,))?
            .extract::<bool>()?;
        // Find number of expected methods (a Python implementation could not accept the context).
        let func_args = inspect
            .call_method1("getargs", (func.getattr(py, "__code__")?,))?
            .getattr("args")?
            .extract::<Vec<String>>()?;
        let func = PyHandler {
            func,
            is_coroutine,
            args: func_args.len(),
        };
        tracing::info!(
            "Registering {} function `{}` for operation {} with {} arguments",
            if func.is_coroutine { "async" } else { "sync" },
            name,
            func.func,
            func.args
        );
        // Insert the handler in the handlers map.
        self.handlers
            .insert(String::from(name), std::sync::Arc::new(func));
        Ok(())
    }

    /// Register a new context object inside the Rust state.
    #[pyo3(text_signature = "($self, context)")]
    pub fn context(&mut self, _py: Python, context: PyObject) {
        self.context = Some(Arc::new(context));
    }

    /// This method is here because it is meant to be overriden by the code-generated
    /// `App` structure (see PythonServerApplicationGenerator.kt) with the code needed
    /// to build the [aws_smithy_http_server::Router] and register the operations on it.
    #[pyo3(text_signature = "($self)")]
    pub fn router(&self, _py: Python) -> Option<PyObject> {
        None
    }

    /// Main entrypoint: start the server on multiple workers.
    ///
    /// The multiprocessing server is achieved using the ability of a Python interpreter
    /// to clone and start itself as a new process.
    /// The shared sockets is created and Using the [multiprocessing::Process] module, multiple
    /// workers with the method `self.start_single_python_worker()` as target are started.
    ///
    /// [multiprocessing::Process]: https://docs.python.org/3/library/multiprocessing.html
    #[pyo3(text_signature = "($self, address, port, backlog, workers)")]
    fn run(
        &mut self,
        py: Python,
        address: Option<String>,
        port: Option<i32>,
        backlog: Option<i32>,
        workers: Option<usize>,
    ) -> PyResult<()> {
        let mp = py.import("multiprocessing")?;
        mp.call_method0("allow_connection_pickling")?;
        let address = address.unwrap_or_else(|| String::from("127.0.0.1"));
        let port = port.unwrap_or(8080);
        let socket = SharedSocket::new(address, port, backlog)?;
        for idx in 0..workers.unwrap_or_else(num_cpus::get) {
            let sock = socket.try_clone()?;
            let process = mp.getattr("Process")?;
            let handle = process.call1((
                py.None(),
                self.clone().into_py(py).getattr(py, "start_hyper_worker")?,
                format!("smithy-rs[{}]", idx),
                (sock.into_py(py), idx),
            ))?;
            handle.call_method0("start")?;
        }
        Ok(())
    }
}
