use crate::eval::PrimitiveType;
use crate::ir::{Expression, Primitive, Program, TopLevel, Type, Value, ValueOrRef, Variable};
use crate::syntax::Location;
use crate::util::scoped_map::ScopedMap;
use proptest::strategy::{NewTree, Strategy, ValueTree};
use proptest::test_runner::{TestRng, TestRunner};
use rand::distributions::{Distribution, WeightedIndex};
use rand::seq::SliceRandom;
use rand::Rng;
use std::str::FromStr;

lazy_static::lazy_static! {
    static ref PROGRAM_LENGTH_DISTRIBUTION: WeightedIndex<usize> = WeightedIndex::new([
        0, // % chance of 0
        10, // % chance of 1
        10, // % chance of 2
        15, // % chance of 3
        10, // % chance of 4
        10, // % chance of 5
        10, // % chance of 6
        5, // % chance of 7
        5, // % chance of 8
        5, // % chance of 9
        5, // % chance of 10
        5, // % chance of 11
        3, // % chance of 12
        3, // % chance of 13
        3, // % chance of 14
        1, // % chance of 15
    ]).unwrap();

    static ref BLOCK_LENGTH_DISTRIBUTION: WeightedIndex<usize> = WeightedIndex::new([
        1, // % chance of 0
        10, // % chance of 1
        20, // % chance of 2
        15, // % chance of 3
        10, // % chance of 4
        8, // % chance of 5
        8, // % chance of 6
        5, // % chance of 7
        5, // % chance of 8
        4, // % chance of 9
        3, // % chance of 10
        3, // % chance of 11
        3, // % chance of 12
        2, // % chance of 13
        2, // % chance of 14
        1, // % chance of 15
    ]).unwrap();

    static ref FUNCTION_ARGUMENTS_DISTRIBUTION: WeightedIndex<usize> = WeightedIndex::new([
        5, // % chance of 0
        20, // % chance of 1
        20, // % chance of 2
        20, // % chance of 3
        15, // % chance of 4
        10, // % chance of 5
        5, // % chance of 6
        2, // % chance of 7
        1, // % chance of 8
        1, // % chance of 9
        1, // % chance of 10
    ]).unwrap();

    static ref STATEMENT_TYPE_DISTRIBUTION: WeightedIndex<usize> = WeightedIndex::new(
        STATEMENT_TYPE_FREQUENCIES.iter().map(|x| x.1)
    ).unwrap();

    static ref EXPRESSION_TYPE_DISTRIBUTION: WeightedIndex<usize> = WeightedIndex::new(
        EXPRESSION_TYPE_FREQUENCIES.iter().map(|x| x.1)
    ).unwrap();

    static ref ARGUMENT_TYPE_DISTRIBUTION: WeightedIndex<usize> = WeightedIndex::new(
        ARGUMENT_TYPE_FREQUENCIES.iter().map(|x| x.1)
    ).unwrap();

    static ref VALUE_TYPE_DISTRIBUTION: WeightedIndex<usize> = WeightedIndex::new(
        VALUE_TYPE_FREQUENCIES.iter().map(|x| x.1)
    ).unwrap();
}

static STATEMENT_TYPE_FREQUENCIES: &[(StatementType, usize)] = &[
    (StatementType::Binding, 3),
    (StatementType::Function, 1),
    (StatementType::Expression, 2),
];

static EXPRESSION_TYPE_FREQUENCIES: &[(ExpressionType, usize)] = &[
    (ExpressionType::Atomic, 50),
    (ExpressionType::Cast, 5),
    (ExpressionType::Primitive, 5),
    (ExpressionType::Block, 10),
    (ExpressionType::Print, 10),
    (ExpressionType::Bind, 20),
];

static ARGUMENT_TYPE_FREQUENCIES: &[(Type, usize)] = &[
    (Type::Primitive(PrimitiveType::U8), 1),
    (Type::Primitive(PrimitiveType::U16), 1),
    (Type::Primitive(PrimitiveType::U32), 1),
    (Type::Primitive(PrimitiveType::U64), 1),
    (Type::Primitive(PrimitiveType::I8), 1),
    (Type::Primitive(PrimitiveType::I16), 1),
    (Type::Primitive(PrimitiveType::I32), 1),
    (Type::Primitive(PrimitiveType::I64), 1),
];

enum StatementType {
    Binding,
    Function,
    Expression,
}

enum ExpressionType {
    Atomic,
    Cast,
    Primitive,
    Block,
    Print,
    Bind,
}

// this knowingly excludes void
static VALUE_TYPE_FREQUENCIES: &[(ValueType, usize)] = &[
    (ValueType::I8, 1),
    (ValueType::I16, 1),
    (ValueType::I32, 1),
    (ValueType::I64, 1),
    (ValueType::U8, 1),
    (ValueType::U16, 1),
    (ValueType::U32, 1),
    (ValueType::U64, 1),
];

