//! The program, as C: its types, its statics and its functions
//!
//! `Program::new` finds all of them before any C is written. Writing a function only reads
//! the program, so the functions can be written in any order.

use std::collections::{BTreeMap, BTreeSet};

use super::functions::Function;
use super::types::{CType, ClassField, ClassKind, Member, Types};
use crate::{
    flatten::ModuleBindings,
    infer::{LENGTH_FIELD, Type, TypeCache, join, module_binding},
    libqjs::parse::{Atom, AtomRef, JsValue, opcodes::OpCode},
    regify::{Operand, RegFn, RegInstr},
    transform::{PROTO_FIELD, PROTOTYPE_FIELD},
};

/// Strings that the C code can need for any program: the names that `typeof` gives, and the
/// results of converting undefined, null and the bools to strings
const FIXED_STRINGS: &[&str] = &[
    "undefined",
    "null",
    "true",
    "false",
    "boolean",
    "number",
    "string",
    "function",
    "object",
];

pub struct Program<'a> {
    // The input: the JS functions, and what infer found
    pub fns: &'a [RegFn],
    pub cache: &'a TypeCache,
    pub atoms: &'a [Atom],
    bindings: &'a ModuleBindings,
    pub entry: usize, // the walk of the entry: the JS function that runs the program

    // The output
    pub types: Types,
    pub strings: Vec<String>, // statics: the text of each string literal
    pub globals: BTreeMap<usize, CType>, // statics: the module bindings, which are locals of the entry, and their C types
    pub functions: Vec<Function>,
}

impl<'a> Program<'a> {
    pub fn new(
        fns: &'a [RegFn],
        cache: &'a TypeCache,
        atoms: &'a [Atom],
        bindings: &'a ModuleBindings,
    ) -> Self {
        let mut program = Program {
            fns,
            cache,
            atoms,
            bindings,
            entry: cache.find(fns.len() - 1, None, &[], fns),
            types: Types::new(&cache.memory),
            strings: vec![],
            globals: BTreeMap::new(),
            functions: vec![],
        };

        program.add_globals();
        program.add_functions();
        program.add_strings();
        program.add_unions();
        program
    }

    /// Each module binding is a C global, so that any function can use the stdlib's
    fn add_globals(&mut self) {
        let entry = self.fns.len() - 1;

        for &local in self.bindings.values() {
            let ty = self.cache.memory.get(module_binding(entry, local));
            let ctype = self.types.add(ty);
            self.globals.insert(local, ctype);
        }
    }

    /// One C function for each walk of a JS function in infer's last round, one for each
    /// lookup, and the host's two entry points
    fn add_functions(&mut self) {
        let walks = (0..self.cache.walks.len()).filter(|&walk| self.cache.is_live(walk));
        let mut lookups = BTreeSet::new();

        for walk in walks.clone() {
            self.add_lookups(walk, &mut lookups);
        }

        let lookups = lookups
            .into_iter()
            .map(|(class, field)| Function::Lookup(class, field));
        self.functions.extend(walks.map(Function::Js));
        self.functions.extend(lookups);
        self.functions.extend([Function::Start, Function::Complete]);
    }

    /// The lookups that a JS function's walk needs: one for each field that it gets, on each
    /// class that can lack the field
    fn add_lookups(&self, walk: usize, lookups: &mut BTreeSet<(usize, AtomRef)>) {
        for instr in self.instrs(walk) {
            let field = match instr.op {
                OpCode::GetField => instr.imm.atom(),
                OpCode::GetLength => LENGTH_FIELD,
                _ => continue,
            };

            let ty = self.slot_type(walk, instr.args[0]);

            for member in self.types.type_members(&ty) {
                if let Some(class) = self.field_class(member, field) {
                    self.add_lookup(class, field, lookups);
                }
            }
        }
    }

