use clingo::{ClingoError, Control, FactBase, Part, ShowType, SolveMode, Symbol, ToSymbol};
use serde::Serialize;

use crate::{SolveError, SolveRequest, SolveResponse};

// Potassco solution for n-queens.
const BASE_PROGRAM: &str = r#"
    { queen (I, 1..n) } = 1 :- I = 1..n.
    { queen (1..n, J) } = 1 :- J = 1..n.
    :- { queen (D - J, J) } > 1, D = 2..2 * n.
    :- { queen (D + J, J) } > 1, D = 1 - n..n - 1.
"#;

#[derive(Debug, Clone, Copy, PartialEq, Eq)]
pub struct Queen(OneBased<u32>, OneBased<u32>);

#[derive(Debug, Clone, Copy, PartialEq, Eq, Serialize)]
#[serde(transparent)]
pub struct ZeroBased<T>(T);

#[derive(Debug, Clone, Copy, PartialEq, Eq, Serialize)]
#[serde(transparent)]
pub struct OneBased<T>(T);

impl ToSymbol for Queen {
    fn symbol(&self) -> Result<Symbol, ClingoError> {
        let args = &[
            Symbol::create_number(self.0.raw() as i32),
            Symbol::create_number(self.1.raw() as i32),
        ];
        Symbol::create_function("queen", args, true)
    }
}

impl OneBased<u32> {
    pub fn new(num: u32) -> Self {
        debug_assert!(num >= 1);
        OneBased(num)
    }
    pub fn raw(&self) -> u32 {
        self.0
    }
}

impl ZeroBased<u32> {
    pub fn new(num: u32) -> Self {
        ZeroBased(num)
    }
    pub fn raw(&self) -> u32 {
        self.0
    }
}

pub fn solve(mut req: SolveRequest) -> Result<SolveResponse, SolveError> {
    // Setup
    let mut ctl = Control::new(vec![])?;
    ctl.add("base", &["n"], BASE_PROGRAM)?;
    let base_parameter = [Symbol::create_number(req.size as i32)];
    let base = Part::new("base", &base_parameter)?;
    let mut fact_base = FactBase::new();
    // Add fixed queens to fact base
    req.fixed_queens
        .drain(..)
        .map(|(x, y)| Queen(x.into(), y.into()))
        .for_each(|queen| fact_base.insert(&queen));
    ctl.add_facts(&fact_base);
    // Ground
    ctl.ground(&[base])?;
    // Solve
    let mut handle = ctl.solve(SolveMode::YIELD | SolveMode::ASYNC, &[])?;
    handle.resume()?;
    let model = handle.model()?.ok_or(SolveError::NoSolution)?;
    // Parse model
    let mut response = SolveResponse {
        id: req.id,
        queens: vec![],
    };
    for atom in model.symbols(ShowType::SHOWN)? {
        match atom.name()? {
            "queen" => {
                let params = atom.arguments()?;
                debug_assert!(params.len() == 2);
                let x_param = OneBased(params[0].number()? as u32);
                let y_param = OneBased(params[1].number()? as u32);
                response.queens.push((x_param.into(), y_param.into()));
            }
            _ => unreachable!(),
        }
    }
    Ok(response)
}

impl From<OneBased<u32>> for ZeroBased<u32> {
    fn from(num: OneBased<u32>) -> Self {
        ZeroBased(num.0 - 1)
    }
}

impl From<ZeroBased<u32>> for OneBased<u32> {
    fn from(num: ZeroBased<u32>) -> Self {
        OneBased(num.0 + 1)
    }
}