#[cfg(feature = "python")]
pub mod python;
mod response;

use std::{collections::HashMap, io::Cursor, net::Ipv4Addr, sync::Arc};

use daft_recordbatch::RecordBatch;
use http_body_util::{combinators::BoxBody, BodyExt, Full};
use hyper::{
    body::{Bytes, Incoming},
    server::conn::http1,
    service::service_fn,
    Method, Request, Response, StatusCode,
};
use hyper_util::rt::TokioIo;
use include_dir::{include_dir, Dir};
use parking_lot::Mutex;
#[cfg(feature = "python")]
use pyo3::prelude::*;
use serde::{Deserialize, Serialize};
use tokio::{net::TcpStream, spawn};
use uuid::Uuid;

type StrRef = Arc<str>;
type Req<T = Incoming> = Request<T>;
type Res = Response<BoxBody<Bytes, std::io::Error>>;
type ServerResult<T> = Result<T, (StatusCode, anyhow::Error)>;

pub const DEFAULT_SERVER_ADDR: Ipv4Addr = Ipv4Addr::LOCALHOST;
pub const DEFAULT_SERVER_PORT: u16 = 3238;

// Global shared dashboard state for this process.
static GLOBAL_DASHBOARD_STATE: Mutex<Option<DashboardState>> = Mutex::new(None);

// $DASHBOARD_ASSETS_DIR is generated by the build script
// and contains all the static files for the dashboard.
static ASSETS_DIR: Dir = include_dir!("$DASHBOARD_ASSETS_DIR");

trait ResultExt<T, E: Into<anyhow::Error>>: Sized {
    fn with_status_code(self, status_code: StatusCode) -> ServerResult<T>;
    fn with_internal_error(self) -> ServerResult<T>;
}

impl<T, E: Into<anyhow::Error>> ResultExt<T, E> for Result<T, E> {
    fn with_status_code(self, status_code: StatusCode) -> ServerResult<T> {
        self.map_err(|err| (status_code, err.into()))
    }

    fn with_internal_error(self) -> ServerResult<T> {
        self.with_status_code(StatusCode::INTERNAL_SERVER_ERROR)
    }
}

#[derive(Deserialize, Serialize, Debug, Clone, PartialEq, Eq)]
#[serde(rename_all = "snake_case", deny_unknown_fields)]
struct QueryInformation {
    id: StrRef,
    unoptimized_plan: Option<StrRef>,
    optimized_plan: Option<StrRef>,
    plan_time_start: StrRef,
    plan_time_end: StrRef,
    logs: Option<StrRef>,
    run_id: Option<StrRef>,
}

#[derive(Serialize)]
struct CellResponse {
    value: String,
    data_type: String,
}

#[derive(Debug)]
struct DashboardState {
    queries: Vec<QueryInformation>,
    dataframe_previews: HashMap<String, RecordBatch>,
    address: String,
    port: u16,
}

impl DashboardState {
    fn new(address: String, port: u16) -> Self {
        Self {
            queries: Default::default(),
            dataframe_previews: Default::default(),
            address,
            port,
        }
    }

    fn queries(&self) -> &[QueryInformation] {
        &self.queries
    }

    fn add_query(&mut self, query_information: QueryInformation) {
        self.queries.push(query_information);
    }

    fn register_dataframe_preview(&mut self, record_batch: RecordBatch) -> String {
        let id = Uuid::new_v4().to_string();
        self.dataframe_previews.insert(id.clone(), record_batch);
        id
    }

    fn get_dataframe_preview(&self, id: &str) -> Option<RecordBatch> {
        self.dataframe_previews.get(id).cloned()
    }

    fn get_url(&self) -> String {
        format!("http://{}:{}", self.address, self.port)
    }

    fn get_queries_url(&self) -> String {
        format!("{}/api/queries", self.get_url())
    }

    fn get_port(&self) -> u16 {
        self.port
    }
}

async fn deserialize<T: for<'de> Deserialize<'de>>(req: Req) -> ServerResult<Req<T>> {
    let (parts, body) = req.into_parts();
    let bytes = body.collect().await.with_internal_error()?.to_bytes();
    let mut cursor = Cursor::new(bytes);
    let body = serde_json::from_reader(&mut cursor).with_status_code(StatusCode::BAD_REQUEST)?;

    Ok(Request::from_parts(parts, body))
}