    /// The lookup of `field` on `class`, if its objects can lack the field. If they can, the
    /// lookup goes on to their prototypes, so those need lookups too
    fn add_lookup(&self, class: usize, field: AtomRef, lookups: &mut BTreeSet<(usize, AtomRef)>) {
        let defined = self.types.field(class, field).is_some_and(|f| f.defined);

        if defined || !lookups.insert((class, field)) {
            return;
        }

        for proto_class in self.proto_classes(class) {
            self.add_lookup(proto_class, field, lookups);
        }
    }

    /// The text of every string literal: the fixed ones, and the ones in the JS functions
    fn add_strings(&mut self) {
        let mut strings: Vec<String> = FIXED_STRINGS.iter().map(|s| s.to_string()).collect();

        for &function in &self.functions {
            let Function::Js(walk) = function else {
                continue;
            };

            let cpool = &self.fns[self.cache.walks[walk].req.fn_idx].cpool;

            for instr in self.instrs(walk) {
                let text = match instr.op {
                    OpCode::PushAtomValue => instr.imm.atom().text(self.atoms),
                    OpCode::PushEmptyString => String::new(),
                    OpCode::PushConst | OpCode::PushConst8 => match &cpool[instr.imm.const_idx()] {
                        JsValue::String(s) => s.text(),
                        _ => continue,
                    },
                    _ => continue,
                };

                strings.push(text);
            }
        }

        for text in strings {
            if !self.strings.contains(&text) {
                self.strings.push(text);
            }
        }
    }

    /// The union of every C type that the functions use, if it has several members. The
    /// classes' fields have theirs already (see `Types::new`)
    fn add_unions(&mut self) {
        let types: Vec<Type> = (self.functions.iter())
            .flat_map(|function| self.value_types(function))
            .collect();

        for ty in &types {
            self.types.add(ty);
        }
    }

    /// The types of the values in a function: the locals, args and regs of a JS function,
    /// and what it returns, or what a lookup gives
    fn value_types(&self, function: &Function) -> Vec<Type> {
        match *function {
            Function::Js(walk) => {
                let f = &self.fns[self.cache.walks[walk].req.fn_idx];
                let slots = slots(f).into_iter().map(|slot| self.slot_type(walk, slot));
                slots.chain([self.cache.walks[walk].ret.clone()]).collect()
            }
            Function::Lookup(class, field) => vec![self.lookup_type(class, field)],
            Function::Start | Function::Complete => vec![],
        }
    }

    /// The instrs of a JS function that can run in a walk
    fn instrs(&self, walk: usize) -> impl Iterator<Item = &'a RegInstr> + use<'a> {
        let (cache, fns) = (self.cache, self.fns);
        let walk = &cache.walks[walk];
        let f = &fns[walk.req.fn_idx];
        let blocks = f.blocks.iter().zip(&walk.reachable);
        let blocks = blocks.filter(|(_, reachable)| **reachable);
        blocks.flat_map(|(block, _)| &f.code[block.start..block.end])
    }
}

// ────────────────────────────────────── C values ───────────────────────────────────────

/// One combination of members of some operands. When an operand is a union, the C code tests
/// its tag, and then uses the member that the tag says
#[derive(Clone)]
pub struct Case {
    pub operands: Vec<(String, Member)>, // each operand, as C, as a value of one member
    pub tests: Vec<String>, // the C tests that the operands are these members: one for each union
    pub first: bool,        // the first case of its operands: an `if`. The others are `else if`
}

