morph-cli 0.1.0

#![allow(dead_code, clippy::all, unreachable_patterns)]
use swc_common::{FileName, SourceMap, Spanned, sync::Lrc};
use swc_ecma_ast::*;
use swc_ecma_parser::{Parser, StringInput, Syntax, TsSyntax, lexer::Lexer};

#[derive(Debug, Clone)]
pub enum InferredType {
    String,
    Number,
    Boolean,
    Null,
    Undefined,
    StringLiteral(String),
    NumberLiteral(f64),
    BooleanLiteral(bool),
    Array(Box<InferredType>),
    Object(Vec<(String, InferredType)>),
    /// React props object — emitted as `Record<string, unknown>`.
    ReactProps,
    Function {
        params: Vec<InferredType>,
        ret: Box<InferredType>,
    },
    Unknown,
}

/// Parse a JSDoc type string like `{string}` or `{number[]}` into an InferredType.
fn parse_jsdoc_type(raw: &str) -> InferredType {
    let t = raw.trim().trim_start_matches('{').trim_end_matches('}').trim();
    match t {
        "string" => InferredType::String,
        "number" => InferredType::Number,
        "boolean" => InferredType::Boolean,
        "null" => InferredType::Null,
        "undefined" => InferredType::Undefined,
        s if s.ends_with("[]") => {
            let inner = parse_jsdoc_type(&s[..s.len() - 2]);
            InferredType::Array(Box::new(inner))
        }
        _ => InferredType::Unknown,
    }
}

/// Extract `@param {type} name` and `@returns {type}` entries from a JSDoc block.
pub fn extract_jsdoc_hints(source: &str, fn_start: usize) -> JsDocHints {
    let mut hints = JsDocHints::default();
    // Only look in the 400 bytes before the function start.
    let window_start = fn_start.saturating_sub(400);
    let window = &source[window_start..fn_start];
    // Find the last /** ... */ block.
    let Some(block_start) = window.rfind("/**") else { return hints; };
    let Some(block_end) = window[block_start..].find("*/") else { return hints; };
    let block = &window[block_start..block_start + block_end + 2];

    for line in block.lines() {
        let trimmed = line.trim().trim_start_matches('*').trim();
        if trimmed.starts_with("@param") {
            // @param {type} name
            if let Some(brace_start) = trimmed.find('{') {
                if let Some(brace_end) = trimmed[brace_start..].find('}') {
                    let type_str = &trimmed[brace_start..brace_start + brace_end + 1];
                    let after = trimmed[brace_start + brace_end + 1..].trim();
                    let param_name = after.split_whitespace().next().unwrap_or("").to_string();
                    hints.params.push((param_name, parse_jsdoc_type(type_str)));
                }
            }
        } else if trimmed.starts_with("@returns") || trimmed.starts_with("@return") {
            if let Some(brace_start) = trimmed.find('{') {
                if let Some(brace_end) = trimmed[brace_start..].find('}') {
                    let type_str = &trimmed[brace_start..brace_start + brace_end + 1];
                    hints.ret = Some(parse_jsdoc_type(type_str));
                }
            }
        }
    }
    hints
}

#[derive(Default)]
pub struct JsDocHints {
    pub params: Vec<(String, InferredType)>,
    pub ret: Option<InferredType>,
}

/// Unify a list of inferred types to their common supertype.
fn unify_types(types: &[InferredType]) -> InferredType {
    if types.is_empty() {
        return InferredType::Unknown;
    }
    let first = &types[0];
    let all_same = types.iter().all(|t| std::mem::discriminant(t) == std::mem::discriminant(first));
    if all_same {
        // Widen literals to primitives.
        match first {
            InferredType::StringLiteral(_) => InferredType::String,
            InferredType::NumberLiteral(_) => InferredType::Number,
            InferredType::BooleanLiteral(_) => InferredType::Boolean,
            other => other.clone(),
        }
    } else {
        InferredType::Unknown
    }
}

