//! Phase 129-C: Post-if continuation lowering with post_k
//!
//! ## Responsibility
//!
//! - Lower `Seq([If, Post...])` patterns to join_k + post_k continuation
//! - join_k merges environments from then/else branches
//! - post_k executes post-if statements and returns
//! - PHI-free: all merging done via env arguments
//!
//! ## Contract
//!
//! - Input: StepTree with if-only pattern + post-if statements
//! - Output: JoinModule with:
//!   - main: condition check → TailCall(k_then/k_else, env)
//!   - k_then: then statements → TailCall(join_k, env_then)
//!   - k_else: else statements → TailCall(join_k, env_else)
//!   - join_k: merge → TailCall(post_k, merged_env)
//!   - post_k: post-if statements → Ret
//!
//! ## Scope
//!
//! - Post-if: Return(Variable) only (Phase 124/125/126 baseline)
//! - If body: Assign(int literal) only (Phase 128 baseline)
//! - Condition: minimal compare only (Phase 123 baseline)
//!
//! ## Fail-Fast
//!
//! - Out of scope → Ok(None) (fallback to legacy)
//! - In scope but conversion failed → Err (with freeze_with_hint in strict mode)

use super::env_layout::EnvLayout;
use super::legacy::LegacyLowerer;
use crate::mir::control_tree::step_tree::{StepNode, StepStmtKind, StepTree};
use crate::mir::join_ir::lowering::carrier_info::JoinFragmentMeta;
use crate::mir::join_ir::lowering::error_tags;
use crate::mir::join_ir::{ConstValue, JoinFunction, JoinFuncId, JoinInst, JoinModule, MirLikeInst};
use crate::mir::ValueId;
use std::collections::BTreeMap;

/// Box-First: Post-if continuation lowering with post_k
pub struct PostIfPostKBuilderBox;

