use std::borrow::BorrowMut;
use std::cell::RefCell;
use std::io;

use ents::{
    DatabaseError, Edge, EdgeDraft, EdgeQuery, EdgeQueryResult, EdgeValue, Ent,
    Id, IncomingEdgeProvider, QueryEdge, ReadEnt, SortOrder,
    TransactionProvider, Transactional,
};
use ents_admin::AdminEnt;
use postgres::types::ToSql;
use postgres::{Config, GenericClient, NoTls, Transaction};
use r2d2::Pool;
use r2d2_postgres::PostgresConnectionManager;

const MAX_EDGES: usize = 100;

pub const SCHEMA_SQL: &str = r#"
CREATE TABLE IF NOT EXISTS entities (
   id BIGINT GENERATED BY DEFAULT AS IDENTITY PRIMARY KEY,
   type TEXT NOT NULL,
   data TEXT NOT NULL
);
CREATE TABLE IF NOT EXISTS edges (
   source BIGINT NOT NULL,
   type BYTEA NOT NULL,
   dest BIGINT NOT NULL,
   PRIMARY KEY (source, type, dest)
);
CREATE INDEX IF NOT EXISTS edges_dest_idx ON edges (dest, source, type);
"#;

fn other_error(message: impl Into<String>) -> DatabaseError {
    DatabaseError::Other {
        source: Box::new(io::Error::other(message.into())),
    }
}

fn id_to_i64(id: Id) -> Result<i64, DatabaseError> {
    i64::try_from(id)
        .map_err(|_| other_error(format!("id {id} exceeds BIGINT range")))
}

fn i64_to_id(id: i64) -> Result<Id, DatabaseError> {
    u64::try_from(id).map_err(|_| other_error(format!("id {id} is negative")))
}

fn usize_to_i64(value: usize, name: &str) -> Result<i64, DatabaseError> {
    i64::try_from(value)
        .map_err(|_| other_error(format!("{name} exceeds BIGINT range")))
}

fn parse_ent_row(
    id: i64,
    data_json: &str,
) -> Result<Box<dyn Ent>, DatabaseError> {
    let mut ent =
        serde_json::from_str::<Box<dyn Ent>>(data_json).map_err(|e| {
            DatabaseError::Other {
                source: Box::new(e),
            }
        })?;
    ent.set_id(i64_to_id(id)?);
    Ok(ent)
}

pub fn init_schema<C>(client: &mut C) -> Result<(), DatabaseError>
where
    C: GenericClient,
{
    client
        .batch_execute(SCHEMA_SQL)
        .map_err(|e| DatabaseError::Other {
            source: Box::new(e),
        })
}

pub type PostgresPool = Pool<PostgresConnectionManager<NoTls>>;

pub struct Txn<'conn> {
    tx: RefCell<Transaction<'conn>>,
}

impl<'conn> Txn<'conn> {
    pub fn new(tx: Transaction<'conn>) -> Self {
        Self {
            tx: RefCell::new(tx),
        }
    }

    fn insert<E: Ent>(&self, ent: &E) -> Result<Id, DatabaseError> {
        let entity_type = ent.typetag_name().to_string();
        let data_json =
            serde_json::to_string(&(ent as &dyn Ent)).map_err(|e| {
                DatabaseError::Other {
                    source: Box::new(e),
                }
            })?;

        let row = self
            .tx
            .borrow_mut()
            .query_one(
                "INSERT INTO entities (type, data) VALUES ($1, $2) RETURNING id",
                &[&entity_type, &data_json],
            )
            .map_err(|e| DatabaseError::Other {
                source: Box::new(e),
            })?;

        i64_to_id(row.get::<_, i64>(0))
    }