impl Program<'_> {
    /// The type of a local, arg or reg of a JS function, in one walk
    pub fn slot_type(&self, walk: usize, slot: Operand) -> Type {
        let addr = self.cache.addr(walk, slot, self.fns);
        self.cache.memory.get(addr).clone()
    }

    /// The C type of a local, arg or reg of a JS function, in one walk
    pub fn slot_ctype(&self, walk: usize, slot: Operand) -> CType {
        self.types.of(&self.slot_type(walk, slot))
    }

    pub fn class_kind(&self, member: Member) -> ClassKind {
        match member {
            Member::Obj(class) => self.types.classes[class].kind,
            _ => ClassKind::Obj,
        }
    }

    /// Whether a value of `member` is a function
    pub fn is_function(&self, member: Member) -> bool {
        let closure = matches!(self.class_kind(member), ClassKind::Closure(_));
        closure || matches!(member, Member::Native(_))
    }

    /// Every combination of the members of `operands`. Each operand is C, and its C type
    pub fn cases(&self, operands: &[(String, CType)]) -> Vec<Case> {
        let mut cases = vec![Case {
            operands: vec![],
            tests: vec![],
            first: true,
        }];

        for (value, ctype) in operands {
            let mut next = vec![];

            for case in &cases {
                for (tag, member) in self.types.members_of(*ctype).into_iter().enumerate() {
                    let mut case = case.clone();

                    match ctype {
                        CType::One(_) => case.operands.push((value.clone(), member)),
                        CType::Union(_) => {
                            case.operands.push((format!("{value}.as.m{tag}"), member));
                            case.tests.push(format!("{value}.tag == {tag}"));
                        }
                    }

                    case.first = next.is_empty();
                    next.push(case);
                }
            }

            cases = next;
        }

        cases
    }

    /// `value`, of `member`, as a value of C type `to`. Types only grow, so `to` holds
    /// `member`, or holds a float if `member` is an int
    pub fn convert(&self, value: &str, member: Member, to: CType) -> String {
        let members = self.types.members_of(to);
        let same = members.iter().position(|&m| m == member);
        let float = members.iter().position(|&m| m == Member::Float);
        let float = float.filter(|_| member == Member::Int);

        let Some(tag) = same.or(float) else {
            panic!("emit: a value of {member:?} where {members:?} was inferred");
        };

        let value = match members[tag] == member {
            true => value.to_string(),
            false => format!("(double){value}"),
        };

        match to {
            CType::One(_) => value,
            CType::Union(id) => format!("(u{id}){{.tag = {tag}, .as.m{tag} = {value}}}"),
        }
    }

    /// A string literal, as C: a pointer to its static
    pub fn string(&self, text: &str) -> String {
        match self.strings.iter().position(|s| s == text) {
            Some(id) => format!("(&str{id})"),
            None => panic!("emit: no static for the string {text:?}"),
        }
    }

    /// `value`, of `member`, as a string, as JS's ToString does. `None` for an object: its
    /// conversion calls its methods, which is not supported
    pub fn to_string(&self, value: &str, member: Member) -> Option<String> {
        Some(match member {
            Member::Undefined => self.string("undefined"),
            Member::Null => self.string("null"),
            Member::Bool => {
                let (t, f) = (self.string("true"), self.string("false"));
                format!("({value} ? {t} : {f})")
            }
            Member::Int => format!("__rt_number_to_string((double){value})"),
            Member::Float => format!("__rt_number_to_string({value})"),
            Member::Str => value.into(),
            Member::Obj(_) | Member::Native(_) => return None,
        })
    }

    /// The value of the stdlib's module binding `name`, as C, and its C type
    pub fn binding(&self, name: &str) -> Option<(String, CType)> {
        let local = *self.bindings.get(name)?;
        Some((format!("(*g{local})"), self.globals[&local]))
    }
}

// ─────────────────────────────────────── Fields ────────────────────────────────────────

