use axum::{
    extract::{ConnectInfo, Extension},
    response::Html,
    routing::{
        delete as axum_delete, get as axum_get, head as axum_head, options as axum_options,
        patch as axum_patch, post as axum_post, put as axum_put,
    },
    Json, Router,
};
use dashmap::DashMap;
use once_cell::sync::Lazy;
use pyo3::prelude::*;
use pyo3::types::{PyAny, PyCFunction, PyDict, PyModule, PyTuple};
use smallvec::SmallVec;
use std::net::SocketAddr;
use std::sync::Arc;
use tokio::net::TcpListener;
use tracing::{debug, error, info, warn, Level};
use tracing_subscriber;
mod openapi;
use openapi::build_openapi_spec;
mod py_handlers;
mod pydantic;
mod utils;

use crate::py_handlers::{run_py_handler_no_args, run_py_handler_with_args};
use crate::pydantic::{is_pydantic_model, register_pydantic_integration};

const SWAGGER_HTML: &str = include_str!("../static/swagger-ui.html");

#[derive(Clone)]
pub struct RouteHandler {
    pub func: Py<PyAny>,
    pub param_validators: Vec<(String, Py<PyAny>)>,
}

pub static ROUTES: Lazy<DashMap<String, RouteHandler>> = Lazy::new(DashMap::new);

static PYTHON_RUNTIME: Lazy<tokio::runtime::Runtime> = Lazy::new(|| {
    let worker_threads = num_cpus::get().max(4).min(16);
    tokio::runtime::Builder::new_multi_thread()
        .worker_threads(worker_threads)
        .thread_name("python-handler")
        .enable_all()
        .build()
        .expect("Failed to create Python runtime")
});

#[derive(Clone)]
struct AppState {
    rt_handle: tokio::runtime::Handle,
}

#[pyclass]
pub struct FastrAPI {
    #[allow(dead_code)]
    router: Arc<DashMap<String, RouteHandler>>,
}

use smartstring::alias::String as SmartString;

// Helper function to parse annotations once
fn parse_param_validators(py: Python, func: &Bound<PyAny>) -> Vec<(String, Py<PyAny>)> {
    let mut validators = Vec::new();

    if let Ok(annotations) = func.getattr("__annotations__") {
        if let Ok(ann_dict) = annotations.cast::<PyDict>() {
            for (key, value) in ann_dict.iter() {
                if let Ok(param_name) = key.extract::<String>() {
                    if param_name != "return" && is_pydantic_model(py, &value) {
                        validators.push((param_name, value.unbind()));
                    }
                }
            }
        }
    }

    validators
}

#[pymethods]
impl FastrAPI {
    #[new]
    fn new() -> Self {
        FastrAPI {
            router: Arc::new(DashMap::new()),
        }
    }

    fn register_route(&self, path: String, func: Py<PyAny>, method: Option<String>) {
        Python::attach(|py| {
            let func_bound = func.bind(py);

            // Parse validators once at registration time
            let param_validators = parse_param_validators(py, func_bound);

            let handler = RouteHandler {
                func: func.clone_ref(py),
                param_validators: param_validators.clone(),
            };

            let method = method.unwrap_or_else(|| "GET".to_string()).to_uppercase();
            let mut key = SmartString::new();
            key.push_str(&method);
            key.push(' ');
            key.push_str(&path);

            let has_validation = !param_validators.is_empty();
            ROUTES.insert((&key).to_string(), handler);
            info!(
                "✅ Registered route [{}] (validators: {})",
                key,
                param_validators.len()
            );
        });
    }