    fn update_internal(
        &self,
        id: Id,
        ent: Box<dyn Ent>,
        expected_last_updated: Option<u64>,
    ) -> Result<bool, DatabaseError> {
        let entity_type = ent.typetag_name().to_string();
        let data_json =
            serde_json::to_string(&ent).map_err(|e| DatabaseError::Other {
                source: Box::new(e),
            })?;
        let id = id_to_i64(id)?;
        let expected_last_updated =
            expected_last_updated.map(id_to_i64).transpose()?;

        let rows_affected = self
            .tx
            .borrow_mut()
            .execute(
                r#"
                UPDATE entities
                SET data = $1, type = $2
                WHERE
                    id = $3 AND
                    (
                        $4::BIGINT IS NULL OR
                        COALESCE((data::jsonb ->> 'last_updated')::BIGINT, -1) = $4
                    )
                "#,
                &[&data_json, &entity_type, &id, &expected_last_updated],
            )
            .map_err(|e| DatabaseError::Other {
                source: Box::new(e),
            })?;

        Ok(rows_affected > 0)
    }
}

impl<'conn> ReadEnt for Txn<'conn> {
    fn get(&self, id: Id) -> Result<Option<Box<dyn Ent>>, DatabaseError> {
        let row = self
            .tx
            .borrow_mut()
            .query_opt(
                "SELECT id, data FROM entities WHERE id = $1",
                &[&id_to_i64(id)?],
            )
            .map_err(|e| DatabaseError::Other {
                source: Box::new(e),
            })?;

        match row {
            Some(row) => {
                let id = row.get::<_, i64>(0);
                let data_json = row.get::<_, String>(1);
                Ok(Some(parse_ent_row(id, &data_json)?))
            }
            None => Ok(None),
        }
    }
}

impl<'conn> Transactional for Txn<'conn> {
    fn create<E: Ent>(&self, mut ent: E) -> Result<Id, DatabaseError> {
        let id = self.insert(&ent)?;
        ent.set_id(id);
        ent.setup_edges(self).map_err(|e| DatabaseError::Other {
            source: Box::new(e),
        })?;
        Ok(id)
    }

    fn delete(&self, id: Id) -> Result<(), DatabaseError> {
        let id = id_to_i64(id)?;

        self.tx
            .borrow_mut()
            .execute("DELETE FROM edges WHERE dest = $1", &[&id])
            .map_err(|e| DatabaseError::Other {
                source: Box::new(e),
            })?;

        self.tx
            .borrow_mut()
            .execute("DELETE FROM entities WHERE id = $1", &[&id])
            .map_err(|e| DatabaseError::Other {
                source: Box::new(e),
            })?;

        Ok(())
    }

    fn create_edge(&self, edge: EdgeValue) -> Result<(), DatabaseError> {
        let source = id_to_i64(edge.source)?;
        let dest = id_to_i64(edge.dest)?;

        self.tx
            .borrow_mut()
            .execute(
                "INSERT INTO edges (source, type, dest) VALUES ($1, $2, $3)",
                &[&source, &edge.sort_key, &dest],
            )
            .map_err(|e| DatabaseError::Other {
                source: Box::new(e),
            })?;

        Ok(())
    }

    fn update<T: Ent, F: FnOnce(&mut T), B: BorrowMut<T>>(
        &self,
        mut ent0: B,
        mutator: F,
    ) -> Result<bool, DatabaseError> {
        let ent = ent0.borrow_mut();
        let draft0 = T::EdgeProvider::draft(ent);
        let ent_id = ent.id();
        let expected_last_updated = ent.last_updated();

        mutator(ent);
        ent.mark_updated().map_err(|e| DatabaseError::Other {
            source: Box::new(e),
        })?;

        let draft1 = T::EdgeProvider::draft(ent);

        if draft0 == draft1 {
            return self.update_internal(
                ent.id(),
                dyn_clone::clone_box(ent as &dyn Ent),
                Some(expected_last_updated),
            );
        }

        let edge0 = draft0
            .check(self)
            .map(|edges| {
                edges
                    .into_iter()
                    .map(|edge| edge.with_dest(ent_id))
                    .collect::<Vec<_>>()
            })
            .map_err(|e| DatabaseError::Other {
                source: Box::new(e),
            })?;
        let edge1 = draft1
            .check(self)
            .map(|edges| {
                edges
                    .into_iter()
                    .map(|edge| edge.with_dest(ent_id))
                    .collect::<Vec<_>>()
            })
            .map_err(|e| DatabaseError::Other {
                source: Box::new(e),
            })?;

        let updated = self.update_internal(
            ent.id(),
            dyn_clone::clone_box(ent as &dyn Ent),
            Some(expected_last_updated),
        )?;

        if updated {
            for edge in edge0 {
                let source = id_to_i64(edge.source)?;
                let dest = id_to_i64(edge.dest)?;
                self.tx
                    .borrow_mut()
                    .execute(
                        "DELETE FROM edges WHERE source = $1 AND type = $2 AND dest = $3",
                        &[&source, &edge.sort_key, &dest],
                    )
                    .map_err(|e| DatabaseError::Other {
                        source: Box::new(e),
                    })?;
            }

            for edge in edge1 {
                self.create_edge(edge)?;
            }
        }

        Ok(updated)
    }

