use std::sync::{Arc, RwLock};

use crate::{
    Error, Result,
    op::{
        BinaryOpKind, ColumnSelector, JoinKind, Literal, Op, OpPool, OpRef, ReduceKind, UnaryOpKind,
    },
    programmetadata::ProgramMetadata,
    programsignature::ProgramSignature,
};
use autark_schema::Schema;

#[derive(Clone)]
pub struct Program {
    pub op_pool: Arc<RwLock<OpPool>>,
    // metdata is too behind an Arc<RwLock<_>> as, like with op_pool, it is expected of a program to include a single, consolidated, metadata.
    pub metadata: Arc<RwLock<ProgramMetadata>>,
    pub signature: Arc<RwLock<ProgramSignature>>,
    pub root: Option<OpRef>,
}

impl Program {
    pub fn new() -> Program {
        let op_pool = Arc::new(RwLock::new(OpPool::new(1024)));
        Program {
            op_pool,
            metadata: Arc::new(RwLock::new(ProgramMetadata::default())),
            signature: Arc::new(RwLock::new(ProgramSignature::default())),
            root: None,
        }
    }

    pub fn len(&self) -> Result<usize> {
        Ok(self.op_pool.read().map_err(|_| Error::PoisonedLock)?.len())
    }

    pub fn signature(&self) -> Result<ProgramSignature> {
        Ok(self
            .signature
            .read()
            .map_err(|_| Error::PoisonedLock)?
            .clone())
    }

    pub fn ir(&self) -> Result<String> {
        Ok(serde_json::to_string(self)?)
    }

    fn with_generic(&self, op: Op) -> Result<Program> {
        let opref = self
            .op_pool
            .write()
            .map_err(|_| Error::PoisonedLock)?
            .insert(op);

        Ok(Self {
            op_pool: self.op_pool.clone(),
            metadata: self.metadata.clone(),
            signature: self.signature.clone(),

            root: Some(opref),
        })
    }

    fn root(&self) -> Result<OpRef> {
        match self.root {
            Some(x) => Ok(x),
            None => Err(Error::ProvidedEmptyProgram),
        }
    }

    fn source(&self, on: Option<OpRef>, index: usize, selector: ColumnSelector) -> Result<Program> {
        self.with_generic(Op::Columns {
            on,
            index,
            selector,
        })
    }

    pub fn columns(&self, index: Option<usize>, selector: ColumnSelector) -> Result<Program> {
        let (on, index) = match (self.root, index) {
            (_, Some(index)) => (None, index),
            (Some(root), None) => (Some(root), 0),
            (None, None) => (None, 0),
        };
        if on.is_none() {
            let mut signature = self.signature.write().map_err(|_| Error::PoisonedLock)?;
            match &selector {
                ColumnSelector::All => signature.mark_all(index),
                ColumnSelector::Names(columns) => {
                    let columns = columns.iter().map(String::as_str).collect::<Vec<_>>();
                    signature.pop(index, &columns);
                }
                ColumnSelector::Indices(_) => signature.ensure(index),
            }
        }
        self.source(on, index, selector)
    }

    #[inline]
    pub fn cols(&self, index: Option<usize>, selector: ColumnSelector) -> Result<Program> {
        self.columns(index, selector)
    }

    pub fn concat(&self, who: &[Program]) -> Result<Program> {
        let mut op_refs = Vec::with_capacity(who.len() + 1);
        op_refs.push(self.root()?);
        for program in who {
            op_refs.push(program.root()?);
        }

        self.with_generic(Op::Concat { who: op_refs })
    }

    pub fn slice(&self, start: Option<isize>, end: Option<isize>) -> Result<Program> {
        self.with_generic(Op::Slice {
            on: self.root()?,
            start,
            end,
        })
    }

    pub fn rolling(&self, n: usize) -> Result<Program> {
        self.with_generic(Op::Rolling {
            on: self.root()?,
            n,
        })
    }

    pub fn cumsum(&self) -> Result<Program> {
        self.unaryop(UnaryOpKind::Cumsum)
    }

    pub fn order_by(&self, by: Option<Program>, ascending: bool) -> Result<Program> {
        let by = by.unwrap_or_else(|| self.clone());
        self.with_generic(Op::OrderBy {
            what: self.root()?,
            by: by.root()?,
            ascending,
        })
    }

    pub fn alias(&self, name: &str, schema: Option<Schema>) -> Result<Program> {
        if let Some(schema) = schema {
            self.metadata
                .write()
                .map_err(|_| Error::PoisonedLock)?
                .schema_map
                .insert(name.to_string(), schema.into_arrow());
        }

        self.with_generic(Op::Output {
            name: name.to_string(),
            value: self.root()?,
        })
    }

    pub fn binaryop(&self, rhs: Program, kind: BinaryOpKind) -> Result<Program> {
        self.with_generic(Op::BinaryOp {
            kind,
            lhs: self.root()?,
            rhs: rhs.root()?,
        })
    }

    pub fn filter(&self, mask: Program) -> Result<Program> {
        self.with_generic(Op::Filter {
            on: self.root()?,
            mask: mask.root()?,
        })
    }

    pub fn r#where(&self, mask: Program, else_: Program) -> Result<Program> {
        self.with_generic(Op::Where {
            mask: mask.root()?,
            then: self.root()?,
            else_: else_.root()?,
        })
    }

    pub fn lit(&self, value: Literal) -> Result<Program> {
        self.with_generic(Op::Constant { value })
    }

    pub fn const_f64(&self, value: f64) -> Result<Program> {
        self.lit(Literal::F64(value))
    }

    pub fn const_string(&self, value: impl Into<String>) -> Result<Program> {
        self.lit(Literal::String(value.into()))
    }

    pub fn reduce(&self, kind: ReduceKind) -> Result<Program> {
        self.with_generic(Op::Reduce {
            kind,
            on: self.root()?,
        })
    }

    pub fn unaryop(&self, kind: UnaryOpKind) -> Result<Program> {
        self.with_generic(Op::UnaryOp {
            kind,
            on: self.root()?,
        })
    }

    pub fn group_by(&self, keys: Option<Program>) -> Result<Program> {
        let keys = keys.unwrap_or_else(|| self.clone());
        let keys = keys.root()?;
        let values = self.root()?;
        self.with_generic(Op::GroupBy { keys, values })
    }

    pub fn group_reduce(&self, keys: Option<Program>, kinds: Vec<ReduceKind>) -> Result<Program> {
        let keys = keys.unwrap_or_else(|| self.clone());
        let keys = keys.root()?;
        let values = self.root()?;
        self.with_generic(Op::GroupReduce {
            keys,
            values,
            kinds,
        })
    }

    pub fn join(
        &self,
        right: Program,
        left_on: Program,
        right_on: Program,
        kind: JoinKind,
    ) -> Result<Program> {
        self.with_generic(Op::Join {
            left: self.root()?,
            right: right.root()?,
            left_on: left_on.root()?,
            right_on: right_on.root()?,
            kind,
        })
    }

    // pub fn const_f64(&mut self, value: f64) -> OpRef {
    //     self.op_pool.insert(Op::Constant { value: Literal::F64(value) })
    // }
}