/// Heuristic: infer a parameter type from its name.
fn infer_param_name(name: &str) -> Option<InferredType> {
    match name {
        // React-specific
        "props" => Some(InferredType::ReactProps),
        "children" => Some(InferredType::Unknown),
        // Callbacks
        n if n.starts_with("on") && n.len() > 2 => Some(InferredType::Function {
            params: vec![InferredType::Unknown],
            ret: Box::new(InferredType::Undefined),
        }),
        n if n.starts_with("handle") && n.len() > 6 => Some(InferredType::Function {
            params: vec![InferredType::Unknown],
            ret: Box::new(InferredType::Undefined),
        }),
        "callback" | "cb" | "done" | "next" => Some(InferredType::Function {
            params: vec![InferredType::Unknown],
            ret: Box::new(InferredType::Undefined),
        }),
        // Common primitives by convention
        "id" | "key" => Some(InferredType::String),
        "count" | "index" | "idx" | "size" | "len" | "length" | "offset" | "page" => {
            Some(InferredType::Number)
        }
        "enabled" | "disabled" | "visible" | "hidden" | "active" | "checked" | "loading" => {
            Some(InferredType::Boolean)
        }
        "name" | "label" | "title" | "message" | "text" | "url" | "path" | "src" | "href" => {
            Some(InferredType::String)
        }
        _ => None,
    }
}

impl InferredType {
    #[allow(dead_code)]
    pub fn as_ts_string(&self) -> String {
        match self {
            Self::String => "string".into(),
            Self::Number => "number".into(),
            Self::Boolean => "boolean".into(),
            Self::Null => "null".into(),
            Self::Undefined => "undefined".into(),
            Self::StringLiteral(_) => "string".into(),
            Self::NumberLiteral(_) => "number".into(),
            Self::BooleanLiteral(_) => "boolean".into(),
            Self::Array(inner) => format!("{}[]", inner.as_ts_string()),
            Self::Object(fields) => {
                if fields.is_empty() {
                    "Record<string, unknown>".into()
                } else {
                    let props: Vec<String> = fields
                        .iter()
                        .map(|(k, v)| format!("{}: {}", k, v.as_ts_string()))
                        .collect();
                    format!("{{ {} }}", props.join("; "))
                }
            }
            Self::ReactProps => "Record<string, unknown>".into(),
            Self::Function { params, ret } => {
                let param_strs: Vec<String> = params
                    .iter()
                    .enumerate()
                    .map(|(i, p)| format!("arg{}: {}", i, p.as_ts_string()))
                    .collect();
                format!("({}) => {}", param_strs.join(", "), ret.as_ts_string())
            }
            Self::Unknown => "unknown".into(),
        }
    }
}

#[derive(Debug, Clone, Copy, PartialEq)]
pub enum Confidence {
    Low,
    Medium,
    High,
}

impl Confidence {
    #[allow(dead_code)]
    pub fn score(&self) -> u8 {
        match self {
            Self::Low => 33,
            Self::Medium => 66,
            Self::High => 100,
        }
    }
}

#[derive(Debug, Clone)]
pub struct InferenceResult {
    pub start: usize,
    pub end: usize,
    pub inferred: InferredType,
    pub confidence: Confidence,
    #[allow(dead_code)]
    pub unsafe_marker: bool,
}

pub struct InferenceEngine {
    strict: bool,
}

impl InferenceEngine {
    pub fn new(strict: bool) -> Self {
        Self { strict }
    }

    pub fn infer(&self, source: &str) -> Vec<InferenceResult> {
        let cm: Lrc<SourceMap> = Lrc::new(SourceMap::default());
        let fm = cm.new_source_file(FileName::Custom("inline".into()).into(), source.to_string());
        let syntax = if source.contains("interface ") || source.contains(": ") {
            Syntax::Typescript(TsSyntax::default())
        } else {
            Syntax::Es(Default::default())
        };
        let lexer = Lexer::new(syntax, Default::default(), StringInput::from(&*fm), None);
        let mut parser = Parser::new_from(lexer);
        parser
            .parse_module()
            .map(|m| self.infer_module(&m))
            .unwrap_or_default()
    }

    fn infer_module(&self, module: &Module) -> Vec<InferenceResult> {
        let mut results = Vec::new();
        for item in &module.body {
            if let ModuleItem::Stmt(Stmt::Decl(Decl::Var(v))) = item {
                results.extend(self.infer_var_decl(v));
            }
            if let ModuleItem::Stmt(Stmt::Decl(Decl::Fn(f))) = item {
                results.extend(self.infer_fn_decl(f));
            }
        }
        results
    }