impl Program<'_> {
    /// The C name of a field
    pub fn field_name(&self, field: AtomRef) -> String {
        let text: String = (field.text(self.atoms).chars().take(24))
            .map(|c| if c.is_ascii_alphanumeric() { c } else { '_' })
            .collect();

        match field {
            AtomRef::TaggedInt(n) => format!("f_{text}_i{n}"),
            AtomRef::Predefined(n) => format!("f_{text}_p{n}"),
            AtomRef::FileLocal(n) => format!("f_{text}_a{n}"),
        }
    }

    /// A field of `obj`, as C to get or set it. A field that shares a location holds a
    /// pointer to it
    pub fn field_access(&self, obj: &str, field: &ClassField) -> String {
        let name = self.field_name(field.name);

        match field.shared {
            true => format!("(*{obj}->{name})"),
            false => format!("{obj}->{name}"),
        }
    }

    /// The value of `field` on `obj`, an object of `class`, and its C type. If every object
    /// of the class has the field, that is the field itself. Else it is a call of the field's
    /// lookup
    pub fn field_value(&self, class: usize, field: AtomRef, obj: &str) -> (String, CType) {
        match self.types.field(class, field) {
            Some(f) if f.defined => (self.field_access(obj, f), f.ctype),
            _ => {
                let lookup = Function::Lookup(class, field).name(self);
                let ctype = self.types.of(&self.lookup_type(class, field));
                (format!("{lookup}({obj})"), ctype)
            }
        }
    }

    /// The class of the object that has `field`, for a value of `member`. That is the value
    /// itself if it is an object, or the stdlib's prototype of its kind if it is a primitive
    /// (see `builtin_proto`). A string's length and an array's are not fields, so they have
    /// none
    pub fn field_class(&self, member: Member, field: AtomRef) -> Option<usize> {
        let length = field == LENGTH_FIELD;

        match (member, self.class_kind(member)) {
            (Member::Str, _) | (Member::Obj(_), ClassKind::Array) if length => None,
            (Member::Obj(class), _) => Some(class),
            (Member::Undefined | Member::Null | Member::Native(_), _) => None,
            _ => self.builtin_proto(member).map(|(_, class)| class),
        }
    }

    /// The type that the lookup of `field` on `class` gives: the join of the field's types on
    /// the class's objects, and on their prototypes (see `Memory::read_field`)
    pub fn lookup_type(&self, class: usize, field: AtomRef) -> Type {
        let objs = &self.types.classes[class].objs;
        let memory = &self.cache.memory;
        let types = objs
            .iter()
            .map(|&obj| memory.read_field(&Type::Ref(obj), field));
        types.fold(Type::Unknown, |all, ty| join(&all, &ty))
    }

    /// The prototype of a primitive or an array, as C, and its class: the `prototype` field of
    /// the stdlib's String, Number, Boolean or Array (see `Memory::builtin_proto`)
    pub fn builtin_proto(&self, member: Member) -> Option<(String, usize)> {
        let name = match (member, self.class_kind(member)) {
            (Member::Str, _) => "String",
            (Member::Int | Member::Float, _) => "Number",
            (Member::Bool, _) => "Boolean",
            (_, ClassKind::Array) => "Array",
            _ => return None,
        };

        let (stdlib_class, ctype) = self.binding(name)?;
        let proto = self.types.field(ctype.class()?, PROTOTYPE_FIELD)?;
        let value = self.field_access(&stdlib_class, proto);
        Some((value, proto.ctype.class()?))
    }

    /// The classes that the prototype of an object of `class` can be. An array's prototype
    /// is the stdlib's. Another object's is its `__proto__` field, if it has one
    pub fn proto_classes(&self, class: usize) -> Vec<usize> {
        if self.types.classes[class].kind == ClassKind::Array {
            let proto = self.builtin_proto(Member::Obj(class));
            return proto.map(|(_, class)| class).into_iter().collect();
        }

        let Some(proto) = self.types.field(class, PROTO_FIELD) else {
            return vec![];
        };

        let members = self.types.members_of(proto.ctype);
        let classes = members.into_iter().map(|member| CType::One(member).class());
        classes.flatten().collect()
    }
}

/// Every local, arg and reg of a JS function
pub fn slots(f: &RegFn) -> Vec<Operand> {
    let locals = (0..f.local_count).map(Operand::Local);
    let args = (0..f.arg_count).map(Operand::Arg);
    let regs = (0..f.reg_count).map(Operand::Reg);
    locals.chain(args).chain(regs).collect()
}