use crate::*;
use crate::gen::schema::*;
use crate::syntax::block::Item;
use crate::syntax::block::escape_string;
use crate::syntax::block::escape_bytes;
use crate::syntax::block::constructors::*;

use preserves::value::AtomClass;
use preserves::value::CompoundClass;
use preserves::value::NestedValue;
use preserves::value::ValueClass;

use super::context::{ModuleContext, FunctionContext};
use super::names;
use super::types::*;

#[derive(Debug)]
pub struct ReaderPlugin;

impl compiler::Plugin for ReaderPlugin {
    fn generate(&self, module_ctxt: &mut ModuleContext, definition_name: &str, definition: &Definition) {
        if module_ctxt.mode.is_none() {
            gen_definition_reader(module_ctxt, definition_name, definition)
        }
    }
}

#[derive(Clone)]
struct BoundaryTracker {
    tracker_name: String,
    item_expr: &'static str,
}

impl BoundaryTracker {
    fn unwrap(
        ctxt: &mut FunctionContext,
        body: &mut Vec<Item>,
        open_expr: &'static str,
        e: Option<&BoundaryTracker>,
    ) -> Self {
        match e {
            None => Self::new(ctxt, body, open_expr, "_support::B::Item::SequenceValue"),
            Some(b) => b.clone(),
        }
    }

    fn new(
        ctxt: &mut FunctionContext,
        body: &mut Vec<Item>,
        open_expr: &'static str,
        item_expr: &'static str,
    ) -> Self {
        let tracker_name = ctxt.gentempname();
        body.push(item(open_expr));
        body.push(item(seq!["let mut ", tracker_name.clone(), " = _support::B::Type::default();"]));
        BoundaryTracker {
            tracker_name,
            item_expr,
        }
    }

    fn emit_boundary(&self, body: &mut Vec<Item>) {
        body.push(item(seq![self.tracker_name.clone(), ".shift(Some(", self.item_expr, "));"]));
        body.push(item(seq!["r.boundary(&", self.tracker_name.clone(), ")?;"]));
    }

    fn emit_loop(&self, body: &mut Vec<Item>, inner: Vec<Item>) {
        body.push(item(seq![
            "while !r.close_compound", parens![
                seq!["&mut ", self.tracker_name.clone()],
                seq!["&", self.item_expr]], "? ",
            block(inner)]))
    }
}

pub fn gen_definition_reader(m: &mut ModuleContext, n: &str, d: &Definition) {
    m.define_function(
        n,
        |mut ctxt| {
            let mut body = vec![];

            match d {
                Definition::Or { pattern_0, pattern_1, pattern_n } => {
                    let mut ps = vec![&**pattern_0, &**pattern_1];
                    ps.extend(pattern_n);
                    ctxt.define_atom(&mut body, "_mark", item("r.mark()?"));
                    for NamedAlternative { variant_label: name, pattern: pat } in ps {
                        let fname = seq!["read_", names::render_fieldname(n), "_", names::render_fieldname(name)];
                        let ctorname = item(name![names::render_constructor(n), names::render_constructor(name)]);
                        ctxt.m.define_function(
                            &(n.to_owned() + "::" + name),
                            |mut ctxt| {
                                let mut body = Vec::new();
                                let dest = pattern_reader(&mut ctxt, pat, None, &mut body);
                                let dest = dest.as_ref().map(String::as_str);
                                construct(&ctxt, ctorname, false, &pattern_type(pat), dest, &mut body);
                                item(seq![
                                    "fn ", fname.clone(), anglebrackets![
                                        "'de",
                                        "R: _support::Reader<'de, _Ptr, _Any>"],
                                    "(r: &mut R) -> ",
                                    "std::result::Result<", names::render_constructor(n), ", _support::ParseError> ",
                                    block(body)])
                            });
                        body.push(item(seq![
                            "match ", fname, "(r) { ",
                            "Err(e) if e.is_conformance_error() => r.restore(&_mark)?, ",
                            "result => return result }"]));
                    }
                    body.push(item(seq![ctxt.err_code()]));
                }
                Definition::And { pattern_0, pattern_1, pattern_n } => {
                    let mut ps = vec![&**pattern_0, &**pattern_1];
                    ps.extend(pattern_n);
                    for e in &ps {
                        named_pattern_reader(&mut ctxt, e, None, &mut body);
                    }
                    construct(&ctxt, item(names::render_constructor(n)), true, &record_type(&ps), None, &mut body);
                }
                Definition::Pattern(p) => {
                    let dest = pattern_reader(&mut ctxt, p, None, &mut body);
                    let dest = dest.as_ref().map(String::as_str);
                    construct(&ctxt, item(names::render_constructor(n)), true, &pattern_type(p), dest, &mut body);
                }
            }

            item(seq![
                "impl", anglebrackets!["'de", "R: _support::Reader<'de, _Ptr, _Any>"], " ",
                "_support::Deserialize", anglebrackets!["'de", "_Ptr", "_Any", "R"], " ",
                "for ", names::render_constructor(n), " ", block![
                    seq!["fn deserialize(r: &mut R) -> ",
                         "std::result::Result<Self, _support::ParseError> ",
                         block(body)]]])
        });
}