fn parse_query_params(query: Option<&str>) -> HashMap<String, String> {
    query
        .unwrap_or("")
        .split('&')
        .filter_map(|param| {
            let mut parts = param.splitn(2, '=');
            match (parts.next(), parts.next()) {
                (Some(key), Some(value)) => Some((key.to_string(), value.to_string())),
                _ => None,
            }
        })
        .collect()
}

fn serve_cell_content(row: usize, col: usize, record_batch: &RecordBatch) -> ServerResult<Res> {
    if row >= record_batch.len() || col >= record_batch.num_columns() {
        return Err((
            StatusCode::BAD_REQUEST,
            anyhow::anyhow!("Row or column index out of bounds"),
        ));
    }

    let column = record_batch.get_column(col);
    let cell_html = daft_recordbatch::html_value(column, row, false);

    let response = CellResponse {
        value: cell_html,
        data_type: format!("{:?}", column.data_type()),
    };

    Ok(response::with_body(StatusCode::OK, response))
}

fn generate_interactive_html(
    record_batch: &RecordBatch,
    df_id: &str,
    host: &str,
    port: u16,
) -> String {
    // Start with the basic table HTML from repr_html
    let table_html = record_batch.repr_html();
    // Build the complete interactive HTML with side pane layout
    let mut html = vec![r#"
        <style>
        .dashboard-container {
            display: flex;
            gap: 20px;
            max-width: 100%;
            height: 100%;
        }
        .table-container {
            flex: 1;
            overflow: auto;
        }
        .side-pane {
            width: 35%;
            max-height: 500px;
            border: 1px solid;
            border-radius: 4px;
            padding: 15px;
            display: none;
            overflow: auto;
        }
        .side-pane.visible {
            display: flex;
            flex-direction: column;
        }
        .side-pane-header {
            display: flex;
            justify-content: space-between;
            align-items: center;
            margin-bottom: 10px;
            padding-bottom: 10px;
            border-bottom: 1px solid;
        }
        .side-pane-title {
            font-weight: bold;
        }
        .close-button {
            cursor: pointer;
        }
        .side-pane-content {
            word-wrap: break-word;
            overflow: auto;
        }
        .dataframe td.clickable {
            cursor: pointer;
            transition: background-color 0.2s;
        }
        .dataframe td.clickable:hover {
            opacity: 0.8;
        }
        .dataframe td.clickable.selected {
            opacity: 0.6;
        }
        </style>
        <div class="dashboard-container">
            <div class="table-container">
        "#
    .to_string()];

    // Add the table HTML with ID
    html.push(format!(
        r#"<div id="dataframe-{}">{}</div>"#,
        df_id, table_html
    ));

    // Add the side pane
    html.push(format!(
        r#"
            </div>
            <div class="side-pane" id="side-pane-{df_id}">
                <div class="side-pane-header">
                    <div class="side-pane-title" id="side-pane-title-{df_id}">Cell Details</div>
                    <button class="close-button" id="close-button-{df_id}">×</button>
                </div>
                <div class="side-pane-content" id="side-pane-content-{df_id}">
                    <p style="font-style: italic;">Click on a cell to view its full content</p>
                </div>
            </div>
        </div>
    "#,
    ));

    // Add JavaScript for side pane functionality
    html.push(format!(
        r#"
        <script>
        (function() {{
            const serverUrl = 'http://{}:{}';
            const dfId = '{}';
            const dataframeElement = document.getElementById('dataframe-' + dfId);
            const cells = dataframeElement ? dataframeElement.querySelectorAll('td') : [];
            const sidePane = document.getElementById('side-pane-' + dfId);
            const sidePaneTitle = document.getElementById('side-pane-title-' + dfId);
            const sidePaneContent = document.getElementById('side-pane-content-' + dfId);
            const closeButton = document.getElementById('close-button-' + dfId);
            let selectedCell = null;

            // Function to check if the server is available
            async function isServerAvailable() {{
                try {{
                    const response = await fetch(serverUrl + '/api/ping', {{ method: 'HEAD' }});
                    return response.ok;
                }} catch (e) {{
                    return false;
                }}
            }}

            function closeSidePane(paneId) {{
                const pane = document.getElementById('side-pane-' + paneId);
                if (pane) {{
                    pane.classList.remove('visible');
                    if (selectedCell) {{
                        selectedCell.classList.remove('selected');
                        selectedCell = null;
                    }}
                }}
            }}

            function showSidePane(row, col, content) {{
                sidePaneTitle.textContent = 'Cell (' + row + ', ' + col + ')';
                sidePaneContent.innerHTML = content;
                sidePane.classList.add('visible');
            }}

            function showLoadingContent() {{
                sidePaneContent.innerHTML = '<div style="text-align:center; padding:20px;"><span style="font-style:italic">Loading full content...</span></div>';
            }}

            isServerAvailable().then((available) => {{
                if (!available) {{
                    // Server is unavailable, do not attach any handlers
                    return;
                }}

                // Add event listener for close button
                if (closeButton) {{
                    closeButton.addEventListener('click', function() {{
                        closeSidePane(dfId);
                    }});
                }}

                cells.forEach((cell) => {{
                    // Skip cells that do not have data-row and data-col attributes (e.g., ellipsis row)
                    const rowAttr = cell.getAttribute('data-row');
                    const colAttr = cell.getAttribute('data-col');
                    if (rowAttr === null || colAttr === null) return;

                    const row = parseInt(rowAttr);
                    const col = parseInt(colAttr);
                    cell.classList.add('clickable');

                    cell.onclick = function() {{
                        // Remove selection from previously selected cell
                        if (selectedCell && selectedCell !== cell) {{
                            selectedCell.classList.remove('selected');
                        }}

                        // Toggle selection for current cell
                        if (selectedCell === cell) {{
                            cell.classList.remove('selected');
                            selectedCell = null;
                            closeSidePane(dfId);
                            return;
                        }} else {{
                            cell.classList.add('selected');
                            selectedCell = cell;
                        }}

                        // Show the side pane immediately
                        showSidePane(row, col, '');

                        // Set a timeout to show loading content after 1 second
                        const loadingTimeout = setTimeout(() => {{
                            showLoadingContent();
                        }}, 100);

                        // Fetch the cell content
                        fetch(serverUrl + '/api/dataframes/' + dfId + '/cell?row=' + row + '&col=' + col)
                            .then(response => response.json())
                            .then(data => {{
                                clearTimeout(loadingTimeout);
                                showSidePane(row, col, data.value);
                            }})
                            .catch(err => {{
                                clearTimeout(loadingTimeout);
                                // Get the original cell content from the table
                                const cell = selectedCell;
                                if (cell) {{
                                    const originalContent = cell.innerHTML;
                                    showSidePane(row, col, originalContent);
                                }}
                            }});
                    }};
                }});
            }});
        }})();
        </script>
        "#,
        host, port, df_id
    ));

    html.join("")
}