#[derive(Copy, Clone)]
enum ValueType {
    I8,
    I16,
    I32,
    I64,
    U8,
    U16,
    U32,
    U64,
    Void,
}

impl From<PrimitiveType> for ValueType {
    fn from(value: PrimitiveType) -> Self {
        match value {
            PrimitiveType::U8 => ValueType::U8,
            PrimitiveType::U16 => ValueType::U16,
            PrimitiveType::U32 => ValueType::U32,
            PrimitiveType::U64 => ValueType::U64,
            PrimitiveType::I8 => ValueType::I8,
            PrimitiveType::I16 => ValueType::I16,
            PrimitiveType::I32 => ValueType::I32,
            PrimitiveType::I64 => ValueType::I64,
            PrimitiveType::Void => ValueType::Void,
        }
    }
}

#[derive(Debug, Default)]
pub struct ProgramGenerator {}

impl Strategy for ProgramGenerator {
    type Tree = ProgramTree;
    type Value = Program<Type>;

    fn new_tree(&self, runner: &mut TestRunner) -> NewTree<Self> {
        NewTree::<ProgramGenerator>::Ok(ProgramTree::new(runner.new_rng()))
    }
}

pub struct ProgramTree {
    _rng: TestRng,
    current: Program<Type>,
}

impl ProgramTree {
    fn new(mut rng: TestRng) -> Self {
        let mut items = vec![];
        let program_length = PROGRAM_LENGTH_DISTRIBUTION.sample(&mut rng);
        let mut env = ScopedMap::new();

        for _ in 0..program_length {
            match STATEMENT_TYPE_FREQUENCIES[STATEMENT_TYPE_DISTRIBUTION.sample(&mut rng)].0 {
                StatementType::Binding => {
                    let binding = generate_random_binding(&mut rng, &mut env);
                    items.push(TopLevel::Statement(binding));
                }
                StatementType::Expression => {
                    let expr = generate_random_expression(&mut rng, &mut env);
                    items.push(TopLevel::Statement(expr));
                }
                StatementType::Function => {
                    env.new_scope();
                    let name = generate_random_name(&mut rng);
                    let mut args = vec![];
                    let arg_count = FUNCTION_ARGUMENTS_DISTRIBUTION.sample(&mut rng);
                    for _ in 0..arg_count {
                        let name = generate_random_name(&mut rng);
                        let ty = generate_random_argument_type(&mut rng);
                        args.push((name, ty));
                    }
                    let body = generate_random_expression(&mut rng, &mut env);
                    let rettype = body.type_of();
                    env.release_scope();
                    items.push(TopLevel::Function(name, args, rettype, body))
                }
            }
        }

        let current = Program { items };

        ProgramTree { _rng: rng, current }
    }
}

impl ValueTree for ProgramTree {
    type Value = Program<Type>;

    fn current(&self) -> Self::Value {
        self.current.clone()
    }

    fn simplify(&mut self) -> bool {
        unimplemented!()
    }

    fn complicate(&mut self) -> bool {
        unimplemented!()
    }
}

#[derive(Debug)]
struct ExpressionGenerator {}

impl Strategy for ExpressionGenerator {
    type Tree = ExpressionTree;
    type Value = Expression<Type>;

    fn new_tree(&self, runner: &mut TestRunner) -> NewTree<Self> {
        NewTree::<ExpressionGenerator>::Ok(ExpressionTree::new(runner.new_rng()))
    }
}

struct ExpressionTree {
    _rng: TestRng,
    current: Expression<Type>,
}

impl ValueTree for ExpressionTree {
    type Value = Expression<Type>;

    fn current(&self) -> Self::Value {
        self.current.clone()
    }

    fn simplify(&mut self) -> bool {
        unimplemented!()
    }

    fn complicate(&mut self) -> bool {
        unimplemented!()
    }
}

impl ExpressionTree {
    fn new(mut rng: TestRng) -> Self {
        let mut env = ScopedMap::new();
        let current = generate_random_expression(&mut rng, &mut env);

        ExpressionTree { _rng: rng, current }
    }
}