fn construct(
    ctxt: &FunctionContext,
    ctorname: Item,
    is_struct: bool,
    ty: &TSimple,
    dest: Option<&str>,
    body: &mut Vec<Item>,
) {
    match ty {
        TSimple::Field(TField::Unit) =>
            body.push(item(seq!["Ok(", ctorname, ")"])),
        TSimple::Field(fieldty) =>
            body.push(item(seq!["Ok(", ctorname, parens![store_wrap(is_struct, fieldty, dest.unwrap())], ")"])),
        TSimple::Record(_) =>
            body.push(item(seq!["Ok(", ctorname, " ", braces(
                ctxt.captures.iter().map(
                    |c| item(seq![c.field_name.clone(), ": ", store_wrap(is_struct, &c.ty, &c.source_expr)])).collect()),
                                ")"])),
    }
}

fn store_wrap(is_struct: bool, ty: &TField, expr: &str) -> String {
    match ty {
        TField::Unit
            | TField::Array(_)
            | TField::Set(_)
            | TField::Map(_, _) => expr.to_owned(),
        TField::Ref(_) =>
            if is_struct {
                expr.to_owned()
            } else {
                format!("std::boxed::Box::new({})", expr)
            },
        TField::Base(_) => expr.to_owned(),
    }
}

fn group_by<T, K, Key>(mut items: Vec<T>, mut key: Key) -> Vec<(K, Vec<T>)>
where
    K: Ord + Clone,
    Key: FnMut(&T) -> K,
{
    let mut result = Vec::new();
    let mut current_key: Option<K> = None;
    let mut buf = Vec::new();
    items.sort_by(|a, b| key(a).cmp(&key(b)));
    for (k, v) in items.into_iter().map(|t| (key(&t), t)) {
        match current_key.cmp(&Some(k.clone())) {
            std::cmp::Ordering::Equal => (),
            std::cmp::Ordering::Less |
            std::cmp::Ordering::Greater => {
                if let Some(k) = current_key {
                    result.push((k, std::mem::take(&mut buf)));
                }
                current_key = Some(k);
            }
        }
        buf.push(v)
    }
    if let Some(k) = current_key {
        result.push((k.clone(), std::mem::take(&mut buf)));
    }
    result
}

type LiteralContinuation = Box<dyn FnOnce(&mut FunctionContext, &mut Vec<Item>) -> ()>;
type LiteralCases = Vec<(_Any, LiteralContinuation)>;
type LiteralSeqCases = Vec<(Vec<_Any>, LiteralContinuation)>;