async fn http_server_application(req: Req) -> ServerResult<Res> {
    let request_path = req.uri().path();
    let paths = request_path
        .split('/')
        .filter(|segment| !segment.is_empty())
        .collect::<Vec<_>>();

    Ok(match (req.method(), paths.as_slice()) {
        (&Method::POST, ["api", "queries"]) => {
            let req = deserialize::<QueryInformation>(req).await?;
            GLOBAL_DASHBOARD_STATE
                .lock()
                .as_mut()
                .expect("Dashboard state should be initialized if server is running")
                .add_query(req.into_body());
            response::empty(StatusCode::OK)
        }
        (&Method::GET, ["api", "queries"]) => {
            let state = GLOBAL_DASHBOARD_STATE.lock();
            let query_informations = state
                .as_ref()
                .expect("Dashboard state should be initialized if server is running")
                .queries();

            response::with_body(StatusCode::OK, query_informations)
        }
        (&Method::GET, ["api", "dataframes", dataframe_id, "cell"]) => {
            let (row, col) = {
                let query = req.uri().query();
                let params = parse_query_params(query);
                (
                    params
                        .get("row")
                        .and_then(|r| r.parse().ok())
                        .ok_or_else(|| {
                            (
                                StatusCode::BAD_REQUEST,
                                anyhow::anyhow!("Invalid row parameter"),
                            )
                        })?,
                    params
                        .get("col")
                        .and_then(|c| c.parse().ok())
                        .ok_or_else(|| {
                            (
                                StatusCode::BAD_REQUEST,
                                anyhow::anyhow!("Invalid col parameter"),
                            )
                        })?,
                )
            };
            let record_batch = {
                let state = GLOBAL_DASHBOARD_STATE.lock();
                state
                    .as_ref()
                    .expect("Dashboard state should be initialized if server is running")
                    .get_dataframe_preview(dataframe_id)
                    .ok_or_else(|| {
                        (
                            StatusCode::NOT_FOUND,
                            anyhow::anyhow!("DataFrame not found: {}", dataframe_id),
                        )
                    })?
            };

            serve_cell_content(row, col, &record_batch)?
        }
        (&Method::HEAD, ["api", "ping"]) | (&Method::GET, ["api", "ping"]) => {
            response::empty(StatusCode::OK)
        }
        (_, ["api", ..]) => response::empty(StatusCode::NOT_FOUND),

        // All other paths (that don't start with "api") will be treated as web-server requests.
        (&Method::GET, _) => {
            let request_path = req.uri().path();
            let path = request_path.trim_start_matches('/');

            let path = if path.is_empty() { "index.html" } else { path };

            // Try to get the file directly
            let file = ASSETS_DIR.get_file(path).or_else(|| {
                // If not found and doesn't end with .html, try with .html extension
                if !std::path::Path::new(path)
                    .extension()
                    .is_some_and(|ext| ext.eq_ignore_ascii_case("html"))
                {
                    ASSETS_DIR.get_file(format!("{}.html", path))
                } else {
                    None
                }
            });

            match file {
                Some(file) => {
                    let content_type = match file.path().extension().and_then(|ext| ext.to_str()) {
                        Some("html") => "text/html",
                        Some("css") => "text/css",
                        Some("js") => "application/javascript",
                        Some("png") => "image/png",
                        Some("jpg") | Some("jpeg") => "image/jpeg",
                        _ => "application/octet-stream",
                    };

                    let bytes = Bytes::copy_from_slice(file.contents());

                    Response::builder()
                        .status(StatusCode::OK)
                        .header("Content-Type", content_type)
                        .header("Access-Control-Allow-Origin", "*")
                        .header("Access-Control-Allow-Methods", "GET, POST")
                        .header("Access-Control-Allow-Headers", "Content-Type")
                        .body(
                            Full::new(bytes)
                                .map_err(|infallible| match infallible {})
                                .boxed(),
                        )
                        .unwrap()
                }
                None => response::empty(StatusCode::NOT_FOUND),
            }
        }
        _ => response::empty(StatusCode::NOT_FOUND),
    })
}