fn generate_random_expression(
    rng: &mut TestRng,
    env: &mut ScopedMap<Variable, Type>,
) -> Expression<Type> {
    match EXPRESSION_TYPE_FREQUENCIES[EXPRESSION_TYPE_DISTRIBUTION.sample(rng)].0 {
        ExpressionType::Atomic => Expression::Atomic(generate_random_valueref(rng, env, None)),

        ExpressionType::Bind => generate_random_binding(rng, env),

        ExpressionType::Block => {
            let num_stmts = BLOCK_LENGTH_DISTRIBUTION.sample(rng);
            let mut stmts = Vec::new();
            let mut last_type = Type::Primitive(PrimitiveType::Void);

            env.new_scope();
            for _ in 0..num_stmts {
                let next = generate_random_expression(rng, env);
                last_type = next.type_of();
                stmts.push(next);
            }
            env.release_scope();

            Expression::Block(Location::manufactured(), last_type, stmts)
        }

        ExpressionType::Cast => {
            let inner = generate_random_valueref(rng, env, None);

            match inner.type_of() {
                // nevermind
                Type::Function(_, _) => Expression::Atomic(inner),
                Type::Primitive(primty) => {
                    let to_type = primty
                        .allowed_casts()
                        .choose(rng)
                        .expect("actually chose type");
                    Expression::Cast(Location::manufactured(), Type::Primitive(*to_type), inner)
                }
            }
        }

        ExpressionType::Primitive => {
            let base_expr = generate_random_valueref(rng, env, None);
            let out_type = base_expr.type_of();

            match out_type {
                Type::Function(_, _) => Expression::Atomic(base_expr),
                Type::Primitive(primty) => match primty.valid_operators().choose(rng) {
                    None => Expression::Atomic(base_expr),
                    Some((operator, arg_count)) => {
                        let primop = Primitive::from_str(operator).expect("chose valid primitive");
                        let mut args = vec![base_expr];

                        while args.len() < *arg_count {
                            args.push(generate_random_valueref(rng, env, Some(primty)));
                        }

                        Expression::Primitive(Location::manufactured(), out_type, primop, args)
                    }
                },
            }
        }

        ExpressionType::Print => {
            let possible_variables = env
                .bindings()
                .iter()
                .filter_map(|(variable, ty)| {
                    if ty.is_printable() {
                        Some(variable.clone())
                    } else {
                        None
                    }
                })
                .collect::<Vec<Variable>>();

            if possible_variables.is_empty() {
                generate_random_binding(rng, env)
            } else {
                Expression::Print(
                    Location::manufactured(),
                    possible_variables.choose(rng).unwrap().clone(),
                )
            }
        }
    }
}

fn generate_random_binding(
    rng: &mut TestRng,
    env: &mut ScopedMap<Variable, Type>,
) -> Expression<Type> {
    let name = generate_random_name(rng);
    let expr = generate_random_expression(rng, env);
    let ty = expr.type_of();
    env.insert(name.clone(), ty.clone());
    Expression::Bind(Location::manufactured(), name, ty, Box::new(expr))
}

fn generate_random_valueref(
    rng: &mut TestRng,
    env: &mut ScopedMap<Variable, Type>,
    target_type: Option<PrimitiveType>,
) -> ValueOrRef<Type> {
    let mut bindings = env.bindings();

    bindings.retain(|_, value| {
        target_type
            .map(|x| value == &Type::Primitive(x))
            .unwrap_or(true)
    });

    if rng.gen() || bindings.is_empty() {
        let value_type = if let Some(target_type) = target_type {
            ValueType::from(target_type)
        } else {
            VALUE_TYPE_FREQUENCIES[VALUE_TYPE_DISTRIBUTION.sample(rng)].0
        };

        // generate a constant
        let val = match value_type {
            ValueType::I8 => Value::I8(None, rng.gen()),
            ValueType::I16 => Value::I16(None, rng.gen()),
            ValueType::I32 => Value::I32(None, rng.gen()),
            ValueType::I64 => Value::I64(None, rng.gen()),
            ValueType::U8 => Value::U8(None, rng.gen()),
            ValueType::U16 => Value::U16(None, rng.gen()),
            ValueType::U32 => Value::U32(None, rng.gen()),
            ValueType::U64 => Value::U64(None, rng.gen()),
            ValueType::Void => Value::Void,
        };

        ValueOrRef::Value(Location::manufactured(), val.type_of(), val)
    } else {
        // generate a reference
        let weighted_keys = bindings.keys().map(|x| (1, x)).collect::<Vec<_>>();
        let distribution = WeightedIndex::new(weighted_keys.iter().map(|x| x.0)).unwrap();
        let var = weighted_keys[distribution.sample(rng)].1.clone();
        let ty = bindings
            .remove(&var)
            .expect("chose unbound variable somehow?");
        ValueOrRef::Ref(Location::manufactured(), ty, var)
    }
}

fn generate_random_name(rng: &mut TestRng) -> Variable {
    let start = rng.gen_range('a'..='z');
    crate::ir::gensym(&format!("{}", start))
}

fn generate_random_argument_type(rng: &mut TestRng) -> Type {
    ARGUMENT_TYPE_FREQUENCIES[ARGUMENT_TYPE_DISTRIBUTION.sample(rng)]
        .0
        .clone()
}