    fn get<'py>(&self, path: String, py: Python<'py>) -> PyResult<Py<PyAny>> {
        self.create_decorator("GET", path, py)
    }

    fn post<'py>(&self, path: String, py: Python<'py>) -> PyResult<Py<PyAny>> {
        self.create_decorator("POST", path, py)
    }

    fn put<'py>(&self, path: String, py: Python<'py>) -> PyResult<Py<PyAny>> {
        self.create_decorator("PUT", path, py)
    }

    fn delete<'py>(&self, path: String, py: Python<'py>) -> PyResult<Py<PyAny>> {
        self.create_decorator("DELETE", path, py)
    }

    fn patch<'py>(&self, path: String, py: Python<'py>) -> PyResult<Py<PyAny>> {
        self.create_decorator("PATCH", path, py)
    }

    fn options<'py>(&self, path: String, py: Python<'py>) -> PyResult<Py<PyAny>> {
        self.create_decorator("OPTIONS", path, py)
    }

    fn head<'py>(&self, path: String, py: Python<'py>) -> PyResult<Py<PyAny>> {
        self.create_decorator("HEAD", path, py)
    }

    fn create_decorator<'py>(
        &self,
        method: &str,
        path: String,
        py: Python<'py>,
    ) -> PyResult<Py<PyAny>> {
        let route_key = format!("{} {}", method, path);

        let decorator = move |args: &Bound<'_, PyTuple>,
                              _kwargs: Option<&Bound<'_, PyDict>>|
              -> PyResult<Py<PyAny>> {
            let py = args.py();
            let func: Py<PyAny> = args.get_item(0)?.extract()?;
            let func_bound = func.bind(py);

            // Parse validators once at decoration time
            let param_validators = parse_param_validators(py, func_bound);

            let handler = RouteHandler {
                func: func.clone_ref(py),
                param_validators: param_validators.clone(),
            };

            ROUTES.insert(route_key.clone(), handler);
            info!(
                "🧩 Added decorated route [{}] (validators: {})",
                route_key,
                param_validators.len()
            );

            Ok(func)
        };

        PyCFunction::new_closure(py, None, None, decorator).map(|f| f.into())
    }

    fn serve(&self, py: Python, host: Option<String>, port: Option<u16>) -> PyResult<()> {
        tracing_subscriber::fmt()
            .with_max_level(Level::DEBUG)
            .with_target(false)
            .init();

        info!("🚀 Starting FastrAPI...");

        let host = host.unwrap_or_else(|| "127.0.0.1".to_string());
        let port = port.unwrap_or(8000);

        let rt_handle = PYTHON_RUNTIME.handle().clone();
        let app_state = AppState { rt_handle };

        let mut app = Router::new();

        println!("🧩 Registered routes:");
        for key in ROUTES.iter() {
            println!("   • {}", key.key());
        }

        // Build OpenAPI spec from routes
        let openapi_spec = build_openapi_spec(py, &ROUTES);
        let openapi_json = Arc::new(
            serde_json::to_value(&openapi_spec).expect("Failed to serialize OpenAPI spec"),
        );

        // Serve OpenAPI JSON manually
        let openapi_json_clone = Arc::clone(&openapi_json);
        app = app.route(
            "/api-docs/openapi.json",
            axum_get(move || {
                let json = Arc::clone(&openapi_json_clone);
                async move { Json(json.as_ref().clone()) }
            }),
        );
        app = app.route("/docs", axum_get(|| async { Html(SWAGGER_HTML) }));
        // Use Arc<str> to avoid cloning on every request
        for entry in ROUTES.iter() {
            let route_key: Arc<str> = entry.key().clone().into();
            let parts: SmallVec<[&str; 2]> = route_key.splitn(2, ' ').collect();

            if parts.len() != 2 {
                warn!("⚠️ Invalid route key format: {}", route_key);
                continue;
            }

            let method = parts[0];
            let path = parts[1].to_string();

            debug!("🔧 Building route: [{} {}]", method, path);

            match method {
                "GET" | "HEAD" | "OPTIONS" => {
                    let route_key = Arc::clone(&route_key);
                    let handler_fn =
                        move |Extension(state): Extension<AppState>,
                              ConnectInfo(_addr): ConnectInfo<SocketAddr>| {
                            let route_key = Arc::clone(&route_key);
                            async move {
                                #[cfg(feature = "verbose-logging")]
                                debug!("{} from {}", route_key, _addr);
                                run_py_handler_no_args(state.rt_handle, route_key).await
                            }
                        };

                    app = match method {
                        "GET" => app.route(&path, axum_get(handler_fn)),
                        "HEAD" => app.route(&path, axum_head(handler_fn)),
                        "OPTIONS" => app.route(&path, axum_options(handler_fn)),
                        _ => app,
                    };
                }
                "POST" | "PUT" | "DELETE" | "PATCH" => {
                    let route_key = Arc::clone(&route_key);
                    let handler_fn =
                        move |Extension(state): Extension<AppState>,
                              ConnectInfo(_addr): ConnectInfo<SocketAddr>,
                              Json(payload): Json<serde_json::Value>| {
                            let route_key = Arc::clone(&route_key);
                            async move {
                                #[cfg(feature = "verbose-logging")]
                                debug!("📥 {} from {}", route_key, _addr);
                                run_py_handler_with_args(state.rt_handle, route_key, payload).await
                            }
                        };

                    app = match method {
                        "POST" => app.route(&path, axum_post(handler_fn)),
                        "PUT" => app.route(&path, axum_put(handler_fn)),
                        "DELETE" => app.route(&path, axum_delete(handler_fn)),
                        "PATCH" => app.route(&path, axum_patch(handler_fn)),
                        _ => app,
                    };
                }
                _ => warn!("⚠️ Ignoring unknown HTTP method: {}", method),
            }
        }

        app = app.layer(axum::Extension(app_state));

        // Add Swagger UI pointing to the JSON endpoint
        // app = app.merge(SwaggerUi::new("/docs").url("/api-docs/openapi.json", openapi_spec));

        py.detach(move || {
            PYTHON_RUNTIME.block_on(async move {
                let addr = format!("{}:{}", host, port);

                let listener = match TcpListener::bind(&addr).await {
                    Ok(l) => l,
                    Err(e) => {
                        error!("Failed to bind to {}: {}", addr, e);
                        return;
                    }
                };

                info!("🚀 FastrAPI running at http://{}", addr);
                info!("📚 Swagger UI available at http://{}/docs", addr);
                info!("📄 OpenAPI spec at http://{}/api-docs/openapi.json", addr);

                if let Err(e) = axum::serve(
                    listener,
                    app.into_make_service_with_connect_info::<SocketAddr>(),
                )
                .await
                {
                    error!("Server error: {}", e);
                }
            });
        });

        Ok(())
    }
}

#[pyfunction]
fn get_decorator(func: Py<PyAny>, path: String) -> PyResult<()> {
    Python::attach(|py| {
        let func_bound = func.bind(py);

        // Parse validators once at registration time
        let param_validators = parse_param_validators(py, func_bound);

        let handler = RouteHandler {
            func: func.clone_ref(py),
            param_validators: param_validators.clone(),
        };

        let key = format!("GET {}", path);
        ROUTES.insert(key.clone(), handler);
        info!(
            "🔗 Registered via get_decorator [{}] (validators: {})",
            key,
            param_validators.len()
        );
    });
    Ok(())
}

#[pymodule(gil_used = false)]
fn fastrapi(m: &Bound<'_, PyModule>) -> PyResult<()> {
    m.add_class::<FastrAPI>()?;
    m.add_function(wrap_pyfunction!(get_decorator, m)?)?;
    register_pydantic_integration(m)?;
    Ok(())
}