fn read_expected_literal_seqs(
    ctxt: &mut FunctionContext,
    body: &mut Vec<Item>,
    possibilities: LiteralSeqCases,
) {
    let grouped = group_by(possibilities, |(vs, _f)| if vs.is_empty() {
        None
    } else {
        Some(vs[0].clone())
    });
    let mut cases = Vec::new();
    let mut nested: LiteralCases = Vec::new();
    for (head, group) in grouped.into_iter() {
        match head {
            None => {
                let mut inner = Vec::new();
                group.into_iter().next().unwrap().1(ctxt, &mut inner);
                cases.push(item(seq!["preserves::value::Token::End => ", block(inner)]));
            },
            Some(h) => {
                let tails = group.into_iter().map(|(mut vs, f)| {
                    vs.remove(0);
                    (vs, f)
                }).collect();
                nested.push((h, Box::new(|ctxt: &mut FunctionContext, b: &'_ mut Vec<Item>| {
                    read_expected_literal_seqs(ctxt, b, tails)
                })));
            }
        }
    }
    cases.extend(read_expected_literals_cases(ctxt, nested));
    body.push(item(seq!["match r.next_token(true)? ", block(cases)]));
}

fn read_expected_literals_cases(
    ctxt: &mut FunctionContext,
    possibilities: LiteralCases,
) -> Vec<Item> {
    let grouped = group_by(possibilities, |(v, _f)| v.value_class());
    let mut cases = grouped.into_iter().map(|(n, group)| {
        match n {
            ValueClass::Atomic(cls) => {
                let mut subcases = Vec::new();
                for p in group {
                    let mut inner = Vec::new();
                    p.1(ctxt, &mut inner);
                    subcases.push(item(seq![
                        format!("preserves::value::Value::{:?}(w)", cls),
                        match cls {
                            AtomClass::Boolean => match p.0.value().to_boolean().unwrap() {
                                true => " if *w".to_owned(),
                                false => " if !*w".to_owned(),
                            },
                            AtomClass::Float |
                            AtomClass::Double =>
                                format!(" if w.0 == {:?}", p.0),
                            AtomClass::SignedInteger =>
                                format!(" if *w == ({:?}).into()", p.0),
                            AtomClass::String =>
                                format!(" if w == {}", escape_string(p.0.value().to_string().unwrap())),
                            AtomClass::ByteString =>
                                format!(" if w == {}", escape_bytes(p.0.value().to_bytestring().unwrap())),
                            AtomClass::Symbol =>
                                format!(" if w == {}", escape_string(p.0.value().to_symbol().unwrap())),
                        },
                        " => ",
                        block(inner)]));
                }
                subcases.push(item(seq!["_ => return ", ctxt.err_code(), "?,"]));
                item(seq!["preserves::value::Token::Atom(v) => match v.value() ", block(subcases)])
            }
            ValueClass::Compound(CompoundClass::Record) => {
                let mut subcases = Vec::new();
                read_expected_literal_seqs(ctxt, &mut subcases, group.into_iter().map(|(v, f)| {
                    let r = v.value().to_record(None).unwrap();
                    (r.0.clone(), f)
                }).collect());
                item(seq![
                    "preserves::value::Token::Compound(preserves::value::CompoundClass::Record) => ",
                    block(subcases)])
            }
            ValueClass::Compound(CompoundClass::Sequence) => {
                let mut subcases = Vec::new();
                read_expected_literal_seqs(ctxt, &mut subcases, group.into_iter().map(|(v, f)| {
                    let s = v.value().to_sequence().unwrap().clone();
                    (s, f)
                }).collect());
                item(seq![
                    "preserves::value::Token::Compound(preserves::value::CompoundClass::Sequence) => ",
                    block(subcases)])
            }
            ValueClass::Compound(CompoundClass::Set) => {
                panic!("Sets in literal constants in Schema not yet supported");
            }
            ValueClass::Compound(CompoundClass::Dictionary) => {
                panic!("Dictionaries in literal constants in Schema not yet supported");
            }
            ValueClass::Embedded => {
                panic!("Embedded values in literal constants in Schema not yet supported");
            }
        }
    }).collect::<Vec<_>>();
    cases.push(item(seq!["_ => return ", ctxt.err_code(), "?,"]));
    cases
}

fn read_expected_literals(
    ctxt: &mut FunctionContext,
    body: &mut Vec<Item>,
    possibilities: LiteralCases,
) {
    let cases = read_expected_literals_cases(ctxt, possibilities);
    body.push(item(seq!["match r.next_token(true)? ", block(cases)]));
}