    fn commit(self) -> Result<(), DatabaseError> {
        self.tx
            .into_inner()
            .commit()
            .map_err(|e| DatabaseError::Other {
                source: Box::new(e),
            })
    }
}

impl<'conn> AdminEnt for Txn<'conn> {
    fn find_edges_by_dest(&self, dest: Id) -> Result<Vec<Edge>, DatabaseError> {
        let rows = self
            .tx
            .borrow_mut()
            .query(
                "SELECT source, type, dest FROM edges WHERE dest = $1 ORDER BY source ASC, type ASC",
                &[&id_to_i64(dest)?],
            )
            .map_err(|e| DatabaseError::Other {
                source: Box::new(e),
            })?;

        rows.into_iter()
            .map(|row| {
                Ok(Edge::new(
                    i64_to_id(row.get::<_, i64>(0))?,
                    row.get::<_, Vec<u8>>(1),
                    i64_to_id(row.get::<_, i64>(2))?,
                ))
            })
            .collect()
    }

    fn remove_edges_by_dest(&self, dest: Id) -> Result<(), DatabaseError> {
        self.tx
            .borrow_mut()
            .execute("DELETE FROM edges WHERE dest = $1", &[&id_to_i64(dest)?])
            .map_err(|e| DatabaseError::Other {
                source: Box::new(e),
            })?;
        Ok(())
    }

    fn create_dyn(&self, ent: Box<dyn Ent>) -> Result<Id, DatabaseError> {
        let entity_type = ent.typetag_name().to_string();
        let data_json =
            serde_json::to_string(&ent).map_err(|e| DatabaseError::Other {
                source: Box::new(e),
            })?;

        let row = self
            .tx
            .borrow_mut()
            .query_one(
                "INSERT INTO entities (type, data) VALUES ($1, $2) RETURNING id",
                &[&entity_type, &data_json],
            )
            .map_err(|e| DatabaseError::Other {
                source: Box::new(e),
            })?;

        i64_to_id(row.get::<_, i64>(0))
    }

    fn update_dyn(&self, ent: Box<dyn Ent>) -> Result<(), DatabaseError> {
        let id = ent.id();
        let updated = self.update_internal(id, ent, None)?;
        if !updated {
            return Err(other_error("entity not found or update failed"));
        }
        Ok(())
    }

    fn list_entities(
        &self,
        entity_type: &str,
        cursor: Option<Id>,
        limit: usize,
    ) -> Result<Vec<Box<dyn Ent>>, DatabaseError> {
        let cursor = cursor.map(id_to_i64).transpose()?.unwrap_or(0);
        let limit = usize_to_i64(limit, "limit")?;
        let rows = self
            .tx
            .borrow_mut()
            .query(
                "SELECT id, data FROM entities WHERE type = $1 AND id > $2 ORDER BY id ASC LIMIT $3",
                &[&entity_type, &cursor, &limit],
            )
            .map_err(|e| DatabaseError::Other {
                source: Box::new(e),
            })?;

        rows.into_iter()
            .map(|row| {
                let id = row.get::<_, i64>(0);
                let data_json = row.get::<_, String>(1);
                parse_ent_row(id, &data_json)
            })
            .collect()
    }
}