impl PostIfPostKBuilderBox {
    /// Try to lower if-with-post pattern to Normalized JoinModule using post_k.
    ///
    /// Returns:
    /// - Ok(Some((module, meta))): Successfully lowered
    /// - Ok(None): Out of scope (fallback to legacy)
    /// - Err(msg): In scope but failed (internal error)
    pub fn lower(
        step_tree: &StepTree,
        env_layout: &EnvLayout,
    ) -> Result<Option<(JoinModule, JoinFragmentMeta)>, String> {
        // Extract if + post pattern
        let (prefix_nodes, if_node, post_nodes) = match Self::extract_if_with_post(&step_tree.root) {
            Some(v) => v,
            None => return Ok(None), // Not an if-with-post pattern
        };

        let env_fields = env_layout.env_fields();

        fn alloc_value_id(next_value_id: &mut u32) -> ValueId {
            let vid = ValueId(*next_value_id);
            *next_value_id += 1;
            vid
        }

        let alloc_env_params =
            |fields: &[String], next_value_id: &mut u32| -> Vec<ValueId> {
                fields
                    .iter()
                    .map(|_| alloc_value_id(next_value_id))
                    .collect()
            };

        let build_env_map = |fields: &[String], params: &[ValueId]| -> BTreeMap<String, ValueId> {
            let mut env = BTreeMap::new();
            for (name, vid) in fields.iter().zip(params.iter()) {
                env.insert(name.clone(), *vid);
            }
            env
        };

        let collect_env_args = |fields: &[String], env: &BTreeMap<String, ValueId>| -> Result<Vec<ValueId>, String> {
            let mut args = Vec::with_capacity(fields.len());
            for name in fields {
                let vid = env.get(name).copied().ok_or_else(|| {
                    error_tags::freeze_with_hint(
                        "phase129/post_k/env_missing",
                        &format!("env missing required field '{name}'"),
                        "ensure env layout and env map are built from the same SSOT field list",
                    )
                })?;
                args.push(vid);
            }
            Ok(args)
        };

        let mut next_value_id: u32 = 1;

        // IDs (stable, dev-only)
        let main_id = JoinFuncId::new(0);
        let k_then_id = JoinFuncId::new(1);
        let k_else_id = JoinFuncId::new(2);
        let join_k_id = JoinFuncId::new(3);
        let post_k_id = JoinFuncId::new(4);

        // main(env)
        let main_params = alloc_env_params(&env_fields, &mut next_value_id);
        let mut env_main = build_env_map(&env_fields, &main_params);
        let mut main_func = JoinFunction::new(main_id, "main".to_string(), main_params);

        // Lower prefix (pre-if) statements into main
        for n in prefix_nodes {
            match n {
                StepNode::Stmt { kind, .. } => match kind {
                    StepStmtKind::Assign { target, value_ast } => {
                        if LegacyLowerer::lower_assign_stmt(
                            target,
                            value_ast,
                            &mut main_func.body,
                            &mut next_value_id,
                            &mut env_main,
                        )
                        .is_err()
                        {
                            return Ok(None);
                        }
                    }
                    StepStmtKind::LocalDecl { .. } => {}
                    _ => {
                        return Ok(None);
                    }
                },
                _ => {
                    return Ok(None);
                }
            }
        }

        // Extract if condition and branches
        let (cond_ast, then_branch, else_branch) = match if_node {
            StepNode::If {
                cond_ast,
                then_branch,
                else_branch,
                ..
            } => (cond_ast, then_branch.as_ref(), else_branch.as_deref()),
            _ => unreachable!(),
        };

        let else_branch = match else_branch {
            Some(b) => b,
            None => return Ok(None), // Phase 129-C requires explicit else
        };

        // Extract branch statements (without return, since post-if handles return)
        let then_stmts = Self::extract_branch_stmts(then_branch)?;
        let else_stmts = Self::extract_branch_stmts(else_branch)?;

        // k_then(env_in): <then_stmts> ; tailcall join_k(env_out)
        let then_params = alloc_env_params(&env_fields, &mut next_value_id);
        let mut env_then = build_env_map(&env_fields, &then_params);
        let mut then_func = JoinFunction::new(k_then_id, "k_then".to_string(), then_params);

        for stmt in then_stmts {
            match stmt {
                StepStmtKind::Assign { target, value_ast } => {
                    if LegacyLowerer::lower_assign_stmt(
                        target,
                        value_ast,
                        &mut then_func.body,
                        &mut next_value_id,
                        &mut env_then,
                    )
                    .is_err()
                    {
                        return Ok(None);
                    }
                }
                StepStmtKind::LocalDecl { .. } => {}
                _ => {
                    return Ok(None); // Unsupported statement
                }
            }
        }

        let then_args = collect_env_args(&env_fields, &env_then)?;
        then_func.body.push(JoinInst::Call {
            func: join_k_id,
            args: then_args,
            k_next: None,
            dst: None,
        });

        // k_else(env_in): <else_stmts> ; tailcall join_k(env_out)
        let else_params = alloc_env_params(&env_fields, &mut next_value_id);
        let mut env_else = build_env_map(&env_fields, &else_params);
        let mut else_func = JoinFunction::new(k_else_id, "k_else".to_string(), else_params);

        for stmt in else_stmts {
            match stmt {
                StepStmtKind::Assign { target, value_ast } => {
                    if LegacyLowerer::lower_assign_stmt(
                        target,
                        value_ast,
                        &mut else_func.body,
                        &mut next_value_id,
                        &mut env_else,
                    )
                    .is_err()
                    {
                        return Ok(None);
                    }
                }
                StepStmtKind::LocalDecl { .. } => {}
                _ => {
                    return Ok(None); // Unsupported statement
                }
            }
        }

        let else_args = collect_env_args(&env_fields, &env_else)?;
        else_func.body.push(JoinInst::Call {
            func: join_k_id,
            args: else_args,
            k_next: None,
            dst: None,
        });

        // join_k(env_phi): tailcall post_k(env_phi)
        let join_k_params = alloc_env_params(&env_fields, &mut next_value_id);
        let env_join_k = build_env_map(&env_fields, &join_k_params);
        let mut join_k_func = JoinFunction::new(join_k_id, "join_k".to_string(), join_k_params);

        let join_k_args = collect_env_args(&env_fields, &env_join_k)?;
        join_k_func.body.push(JoinInst::Call {
            func: post_k_id,
            args: join_k_args,
            k_next: None,
            dst: None,
        });

        // post_k(env): <post_stmts> ; Ret
        let post_k_params = alloc_env_params(&env_fields, &mut next_value_id);
        let mut env_post_k = build_env_map(&env_fields, &post_k_params);
        let mut post_k_func = JoinFunction::new(post_k_id, "post_k".to_string(), post_k_params);

        // Lower post-if statements
        for n in post_nodes {
            match n {
                StepNode::Stmt { kind, .. } => match kind {
                    StepStmtKind::Return { value_ast } => {
                        // Phase 124/125/126: Return(Variable) support
                        if LegacyLowerer::lower_return_value(
                            value_ast,
                            &mut post_k_func.body,
                            &mut next_value_id,
                            &env_post_k,
                            &step_tree.contract,
                        )
                        .is_err()
                        {
                            return Ok(None);
                        }
                    }
                    _ => {
                        return Ok(None); // Unsupported post-if statement
                    }
                },
                _ => {
                    return Ok(None);
                }
            }
        }

        // If no return was emitted, add void return
        if !post_k_func.body.iter().any(|inst| matches!(inst, JoinInst::Ret { .. })) {
            post_k_func.body.push(JoinInst::Ret { value: None });
        }

        // main: cond compare + conditional jump to k_then/k_else
        let (lhs_var, op, rhs_literal) = LegacyLowerer::parse_minimal_compare(&cond_ast.0)?;
        let lhs_vid = env_main.get(&lhs_var).copied().ok_or_else(|| {
            error_tags::freeze_with_hint(
                "phase129/post_k/cond_lhs_missing",
                &format!("condition lhs var '{lhs_var}' not found in env"),
                "ensure the if condition uses a variable from writes or captured inputs",
            )
        })?;

        let rhs_vid = alloc_value_id(&mut next_value_id);
        main_func.body.push(JoinInst::Compute(MirLikeInst::Const {
            dst: rhs_vid,
            value: ConstValue::Integer(rhs_literal),
        }));

        let cond_vid = alloc_value_id(&mut next_value_id);
        main_func.body.push(JoinInst::Compute(MirLikeInst::Compare {
            dst: cond_vid,
            op,
            lhs: lhs_vid,
            rhs: rhs_vid,
        }));

        let main_args = collect_env_args(&env_fields, &env_main)?;
        main_func.body.push(JoinInst::Jump {
            cont: k_then_id.as_cont(),
            args: main_args.clone(),
            cond: Some(cond_vid),
        });
        main_func.body.push(JoinInst::Jump {
            cont: k_else_id.as_cont(),
            args: main_args,
            cond: None,
        });

        // Build module
        let mut module = JoinModule::new();
        module.add_function(main_func);
        module.add_function(then_func);
        module.add_function(else_func);
        module.add_function(join_k_func);
        module.add_function(post_k_func);
        module.entry = Some(main_id);
        module.mark_normalized();

        Ok(Some((module, JoinFragmentMeta::empty())))
    }