    fn infer_var_decl(&self, decl: &VarDecl) -> Vec<InferenceResult> {
        let mut results = Vec::new();
        for declator in &decl.decls {
            if let Some(init) = &declator.init {
                if let Some((inferred, confidence)) = self.infer_expr(init) {
                    if confidence != Confidence::Low || !self.strict {
                        let span = declator.name.span();
                        let end = span.hi.0 as usize;
                        let start = span.lo.0 as usize;
                        let unsafe_marker = matches!(inferred, InferredType::Unknown)
                            || (confidence == Confidence::Low && self.strict);
                        if !unsafe_marker {
                            results.push(InferenceResult {
                                start,
                                end,
                                inferred,
                                confidence,
                                unsafe_marker,
                            });
                        }
                    }
                }
            }
        }
        results
    }

    fn infer_fn_decl(&self, decl: &FnDecl) -> Vec<InferenceResult> {
        self.infer_fn_decl_with_source(decl, "")
    }

    fn infer_fn_decl_with_source(&self, decl: &FnDecl, source: &str) -> Vec<InferenceResult> {
        let fn_start = decl.function.span.lo.0 as usize;
        let jsdoc = if !source.is_empty() {
            extract_jsdoc_hints(source, fn_start.saturating_sub(1))
        } else {
            JsDocHints::default()
        };

        // Build a map of JSDoc param types keyed by name.
        let jsdoc_params: std::collections::HashMap<&str, &InferredType> = jsdoc
            .params
            .iter()
            .map(|(n, t)| (n.as_str(), t))
            .collect();

        let params: Vec<InferredType> = decl
            .function
            .params
            .iter()
            .map(|p| {
                if let Pat::Ident(ident) = &p.pat {
                    let name = ident.sym.as_ref();
                    // 1. JSDoc wins
                    if let Some(jt) = jsdoc_params.get(name) {
                        return (*jt).clone();
                    }
                    // 2. Name heuristics
                    if let Some(ht) = infer_param_name(name) {
                        return ht;
                    }
                }
                InferredType::Unknown
            })
            .collect();

        // Collect all return-statement expressions and unify.
        let ret_type = decl
            .function
            .body
            .as_ref()
            .map(|body| {
                let mut ret_types: Vec<InferredType> = Vec::new();
                for stmt in &body.stmts {
                    if let Stmt::Return(ret) = stmt {
                        if let Some(arg) = &ret.arg {
                            if let Some((t, _)) = self.infer_expr(arg) {
                                ret_types.push(t);
                            }
                        } else {
                            ret_types.push(InferredType::Undefined);
                        }
                    }
                }
                if ret_types.is_empty() {
                    jsdoc.ret.clone().unwrap_or(InferredType::Undefined)
                } else {
                    let unified = unify_types(&ret_types);
                    jsdoc.ret.clone().unwrap_or(unified)
                }
            })
            .unwrap_or(InferredType::Unknown);

        let all_params_known = params.iter().all(|p| !matches!(p, InferredType::Unknown));
        let ret_known = !matches!(ret_type, InferredType::Unknown | InferredType::Undefined);
        let confidence = if all_params_known && ret_known {
            Confidence::High
        } else if all_params_known || ret_known {
            Confidence::Medium
        } else {
            Confidence::Low
        };

        if confidence != Confidence::Low || !self.strict {
            let span = decl.function.span;
            return vec![InferenceResult {
                start: span.lo.0 as usize,
                end: span.hi.0 as usize,
                inferred: InferredType::Function {
                    params,
                    ret: Box::new(ret_type),
                },
                confidence,
                unsafe_marker: confidence == Confidence::Low && self.strict,
            }];
        }
        Vec::new()
    }