impl<'conn> QueryEdge for Txn<'conn> {
    fn find_edges(
        &self,
        source: Id,
        query: EdgeQuery,
    ) -> Result<EdgeQueryResult, DatabaseError> {
        let mut sql = String::from(
            "SELECT source, type, dest FROM edges WHERE source = $1",
        );
        let mut params: Vec<Box<dyn ToSql + Sync>> =
            vec![Box::new(id_to_i64(source)?)];
        let mut next_param = 2;

        if !query.edge_names.is_empty() {
            sql.push_str(" AND type IN (");
            for (idx, edge_name) in query.edge_names.iter().enumerate() {
                if idx > 0 {
                    sql.push_str(", ");
                }
                sql.push('$');
                sql.push_str(&next_param.to_string());
                params.push(Box::new(edge_name.to_vec()));
                next_param += 1;
            }
            sql.push(')');
        }

        if let Some(cursor) = query.cursor {
            let comparator = match query.order {
                SortOrder::Asc => ">",
                SortOrder::Desc => "<",
            };
            sql.push_str(" AND (type, dest) ");
            sql.push_str(comparator);
            sql.push_str(" ($");
            sql.push_str(&next_param.to_string());
            sql.push_str(", $");
            sql.push_str(&(next_param + 1).to_string());
            sql.push(')');
            params.push(Box::new(cursor.sort_key.to_vec()));
            params.push(Box::new(id_to_i64(cursor.destination)?));
            next_param += 2;
        }

        let order_clause = match query.order {
            SortOrder::Asc => " ORDER BY type ASC, dest ASC",
            SortOrder::Desc => " ORDER BY type DESC, dest DESC",
        };
        sql.push_str(order_clause);
        sql.push_str(" LIMIT $");
        sql.push_str(&next_param.to_string());
        params.push(Box::new(usize_to_i64(MAX_EDGES + 1, "edge query limit")?));

        let param_refs: Vec<&(dyn ToSql + Sync)> =
            params.iter().map(|param| param.as_ref()).collect();
        let rows =
            self.tx.borrow_mut().query(&sql, &param_refs).map_err(|e| {
                DatabaseError::Other {
                    source: Box::new(e),
                }
            })?;

        let mut edges: Vec<Edge> = rows
            .into_iter()
            .map(|row| {
                Ok(Edge::new(
                    i64_to_id(row.get::<_, i64>(0))?,
                    row.get::<_, Vec<u8>>(1),
                    i64_to_id(row.get::<_, i64>(2))?,
                ))
            })
            .collect::<Result<_, DatabaseError>>()?;

        let has_more = edges.len() > MAX_EDGES;
        if has_more {
            edges.truncate(MAX_EDGES);
        }

        Ok(EdgeQueryResult { edges, has_more })
    }
}

#[derive(Clone)]
pub struct PostgresDb {
    pool: PostgresPool,
}

impl PostgresDb {
    pub fn new(pool: PostgresPool) -> Self {
        Self { pool }
    }

    pub fn connect(database_url: &str) -> Result<Self, DatabaseError> {
        let config = database_url.parse::<Config>().map_err(|e| {
            DatabaseError::Other {
                source: Box::new(e),
            }
        })?;
        let manager = PostgresConnectionManager::new(config, NoTls);
        let pool = Pool::new(manager).map_err(|e| DatabaseError::Other {
            source: Box::new(e),
        })?;
        Ok(Self::new(pool))
    }

    pub fn init(&self) -> Result<(), DatabaseError> {
        let mut conn = self.pool.get().map_err(|e| DatabaseError::Other {
            source: Box::new(e),
        })?;
        init_schema(&mut *conn)
    }
}

impl From<PostgresPool> for PostgresDb {
    fn from(pool: PostgresPool) -> Self {
        Self::new(pool)
    }
}

impl TransactionProvider for PostgresDb {
    type Tx<'a> = Txn<'a>;

    fn execute<R, F>(&self, func: F) -> Result<R, DatabaseError>
    where
        F: for<'a> FnOnce(Self::Tx<'a>) -> R,
    {
        let mut conn = self.pool.get().map_err(|e| DatabaseError::Other {
            source: Box::new(e),
        })?;
        let tx = conn.transaction().map_err(|e| DatabaseError::Other {
            source: Box::new(e),
        })?;
        Ok(func(Txn::new(tx)))
    }
}
