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

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

use super::context::FunctionContext;
use super::context::ModuleContext;
use super::context::ModuleContextMode;
use super::context::RefRenderStyle;
use super::names;
use super::types::*;

#[derive(Debug)]
pub struct ReaderPlugin;

impl compiler::Plugin for ReaderPlugin {
    fn generate_definition(
        &self,
        module_ctxt: &mut ModuleContext,
        definition_name: &str,
        definition: &Definition,
    ) {
        if let ModuleContextMode::TargetGeneric = module_ctxt.mode {
            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]
            ],
            "? ",
            codeblock(inner)
        ]))
    }
}

pub fn gen_definition_reader(m: &mut ModuleContext, n: &str, d: &Definition) {
    let ty = definition_type(&m.module_path, Purpose::Codegen, n, d);

    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",
                                    seq![ctxt.m.any_type(), ": preserves::value::NestedValue"],
                                    seq!["R: _support::Reader<'de, ", ctxt.m.any_type(), ">"]
                                ],
                                "(r: &mut R) -> ",
                                "std::result::Result<",
                                names::render_constructor(n),
                                ty.generic_arg(ctxt.m),
                                ", _support::ParseError> ",
                                codeblock(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];
                let mut need_restore = false;
                ps.extend(pattern_n);
                ctxt.define_atom(&mut body, "_mark", item("r.mark()?"));
                for e in &ps {
                    if need_restore {
                        body.push(item("r.restore(&_mark)?;"));
                    } else {
                        need_restore = true;
                    }
                    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![seq![ctxt.m.any_type(), ": preserves::value::NestedValue"]],
            " _support::Deserialize",
            anglebrackets![ctxt.m.any_type()],
            " for ",
            names::render_constructor(n),
            ty.generic_arg(ctxt.m),
            " ",
            codeblock![seq![
                "fn deserialize",
                anglebrackets![
                    "'de",
                    seq!["R: _support::Reader<'de, ", ctxt.m.any_type(), ">"]
                ],
                "(r: &mut R) -> ",
                "std::result::Result<Self, _support::ParseError> ",
                codeblock(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(_) | TField::Any | TField::Embedded => 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<(IOValue, LiteralContinuation)>;
type LiteralSeqCases = Vec<(Vec<IOValue>, 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 => ",
                    codeblock(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)? ", codeblock(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::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())),
                        },
                        " => ",
                        codeblock(inner)]));
                }
                subcases.push(item(seq!["_ => return ", ctxt.err_code(), "?,"]));
                item(seq!["preserves::value::Token::Atom(v) => match v.value() ", codeblock(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) => ",
                    codeblock(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) => ",
                    codeblock(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)? ", codeblock(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::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 { .. } => {
            ctxt.define_atom(
                body,
                &dest,
                item("r.demand_next(true)?.value().to_embedded()?.clone()"),
            );
            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 compound_dest = ctxt.gentempname();
            ctxt.with_definite_mode(|ctxt| {
                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![
                    compound_dest.to_owned(),
                    ".push(",
                    store_wrap(true, &field_type(pattern), &item_dest),
                    ");"
                ]));
                ctxt.declare_compound(body, &compound_dest, item("std::vec::Vec::new()"));
                boundary_tracker.emit_loop(body, inner);
            });
            ctxt.define_atom(body, &dest, item(compound_dest));
            dest
        }
        SimplePattern::Setof { pattern } => {
            let compound_dest = ctxt.gentempname();
            ctxt.with_definite_mode(|ctxt| {
                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![
                    compound_dest.to_owned(),
                    ".insert(",
                    store_wrap(true, &field_type(pattern), &item_dest),
                    ");"
                ]));
                ctxt.declare_compound(body, &compound_dest, item("preserves::value::Set::new()"));
                boundary_tracker.emit_loop(body, inner);
            });
            ctxt.define_atom(body, &dest, item(compound_dest));
            dest
        }
        SimplePattern::Dictof { key, value } => {
            let compound_dest = ctxt.gentempname();
            ctxt.with_definite_mode(|ctxt| {
                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![
                    compound_dest.to_owned(),
                    ".insert(",
                    store_wrap(true, &field_type(key), &key_dest),
                    ", ",
                    store_wrap(true, &field_type(value), &value_dest),
                    ");"
                ]));
                ctxt.declare_compound(body, &compound_dest, item("preserves::value::Map::new()"));
                boundary_tracker.item_expr = "_support::B::Item::DictionaryKey";
                boundary_tracker.emit_loop(body, inner);
            });
            ctxt.define_atom(body, &dest, item(compound_dest));
            dest
        }
        SimplePattern::Ref(r) => {
            let tf = name![ctxt.m.render_ref(&**r, RefRenderStyle::Bare), "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
        }
    }
}