    /// Extract if-with-post pattern: (prefix, if_node, post)
    ///
    /// Returns None if not an if-with-post pattern.
    fn extract_if_with_post(node: &StepNode) -> Option<(&[StepNode], &StepNode, &[StepNode])> {
        match node {
            StepNode::Block(nodes) => {
                // Find the last If node
                let if_pos = nodes.iter().position(|n| matches!(n, StepNode::If { .. }))?;

                // Must have post-if statements
                if if_pos == nodes.len() - 1 {
                    return None; // No post-if (this is handled by if_as_last_join_k)
                }

                let if_node = &nodes[if_pos];
                let prefix = &nodes[..if_pos];
                let post = &nodes[if_pos + 1..];

                Some((prefix, if_node, post))
            }
            _ => None,
        }
    }

    /// Extract statements from a branch (excluding return)
    ///
    /// Phase 129-C: branches should not end with return (post-if handles return)
    fn extract_branch_stmts(branch: &StepNode) -> Result<Vec<&StepStmtKind>, String> {
        match branch {
            StepNode::Block(nodes) => {
                let mut stmts = Vec::new();
                for n in nodes {
                    match n {
                        StepNode::Stmt { kind, .. } => {
                            // Skip return statements (handled in post_k)
                            if !matches!(kind, StepStmtKind::Return { .. }) {
                                stmts.push(kind);
                            }
                        }
                        _ => {
                            // Unsupported node in branch
                            return Ok(Vec::new()); // Signal out-of-scope
                        }
                    }
                }
                Ok(stmts)
            }
            StepNode::Stmt { kind, .. } => {
                // Single statement branch
                if matches!(kind, StepStmtKind::Return { .. }) {
                    Ok(Vec::new()) // Empty if only return
                } else {
                    Ok(vec![kind])
                }
            }
            _ => Ok(Vec::new()), // Unsupported
        }
    }
}