    fn infer_expr(&self, expr: &Expr) -> Option<(InferredType, Confidence)> {
        match expr {
            Expr::Lit(lit) => match lit {
                Lit::Str(s) => Some((
                    InferredType::StringLiteral(s.value.to_string()),
                    Confidence::High,
                )),
                Lit::Num(n) => Some((InferredType::NumberLiteral(n.value), Confidence::High)),
                Lit::Bool(b) => Some((InferredType::BooleanLiteral(b.value), Confidence::High)),
                Lit::Null(_) => Some((InferredType::Null, Confidence::High)),
                Lit::BigInt(_) => Some((InferredType::Number, Confidence::Medium)),
                Lit::Regex(_) => Some((InferredType::String, Confidence::Medium)),
                Lit::JSXText(_) => Some((InferredType::String, Confidence::High)),
                _ => None,
            },
            Expr::Ident(ident) => {
                let name = ident.sym.as_ref();
                match name {
                    "undefined" => Some((InferredType::Undefined, Confidence::High)),
                    "null" => Some((InferredType::Null, Confidence::High)),
                    "true" | "false" => Some((InferredType::Boolean, Confidence::High)),
                    _ => Some((InferredType::Unknown, Confidence::Low)),
                }
            }
            Expr::Array(arr) => {
                // Collect types of all non-spread elements and unify.
                let elem_types: Vec<InferredType> = arr
                    .elems
                    .iter()
                    .filter_map(|e| e.as_ref())
                    .filter(|e| e.spread.is_none())
                    .filter_map(|e| self.infer_expr(&e.expr).map(|(t, _)| t))
                    .collect();
                let inner = if elem_types.is_empty() {
                    InferredType::Unknown
                } else {
                    unify_types(&elem_types)
                };
                let conf = if matches!(inner, InferredType::Unknown) {
                    Confidence::Low
                } else {
                    Confidence::Medium
                };
                Some((InferredType::Array(Box::new(inner)), conf))
            }
            Expr::Object(obj) => {
                let mut fields = Vec::new();
                for prop in &obj.props {
                    match prop {
                        PropOrSpread::Prop(box_prop) => match &**box_prop {
                            Prop::KeyValue(kv) => {
                                let key = match &kv.key {
                                    PropName::Str(s) => s.value.to_string(),
                                    PropName::Ident(i) => i.sym.as_ref().to_string(),
                                    PropName::Num(n) => n.value.to_string(),
                                    _ => continue,
                                };
                                let val_type = self
                                    .infer_expr(&kv.value)
                                    .map(|(t, _)| t)
                                    .unwrap_or(InferredType::Unknown);
                                fields.push((key, val_type));
                            }
                            // shorthand: { foo } — value type is unknown without scope
                            Prop::Shorthand(ident) => {
                                let key = ident.sym.as_ref().to_string();
                                let val_type = infer_param_name(&key)
                                    .unwrap_or(InferredType::Unknown);
                                fields.push((key, val_type));
                            }
                            _ => {}
                        },
                        PropOrSpread::Spread(_) => {
                            // spread makes the shape unknown; keep what we have
                        }
                    }
                }
                let conf = if fields.is_empty() { Confidence::Low } else { Confidence::High };
                Some((InferredType::Object(fields), conf))
            }
            Expr::Bin(bin) => match bin.op {
                BinaryOp::Add | BinaryOp::Sub | BinaryOp::Mul | BinaryOp::Div => {
                    Some((InferredType::Number, Confidence::Medium))
                }
                BinaryOp::EqEq | BinaryOp::NotEq | BinaryOp::EqEqEq | BinaryOp::NotEqEq => {
                    Some((InferredType::Boolean, Confidence::High))
                }
                BinaryOp::LogicalAnd | BinaryOp::LogicalOr => self.infer_expr(&bin.right),
                _ => Some((InferredType::Unknown, Confidence::Low)),
            },
            Expr::Unary(unary) => {
                #[allow(deprecated)]
                let (inferred, conf) = match unary.op {
                    UnaryOp::Minus | UnaryOp::Plus => (InferredType::Number, Confidence::Medium),
                    UnaryOp::Void | UnaryOp::TypeOf => {
                        (InferredType::Undefined, Confidence::Medium)
                    }
                    UnaryOp::Bang | UnaryOp::Tilde => (InferredType::Boolean, Confidence::Medium),
                    UnaryOp::Delete => (InferredType::Unknown, Confidence::Low),
                };
                Some((inferred, conf))
            }
            Expr::Tpl(tpl) => {
                let conf = if tpl.quasis.is_empty() {
                    Confidence::High
                } else {
                    Confidence::Medium
                };
                Some((InferredType::String, conf))
            }
            _ => Some((InferredType::Unknown, Confidence::Low)),
        }
    }
}

#[derive(Default)]
pub struct InferenceStats {
    pub total_inferred: usize,
    pub skipped: usize,
    pub uncertain: usize,
    pub by_type: std::collections::HashMap<String, usize>,
}

impl InferenceStats {
    pub fn new() -> Self {
        Self::default()
    }

    #[allow(dead_code)]
    pub fn add(&mut self, result: &InferenceResult) {
        self.total_inferred += 1;
        if result.confidence == Confidence::Low {
            self.uncertain += 1;
        }
        let type_str = result.inferred.as_ts_string();
        *self.by_type.entry(type_str).or_insert(0) += 1;
    }