fn simple_pattern_reader(
    ctxt: &mut FunctionContext,
    p: &SimplePattern,
    boundary_tracker: Option<&BoundaryTracker>,
    body: &mut Vec<Item>,
) -> String {
    let dest = ctxt.gentempname();
    match p {
        SimplePattern::Any => {
            ctxt.define_atom(body, &dest, item("r.demand_next(true)?"));
            dest
        },
        SimplePattern::Atom { atom_kind: k } => {
            let reader = match &**k {
                AtomKind::Boolean => "r.next_boolean()?",
                AtomKind::Float => "r.next_float()?",
                AtomKind::Double => "r.next_double()?",
                AtomKind::SignedInteger => "r.next_signedinteger()?",
                AtomKind::String => "r.next_str()?.into_owned()",
                AtomKind::ByteString => "r.next_bytestring()?.into_owned()",
                AtomKind::Symbol => "r.next_symbol()?.into_owned()",
            };
            ctxt.define_atom(body, &dest, item(reader.to_owned()));
            dest
        },
        SimplePattern::Embedded { .. } => {
            // TODO: Is this right? If so, why doesn't it expect *two* Embedded tags in a row??
            body.push(item("r.open_embedded()?;"));
            ctxt.define_atom(body, &dest, item("r.demand_next(true)?.value().to_embedded()?.clone()"));
            body.push(item("r.close_embedded()?;"));
            dest
        },
        SimplePattern::Lit { value } => {
            let f = Box::new(|_ctxt: &mut FunctionContext, _: &'_ mut Vec<Item>| ());
            read_expected_literals(ctxt, body, vec![(value.clone(), f)]);
            ctxt.define_atom(body, &dest, item("()"));
            dest
        },
        SimplePattern::Seqof { pattern } => {
            let boundary_tracker = BoundaryTracker::unwrap(
                ctxt, body, "r.open_sequence()?;", boundary_tracker);
            let mut inner = Vec::new();
            boundary_tracker.emit_boundary(&mut inner);
            let item_dest = simple_pattern_reader(ctxt, pattern, None, &mut inner);
            inner.push(item(seq![dest.to_owned(), ".push(",
                                 store_wrap(true, &field_type(pattern), &item_dest), ");"]));
            ctxt.declare_compound(body, &dest, item("std::vec::Vec::new()"));
            boundary_tracker.emit_loop(body, inner);
            dest
        },
        SimplePattern::Setof { pattern } => {
            let boundary_tracker = BoundaryTracker::new(
                ctxt, body, "r.open_set()?;", "_support::B::Item::SetValue");
            let mut inner = Vec::new();
            boundary_tracker.emit_boundary(&mut inner);
            let item_dest = simple_pattern_reader(ctxt, pattern, None, &mut inner);
            inner.push(item(seq![dest.to_owned(), ".insert(",
                                 store_wrap(true, &field_type(pattern), &item_dest), ");"]));
            ctxt.declare_compound(body, &dest, item("preserves::value::Set::new()"));
            boundary_tracker.emit_loop(body, inner);
            dest
        },
        SimplePattern::Dictof { key, value } => {
            let mut boundary_tracker = BoundaryTracker::new(
                ctxt, body, "r.open_dictionary()?;", "_support::B::Item::DictionaryKey");
            let mut inner = Vec::new();
            boundary_tracker.emit_boundary(&mut inner);
            let key_dest = simple_pattern_reader(ctxt, key, None, &mut inner);
            boundary_tracker.item_expr = "_support::B::Item::DictionaryValue";
            boundary_tracker.emit_boundary(&mut inner);
            let value_dest = simple_pattern_reader(ctxt, value, None, &mut inner);
            inner.push(item(seq![
                dest.to_owned(), ".insert(",
                store_wrap(true, &field_type(key), &key_dest), ", ",
                store_wrap(true, &field_type(value), &value_dest), ");"]));
            ctxt.declare_compound(body, &dest, item("preserves::value::Map::new()"));
            boundary_tracker.item_expr = "_support::B::Item::DictionaryKey";
            boundary_tracker.emit_loop(body, inner);
            dest
        },
        SimplePattern::Ref(r) => {
            let tf = name![ctxt.m.render_ref(&**r), "deserialize"];
            ctxt.define_atom(body, &dest, item(seq![tf, "(r)?"]));
            dest
        },
    }
}

fn named_pattern_reader(
    ctxt: &mut FunctionContext,
    p: &NamedPattern,
    boundary_tracker: Option<&BoundaryTracker>,
    body: &mut Vec<Item>,
) {
    match p {
        NamedPattern::Anonymous(p) => {
            pattern_reader(ctxt, p, boundary_tracker, body);
        },
        NamedPattern::Named(b) => {
            let Binding { name, pattern} = &**b;
            let dest = simple_pattern_reader(ctxt, pattern, boundary_tracker, body);
            let capture_ty = field_type(pattern);
            ctxt.capture(names::render_fieldname(name), capture_ty, dest);
        }
    }
}

fn pattern_reader(
    ctxt: &mut FunctionContext,
    p: &Pattern,
    boundary_tracker: Option<&BoundaryTracker>,
    body: &mut Vec<Item>,
) -> Option<String> {
    match p {
        Pattern::SimplePattern(s) =>
            Some(simple_pattern_reader(ctxt, s, boundary_tracker, body)),
        Pattern::CompoundPattern(c) => {
            match &**c {
                CompoundPattern::Rec { label, fields } => {
                    let mut boundary_tracker = BoundaryTracker::new(
                        ctxt, body, "r.open_record(None)?;", "_support::B::Item::RecordLabel");
                    boundary_tracker.emit_boundary(body);
                    boundary_tracker.item_expr = "_support::B::Item::RecordField";
                    named_pattern_reader(ctxt, &**label, None, body);
                    named_pattern_reader(ctxt, &**fields, Some(&boundary_tracker), body);
                },
                CompoundPattern::Tuple { patterns } => {
                    let boundary_tracker = BoundaryTracker::unwrap(
                        ctxt, body, "r.open_sequence()?;", boundary_tracker);
                    for p in patterns {
                        boundary_tracker.emit_boundary(body);
                        named_pattern_reader(ctxt, p, None, body);
                    }
                    body.push(item(seq!["r.ensure_complete", parens![
                        boundary_tracker.tracker_name.clone(),
                        seq!["&", boundary_tracker.item_expr]], "?;"]));
                },
                CompoundPattern::TuplePrefix { fixed, variable } => {
                    let boundary_tracker = BoundaryTracker::unwrap(
                        ctxt, body, "r.open_sequence()?;", boundary_tracker);
                    for p in fixed {
                        boundary_tracker.emit_boundary(body);
                        named_pattern_reader(ctxt, p, None, body);
                    }
                    named_pattern_reader(ctxt, &promote(variable), Some(&boundary_tracker), body);
                },
                CompoundPattern::Dict { entries } => {
                    let boundary_tracker = BoundaryTracker::new(
                        ctxt, body, "r.open_dictionary()?;", "_support::B::Item::DictionaryKey");
                    let mut inner = Vec::new();
                    boundary_tracker.emit_boundary(&mut inner);
                    let mut val_boundary_tracker = boundary_tracker.clone();
                    val_boundary_tracker.item_expr = "_support::B::Item::DictionaryValue";
                    body.extend(ctxt.with_indefinite_mode(|ctxt| {
                        read_expected_literals(ctxt, &mut inner, entries.0.iter().map(move |(key_lit, value_pat)| {
                            let value_pat = value_pat.clone();
                            let val_boundary_tracker = val_boundary_tracker.clone();
                            let f: LiteralContinuation = Box::new(
                                move |ctxt: &mut FunctionContext, innerinner: &mut Vec<Item>| {
                                    val_boundary_tracker.emit_boundary(innerinner);
                                    named_pattern_reader(ctxt, &promote(&value_pat), None, innerinner);
                                    innerinner.push(item("continue;"));
                                });
                            (key_lit.clone(), f)
                        }).collect());
                    }));
                    boundary_tracker.emit_loop(body, inner);
                }
            }
            None
        },
    }
}