mod error;
mod ticket;

use std::pin::Pin;
use std::sync::Arc;

use arrow::datatypes::Schema;
use arrow::ipc::writer::IpcWriteOptions;
use arrow_flight::encode::FlightDataEncoderBuilder;
use arrow_flight::flight_service_server::{FlightService, FlightServiceServer};
use arrow_flight::{
    Action, ActionType, Criteria, Empty, FlightData, FlightDescriptor, FlightInfo,
    HandshakeRequest, HandshakeResponse, PollInfo, PutResult, SchemaAsIpc, SchemaResult, Ticket,
};
use futures::stream::{self, Stream, StreamExt};
use tonic::transport::Server;
use tonic::{Request, Response, Status, Streaming};

use crate::core::MurrError;
use crate::service::MurrService;
use log::info;
use ticket::FetchTicket;

pub struct MurrFlightService {
    service: Arc<MurrService>,
}

impl MurrFlightService {
    pub fn new(service: Arc<MurrService>) -> Self {
        Self { service }
    }

    pub async fn serve(self) -> Result<(), MurrError> {
        let addr = self
            .service
            .config()
            .server
            .grpc
            .addr()
            .parse()
            .map_err(|e| MurrError::ConfigParsingError(format!("invalid address: {e}")))?;
        info!("Listening for Flight/gRPC requests on {addr}");
        Server::builder()
            .tcp_nodelay(true)
            .add_service(FlightServiceServer::new(self))
            .serve(addr)
            .await
            .map_err(|e| MurrError::IoError(format!("Flight server error: {e}")))?;

        Ok(())
    }
}

type BoxStream<T> = Pin<Box<dyn Stream<Item = Result<T, Status>> + Send>>;

#[tonic::async_trait]
impl FlightService for MurrFlightService {
    type HandshakeStream = BoxStream<HandshakeResponse>;
    type ListFlightsStream = BoxStream<FlightInfo>;
    type DoGetStream = BoxStream<FlightData>;
    type DoPutStream = BoxStream<PutResult>;
    type DoExchangeStream = BoxStream<FlightData>;
    type DoActionStream = BoxStream<arrow_flight::Result>;
    type ListActionsStream = BoxStream<ActionType>;

    async fn do_get(
        &self,
        request: Request<Ticket>,
    ) -> Result<Response<Self::DoGetStream>, Status> {
        let ticket = request.into_inner();
        let fetch: FetchTicket = serde_json::from_slice(&ticket.ticket)
            .map_err(|e| Status::invalid_argument(format!("invalid ticket JSON: {e}")))?;

        let service = self.service.clone();
        let batch = tokio::task::spawn_blocking(move || {
            let keys: Vec<&str> = fetch.keys.iter().map(String::as_str).collect();
            let columns: Vec<&str> = fetch.columns.iter().map(String::as_str).collect();
            service.read(&fetch.table, &keys, &columns)
        })
        .await
        .map_err(join_to_status)?
        .map_err(Status::from)?;

        let stream = FlightDataEncoderBuilder::new()
            .build(stream::once(async { Ok(batch) }))
            .map(|result| result.map_err(|e| e.into()));

        Ok(Response::new(Box::pin(stream)))
    }

    async fn get_flight_info(
        &self,
        request: Request<FlightDescriptor>,
    ) -> Result<Response<FlightInfo>, Status> {
        let descriptor = request.into_inner();
        let table_name = descriptor
            .path
            .first()
            .ok_or_else(|| Status::invalid_argument("path must contain table name"))?
            .clone();

        let service = self.service.clone();
        let schema = tokio::task::spawn_blocking(move || service.get_schema(&table_name))
            .await
            .map_err(join_to_status)?
            .map_err(Status::from)?;
        let arrow_schema: Schema = (&schema).into();

        let info = FlightInfo::new()
            .try_with_schema(&arrow_schema)
            .map_err(|e| Status::internal(e.to_string()))?
            .with_descriptor(descriptor);

        Ok(Response::new(info))
    }

    async fn get_schema(
        &self,
        request: Request<FlightDescriptor>,
    ) -> Result<Response<SchemaResult>, Status> {
        let descriptor = request.into_inner();
        let table_name = descriptor
            .path
            .first()
            .ok_or_else(|| Status::invalid_argument("path must contain table name"))?
            .clone();

        let service = self.service.clone();
        let schema = tokio::task::spawn_blocking(move || service.get_schema(&table_name))
            .await
            .map_err(join_to_status)?
            .map_err(Status::from)?;
        let arrow_schema: Schema = (&schema).into();
        let options = IpcWriteOptions::default();

        let result = SchemaResult::try_from(SchemaAsIpc::new(&arrow_schema, &options))
            .map_err(|e| Status::internal(e.to_string()))?;

        Ok(Response::new(result))
    }

    async fn list_flights(
        &self,
        _request: Request<Criteria>,
    ) -> Result<Response<Self::ListFlightsStream>, Status> {
        let service = self.service.clone();
        let tables = tokio::task::spawn_blocking(move || service.list_tables())
            .await
            .map_err(join_to_status)?;
        let infos: Vec<Result<FlightInfo, Status>> = tables
            .into_iter()
            .map(|(name, schema)| {
                let arrow_schema: Schema = (&schema).into();
                let descriptor = FlightDescriptor::new_path(vec![name]);
                FlightInfo::new()
                    .try_with_schema(&arrow_schema)
                    .map(|info| info.with_descriptor(descriptor))
                    .map_err(|e| Status::internal(e.to_string()))
            })
            .collect();

        Ok(Response::new(Box::pin(stream::iter(infos))))
    }

    async fn handshake(
        &self,
        _request: Request<Streaming<HandshakeRequest>>,
    ) -> Result<Response<Self::HandshakeStream>, Status> {
        Err(Status::unimplemented("handshake not supported"))
    }

    async fn poll_flight_info(
        &self,
        _request: Request<FlightDescriptor>,
    ) -> Result<Response<PollInfo>, Status> {
        Err(Status::unimplemented("poll_flight_info not supported"))
    }

    async fn do_put(
        &self,
        _request: Request<Streaming<FlightData>>,
    ) -> Result<Response<Self::DoPutStream>, Status> {
        Err(Status::unimplemented("do_put not supported"))
    }

    async fn do_exchange(
        &self,
        _request: Request<Streaming<FlightData>>,
    ) -> Result<Response<Self::DoExchangeStream>, Status> {
        Err(Status::unimplemented("do_exchange not supported"))
    }

    async fn do_action(
        &self,
        _request: Request<Action>,
    ) -> Result<Response<Self::DoActionStream>, Status> {
        Err(Status::unimplemented("do_action not supported"))
    }

    async fn list_actions(
        &self,
        _request: Request<Empty>,
    ) -> Result<Response<Self::ListActionsStream>, Status> {
        Err(Status::unimplemented("list_actions not supported"))
    }
}

fn join_to_status(e: tokio::task::JoinError) -> Status {
    Status::internal(format!("blocking task failed: {e}"))
}