    pub fn print_summary(&self) {
        println!();
        println!("  Inferred: {}", self.total_inferred);
        println!("  Skipped: {}", self.skipped);
        println!("  Uncertain: {}", self.uncertain);
        if !self.by_type.is_empty() {
            print!("  Types: ");
            let types: Vec<String> = self
                .by_type
                .iter()
                .map(|(k, v)| format!("{}({})", k, v))
                .collect();
            println!("{}", types.join(", "));
        }
    }
}

#[cfg(test)]
mod tests {
    use super::*;

    #[test]
    fn test_string() {
        let e = InferenceEngine::new(false);
        let r = e.infer("const x = 'hello';");
        assert!(!r.is_empty());
    }

    #[test]
    fn test_number() {
        let e = InferenceEngine::new(false);
        let r = e.infer("const x = 42;");
        assert!(!r.is_empty());
    }

    #[test]
    fn test_boolean() {
        let e = InferenceEngine::new(false);
        let r = e.infer("const x = true;");
        assert!(!r.is_empty());
    }

    #[test]
    fn test_null() {
        let e = InferenceEngine::new(false);
        let r = e.infer("const x = null;");
        assert!(!r.is_empty());
    }

    #[test]
    fn test_array() {
        let e = InferenceEngine::new(false);
        let r = e.infer("const x = [1, 2, 3];");
        assert!(!r.is_empty());
    }

    #[test]
    fn test_object() {
        let e = InferenceEngine::new(false);
        let r = e.infer("const x = { a: 1, b: 'hi' };");
        assert!(!r.is_empty());
    }

    #[test]
    fn test_ts_string() {
        assert_eq!(InferredType::Number.as_ts_string(), "number");
        assert_eq!(
            InferredType::Array(Box::new(InferredType::Number)).as_ts_string(),
            "number[]"
        );
    }

    #[test]
    fn test_strict_skips_low() {
        let e = InferenceEngine::new(true);
        assert!(e.infer("const x = fn();").is_empty());
    }

    #[test]
    fn test_uniform_array_infers_element_type() {
        let e = InferenceEngine::new(false);
        let r = e.infer("const x = [1, 2, 3];");
        assert!(!r.is_empty());
        assert_eq!(r[0].inferred.as_ts_string(), "number[]");
    }

    #[test]
    fn test_mixed_array_infers_unknown() {
        let e = InferenceEngine::new(false);
        let r = e.infer("const x = [1, 'two', true];");
        // Mixed types → unknown[]
        assert!(!r.is_empty());
        assert_eq!(r[0].inferred.as_ts_string(), "unknown[]");
    }

    #[test]
    fn test_empty_object_emits_record() {
        assert_eq!(InferredType::Object(vec![]).as_ts_string(), "Record<string, unknown>");
    }

    #[test]
    fn test_react_props_emits_record() {
        assert_eq!(InferredType::ReactProps.as_ts_string(), "Record<string, unknown>");
    }

    #[test]
    fn test_function_ts_string() {
        let t = InferredType::Function {
            params: vec![InferredType::String],
            ret: Box::new(InferredType::Number),
        };
        assert!(t.as_ts_string().contains("string"));
        assert!(t.as_ts_string().contains("number"));
    }

    #[test]
    fn test_jsdoc_param_parsing() {
        let source = "/**\n * @param {string} name\n * @returns {number}\n */\nfunction greet(name) { return name.length; }";
        let hints = extract_jsdoc_hints(source, source.find("function").unwrap());
        assert_eq!(hints.params.len(), 1);
        assert_eq!(hints.params[0].0, "name");
        assert_eq!(hints.params[0].1.as_ts_string(), "string");
        assert!(hints.ret.is_some());
        assert_eq!(hints.ret.unwrap().as_ts_string(), "number");
    }

    #[test]
    fn test_param_name_heuristic_callback() {
        let t = infer_param_name("onChange").expect("should infer");
        assert!(t.as_ts_string().contains("=>"));
    }

    #[test]
    fn test_param_name_heuristic_props() {
        let t = infer_param_name("props").expect("should infer");
        assert_eq!(t.as_ts_string(), "Record<string, unknown>");
    }

    #[test]
    fn test_param_name_heuristic_id() {
        let t = infer_param_name("id").expect("should infer");
        assert_eq!(t.as_ts_string(), "string");
    }

    #[test]
    fn test_jsdoc_array_type() {
        let t = parse_jsdoc_type("{string[]}");
        assert_eq!(t.as_ts_string(), "string[]");
    }
}