fn handle_stream(stream: TcpStream) {
    let io = TokioIo::new(stream);
    spawn(async move {
        http1::Builder::new()
            .serve_connection(
                io,
                service_fn(move |request| async move {
                    Ok::<_, std::convert::Infallible>(
                        match http_server_application(request).await {
                            Ok(response) => response,
                            Err((status_code, error)) => {
                                response::with_body(status_code, error.to_string())
                            }
                        },
                    )
                }),
            )
            .await
            .unwrap();
    });
}

#[cfg(feature = "python")]
pub fn register_modules(parent: &Bound<PyModule>) -> PyResult<()> {
    const DAFT_DASHBOARD_ENV_ENABLED: &str = "DAFT_DASHBOARD_ENABLED";
    const DAFT_DASHBOARD_ENV_NAME: &str = "DAFT_DASHBOARD";

    let module = PyModule::new(parent.py(), "dashboard")?;
    module.add_wrapped(wrap_pyfunction!(python::launch))?;
    module.add_wrapped(wrap_pyfunction!(python::register_dataframe_for_display))?;
    module.add_wrapped(wrap_pyfunction!(python::generate_interactive_html))?;
    module.add_wrapped(wrap_pyfunction!(python::get_dashboard_url))?;
    module.add_wrapped(wrap_pyfunction!(python::get_dashboard_queries_url))?;
    // module.add_wrapped(wrap_pyfunction!(python::shutdown))?;
    // module.add_wrapped(wrap_pyfunction!(python::cli))?;
    module.add("DAFT_DASHBOARD_ENV_NAME", DAFT_DASHBOARD_ENV_NAME)?;
    module.add("DAFT_DASHBOARD_ENV_ENABLED", DAFT_DASHBOARD_ENV_ENABLED)?;
    parent.add_submodule(&module)?;

    Ok(())
}
