use std::collections::HashMap;
use std::convert::TryFrom;

use anyhow::{anyhow, Context, Result};
use base64::engine::general_purpose::STANDARD as BASE64_STANDARD;
use base64::Engine as _;
use chrono::{DateTime, Utc};
use eframe::egui;
use serde::Deserialize;
use serde_json::{json, Map, Value};

use crate::api::ApiClient;
use crate::core::components::interactions::ConfirmDialog;
use crate::core::tables::{
    find_asset_by_tag_or_numeric, find_zone_by_code, get_assets_in_zone, get_audit_task_definition,
};

const STATUS_OPTIONS: &[&str] = &[
    "Good",
    "Attention",
    "Faulty",
    "Missing",
    "Retired",
    "In Repair",
    "In Transit",
    "Expired",
    "Unmanaged",
];

const EXCEPTION_WRONG_ZONE: &str = "wrong-zone";
const EXCEPTION_UNEXPECTED_ASSET: &str = "unexpected-asset";
const EXCEPTION_OTHER: &str = "other";
const DEFAULT_MISSING_DETAIL: &str = "Marked missing during audit";

#[derive(Debug, Clone, Copy, PartialEq, Eq)]
pub enum AuditMode {
    FullZone,
    SpotCheck,
}

#[derive(Debug, Clone)]
struct ZoneInfo {
    id: i64,
    zone_code: Option<String>,
    zone_name: String,
    _zone_type: Option<String>,
    audit_timeout_minutes: Option<i64>,
}

impl ZoneInfo {
    fn from_value(value: &Value) -> Result<Self> {
        let id = value
            .get("id")
            .and_then(|v| v.as_i64())
            .ok_or_else(|| anyhow!("Zone record missing id"))?;
        let zone_name = value
            .get("zone_name")
            .and_then(|v| v.as_str())
            .unwrap_or("")
            .to_string();
        Ok(Self {
            id,
            zone_code: value
                .get("zone_code")
                .and_then(|v| v.as_str())
                .map(|s| s.to_string()),
            zone_name,
            _zone_type: value
                .get("zone_type")
                .and_then(|v| v.as_str())
                .map(|s| s.to_string()),
            audit_timeout_minutes: value.get("audit_timeout_minutes").and_then(|v| v.as_i64()),
        })
    }
}

#[derive(Debug, Clone, Copy, PartialEq, Eq)]
enum AuditScanPolicy {
    Required,
    Ask,
    Skip,
}

impl AuditScanPolicy {
    fn from_value(value: Option<&Value>) -> Self {
        match value.and_then(|v| v.as_str()).map(|s| s.to_lowercase()) {
            Some(ref s) if s == "yes" => AuditScanPolicy::Skip,
            Some(ref s) if s == "ask" => AuditScanPolicy::Ask,
            _ => AuditScanPolicy::Required,
        }
    }
}

#[derive(Debug, Clone)]
struct AuditAssetState {
    asset_id: i64,
    asset_numeric_id: Option<i64>,
    asset_tag: String,
    name: String,
    _status_before: Option<String>,
    scan_policy: AuditScanPolicy,
    audit_task_id: Option<i64>,
    expected: bool,
    _expected_zone_id: Option<i64>,
    _actual_zone_id: Option<i64>,
    scanned: bool,
    status_found: String,
    notes: String,
    task_responses: Option<Value>,
    additional_fields: Map<String, Value>,
    exception_type: Option<String>,
    exception_details: Option<String>,
    completed_at: Option<DateTime<Utc>>,
}

impl AuditAssetState {
    fn from_value(value: Value, expected_zone_id: Option<i64>, expected: bool) -> Result<Self> {
        let asset_id = value
            .get("id")
            .and_then(|v| v.as_i64())
            .ok_or_else(|| anyhow!("Asset record missing id"))?;
        let asset_tag = value
            .get("asset_tag")
            .and_then(|v| v.as_str())
            .unwrap_or("")
            .to_string();
        let name = value
            .get("name")
            .and_then(|v| v.as_str())
            .unwrap_or("Unnamed Asset")
            .to_string();
        let status_before = value
            .get("status")
            .and_then(|v| v.as_str())
            .map(|s| s.to_string());
        let scan_policy = AuditScanPolicy::from_value(value.get("no_scan"));
        let status_found = status_before.clone().unwrap_or_else(|| "Good".to_string());
        Ok(Self {
            asset_id,
            asset_numeric_id: value.get("asset_numeric_id").and_then(|v| v.as_i64()),
            asset_tag,
            name,
            _status_before: status_before,
            scan_policy,
            audit_task_id: value.get("audit_task_id").and_then(|v| v.as_i64()),
            expected,
            _expected_zone_id: expected_zone_id,
            _actual_zone_id: value.get("zone_id").and_then(|v| v.as_i64()),
            scanned: matches!(scan_policy, AuditScanPolicy::Skip),
            status_found,
            notes: String::new(),
            task_responses: None,
            additional_fields: Map::new(),
            exception_type: None,
            exception_details: None,
            completed_at: if matches!(scan_policy, AuditScanPolicy::Skip) {
                Some(Utc::now())
            } else {
                None
            },
        })
    }

    fn requires_scan(&self) -> bool {
        self.expected && matches!(self.scan_policy, AuditScanPolicy::Required)
    }

    fn matches_identifier(&self, identifier: &str) -> bool {
        let normalized = identifier.trim().to_lowercase();
        if normalized.is_empty() {
            return false;
        }
        let tag_match = !self.asset_tag.is_empty() && self.asset_tag.to_lowercase() == normalized;
        let numeric_match = self
            .asset_numeric_id
            .map(|n| n.to_string() == normalized)
            .unwrap_or(false);
        tag_match || numeric_match
    }

    fn display_label(&self, mode: AuditMode) -> String {
        let mut label = format!("{} — {}", self.asset_tag, self.name);
        if !self.expected {
            label.push_str(match mode {
                AuditMode::FullZone => " (unexpected)",
                AuditMode::SpotCheck => " (spot check)",
            });
        }
        if matches!(self.scan_policy, AuditScanPolicy::Ask) && !self.scanned {
            label.push_str(" (confirm)");
        } else if !self.requires_scan() {
            label.push_str(" (auto)");
        } else if !self.scanned {
            label.push_str(" (pending)");
        }
        label
    }

    fn set_status(&mut self, status: &str, mark_scanned: bool) {
        self.status_found = status.to_string();
        if mark_scanned {
            self.scanned = true;
            self.completed_at = Some(Utc::now());
        }
        if status == "Missing" {
            if matches!(self.scan_policy, AuditScanPolicy::Ask) && !self.scanned {
                // Leave confirmation-driven assets pending until explicitly handled
                self.status_found = "Missing".to_string();
                self.scanned = false;
                self.completed_at = None;
                return;
            }

            self.exception_type = Some(EXCEPTION_OTHER.to_string());
            if self
                .exception_details
                .as_deref()
                .map(|d| d == DEFAULT_MISSING_DETAIL)
                .unwrap_or(true)
            {
                self.exception_details = Some(DEFAULT_MISSING_DETAIL.to_string());
            }
        } else if self
            .exception_type
            .as_deref()
            .map(|t| t == EXCEPTION_OTHER)
            .unwrap_or(false)
            && self
                .exception_details
                .as_deref()
                .map(|d| d == DEFAULT_MISSING_DETAIL)
                .unwrap_or(false)
        {
            self.exception_type = None;
            self.exception_details = None;
        }
    }
}

#[derive(Debug, Clone)]
struct TaskRunnerState {
    asset_index: usize,
    runner: AuditTaskRunner,
}

#[derive(Debug, Clone)]
struct AuditTaskOutcome {
    status_override: Option<String>,
    additional_fields: Map<String, Value>,
    responses: Value,
}

#[derive(Debug, Clone)]
struct TaskResponseEntry {
    step: i64,
    question: String,
    answer: Value,
}

#[allow(dead_code)]
#[derive(Debug, Clone)]
pub struct AuditCompletion {
    pub audit_id: i64,
    pub status: String,
}

#[derive(Debug, Clone, Copy)]
enum PendingFinalizeIntent {
    FromButton { needs_force: bool },
    FromDialog { force_missing: bool },
}

pub struct AuditWorkflow {
    is_open: bool,
    mode: AuditMode,
    zone_info: Option<ZoneInfo>,
    expected_assets: Vec<AuditAssetState>,
    selected_asset: Option<usize>,
    scan_input: String,
    notes: String,
    audit_name: String,
    started_at: Option<DateTime<Utc>>,
    timeout_minutes: Option<i64>,
    last_error: Option<String>,
    ask_dialog: ConfirmDialog,
    pending_ask_index: Option<usize>,
    cancel_dialog: ConfirmDialog,
    finalize_dialog: ConfirmDialog,
    current_task_runner: Option<TaskRunnerState>,
    cached_tasks: HashMap<i64, AuditTaskDefinition>,
    has_recent_completion: bool,
    completion_snapshot: Option<AuditCompletion>,
    user_id: Option<i64>,
    pending_finalize: Option<PendingFinalizeIntent>,
}

impl AuditWorkflow {
    pub fn new() -> Self {
        Self {
            is_open: false,
            mode: AuditMode::FullZone,
            zone_info: None,
            expected_assets: Vec::new(),
            selected_asset: None,
            scan_input: String::new(),
            notes: String::new(),
            audit_name: String::new(),
            started_at: None,
            timeout_minutes: None,
            last_error: None,
            ask_dialog: ConfirmDialog::new(
                "Confirm Asset",
                "This asset is marked as 'Ask'. Confirm to include it in the audit progress.",
            )
            .dangerous(false)
            .confirm_text("Confirm")
            .cancel_text("Skip"),
            pending_ask_index: None,
            cancel_dialog: ConfirmDialog::new(
                "Cancel Audit",
                "Are you sure you want to cancel the current audit? Progress will be lost.",
            )
            .dangerous(true)
            .confirm_text("Cancel Audit")
            .cancel_text("Keep Working"),
            finalize_dialog: ConfirmDialog::new(
                "Complete Audit",
                "Some required assets have not been scanned. They will be marked as Missing if you continue.",
            )
            .dangerous(true)
            .confirm_text("Mark Missing & Complete")
            .cancel_text("Go Back"),
            current_task_runner: None,
            cached_tasks: HashMap::new(),
            has_recent_completion: false,
            completion_snapshot: None,
            user_id: None,
            pending_finalize: None,
        }
    }

    pub fn is_active(&self) -> bool {
        self.is_open
    }

    pub fn start_zone_audit(
        &mut self,
        api_client: &ApiClient,
        zone_code: &str,
        user_id: i64,
    ) -> Result<()> {
        let zone_value = find_zone_by_code(api_client, zone_code)?
            .ok_or_else(|| anyhow!("Zone '{}' was not found", zone_code))?;
        let zone = ZoneInfo::from_value(&zone_value)?;
        let zone_id = i32::try_from(zone.id).context("Zone identifier exceeds i32 range")?;
        let raw_assets = get_assets_in_zone(api_client, zone_id, Some(1_000))?;

        let mut assets = Vec::with_capacity(raw_assets.len());
        for value in raw_assets {
            let mut state = AuditAssetState::from_value(value, Some(zone.id), true)?;
            if matches!(state.scan_policy, AuditScanPolicy::Skip) {
                state.completed_at = Some(Utc::now());
            }
            assets.push(state);
        }

        self.reset_core_state();
        self.has_recent_completion = false;
        self.completion_snapshot = None;
        self.is_open = true;
        self.mode = AuditMode::FullZone;
        self.zone_info = Some(zone.clone());
        self.expected_assets = assets;
        self.started_at = Some(Utc::now());
        self.timeout_minutes = zone.audit_timeout_minutes;
        self.audit_name = format!("Zone {} Audit", zone.zone_name);
        self.user_id = Some(user_id);
        self.last_error = None;
        self.ensure_skip_assets_recorded();
        Ok(())
    }

    pub fn start_spot_check(&mut self, user_id: i64) {
        self.reset_core_state();
        self.has_recent_completion = false;
        self.completion_snapshot = None;
        self.is_open = true;
        self.mode = AuditMode::SpotCheck;
        self.audit_name = format!("Spot Check {}", Utc::now().format("%Y-%m-%d %H:%M"));
        self.started_at = Some(Utc::now());
        self.user_id = Some(user_id);
        self.last_error = None;
    }

    pub fn show(&mut self, ctx: &egui::Context, api_client: &ApiClient) -> bool {
        if !self.is_open {
            return false;
        }

        let mut keep_open = self.is_open;
        let window_title = match self.mode {
            AuditMode::FullZone => "Zone Audit",
            AuditMode::SpotCheck => "Spot Check",
        };

        let screen_rect = ctx.input(|i| {
            i.viewport().inner_rect.unwrap_or(egui::Rect::from_min_max(
                egui::pos2(0.0, 0.0),
                egui::pos2(1920.0, 1080.0),
            ))
        });
        let mut max_size = screen_rect.size() - egui::vec2(32.0, 32.0);
        max_size.x = max_size.x.max(860.0).min(screen_rect.width());
        max_size.y = max_size.y.max(520.0).min(screen_rect.height());
        let mut default_size = egui::vec2(1040.0, 680.0);
        default_size.x = default_size.x.min(max_size.x);
        default_size.y = default_size.y.min(max_size.y);

        egui::Window::new(window_title)
            .id(egui::Id::new("audit_workflow_window"))
            .collapsible(false)
            .resizable(true)
            .default_size(default_size)
            .max_size(max_size)
            .min_size(egui::vec2(820.0, 520.0))
            .open(&mut keep_open)
            .show(ctx, |ui| {
                if let Some(zone) = &self.zone_info {
                    ui.horizontal(|ui| {
                        ui.heading(format!(
                            "Auditing {} ({})",
                            zone.zone_name,
                            zone.zone_code.as_deref().unwrap_or("no-code")
                        ));
                        if let Some(timeout) = zone.audit_timeout_minutes {
                            ui.add_space(12.0);
                            ui.label(format!("Timeout: {} min", timeout));
                        }
                    });
                } else {
                    ui.heading(&self.audit_name);
                }

                if let Some(err) = &self.last_error {
                    ui.add_space(8.0);
                    ui.colored_label(egui::Color32::RED, err);
                }

                ui.add_space(8.0);
                self.render_scanning(ui, ctx, api_client);
            });

        if !keep_open {
            self.cancel_without_saving();
        }

        if let Some(result) = self.ask_dialog.show_dialog(ctx) {
            self.process_ask_dialog(result, api_client);
        }

        if let Some(result) = self.cancel_dialog.show_dialog(ctx) {
            if result {
                match self.cancel_audit(api_client) {
                    Ok(()) => {
                        keep_open = false;
                    }
                    Err(err) => {
                        self.last_error = Some(err.to_string());
                    }
                }
            }
        }

        if let Some(result) = self.finalize_dialog.show_dialog(ctx) {
            if result {
                if self.trigger_pending_ask(PendingFinalizeIntent::FromDialog {
                    force_missing: true,
                }) {
                    // Ask dialog opened; finalize will continue after confirmations.
                } else if let Err(err) = self.finalize_audit(api_client, true) {
                    self.last_error = Some(err.to_string());
                }
            }
        }

        if let Some(mut state) = self.current_task_runner.take() {
            if let Some(outcome) = state.runner.show(ctx) {
                self.apply_task_outcome(state.asset_index, outcome);
            } else if state.runner.is_open() {
                self.current_task_runner = Some(state);
            }
        }

        if !self.is_open {
            keep_open = false;
        }
        self.is_open = keep_open;
        keep_open
    }

    pub fn take_recent_completion(&mut self) -> Option<AuditCompletion> {
        if self.has_recent_completion {
            self.has_recent_completion = false;
            self.completion_snapshot.take()
        } else {
            None
        }
    }

    fn reset_core_state(&mut self) {
        self.is_open = false;
        self.zone_info = None;
        self.expected_assets.clear();
        self.selected_asset = None;
        self.scan_input.clear();
        self.notes.clear();
        self.audit_name.clear();
        self.started_at = None;
        self.timeout_minutes = None;
        self.last_error = None;
        self.pending_ask_index = None;
        self.current_task_runner = None;
        self.user_id = None;
        self.pending_finalize = None;
        // Preserve cached_tasks so audit tasks are reused between runs
    }

    fn cancel_without_saving(&mut self) {
        self.reset_core_state();
    }

    fn cancel_audit(&mut self, api_client: &ApiClient) -> Result<()> {
        if !self.is_open {
            return Ok(());
        }

        if self.started_at.is_none() {
            self.reset_core_state();
            return Ok(());
        }

        let user_id = self
            .user_id
            .ok_or_else(|| anyhow!("Missing current user id for audit session"))?;
        let started_at = self.started_at.unwrap();
        let cancelled_at = Utc::now();

        let required_total = self.required_total();
        let _scanned_total = self.expected_assets.iter().filter(|a| a.scanned).count();

        let mut found_count = 0;
        let mut missing_assets = Vec::new();
        let mut attention_assets = Vec::new();
        let mut exceptions = Vec::new();
        let mut unexpected_assets = Vec::new();

        for asset in &self.expected_assets {
            if !asset.scanned {
                continue;
            }

            if asset.expected && asset.requires_scan() {
                if asset.status_found != "Missing" {
                    found_count += 1;
                } else {
                    missing_assets.push(json!({
                        "asset_id": asset.asset_id,
                        "asset_tag": asset.asset_tag,
                        "name": asset.name,
                    }));
                }

                if asset.status_found != "Good" && asset.status_found != "Missing" {
                    attention_assets.push(json!({
                        "asset_id": asset.asset_id,
                        "asset_tag": asset.asset_tag,
                        "name": asset.name,
                        "status": asset.status_found,
                    }));
                }
            }

            if let Some(ref exception) = asset.exception_type {
                exceptions.push(json!({
                    "asset_id": asset.asset_id,
                    "asset_tag": asset.asset_tag,
                    "name": asset.name,
                    "type": exception,
                    "details": asset.exception_details,
                }));
            }

            if !asset.expected {
                unexpected_assets.push(json!({
                    "asset_id": asset.asset_id,
                    "asset_tag": asset.asset_tag,
                    "name": asset.name,
                    "status": asset.status_found,
                }));
            }
        }

        let mut issues = Map::new();
        if !missing_assets.is_empty() {
            issues.insert("missing_assets".into(), Value::Array(missing_assets));
        }
        if !attention_assets.is_empty() {
            issues.insert("attention_assets".into(), Value::Array(attention_assets));
        }
        if !exceptions.is_empty() {
            issues.insert("exceptions".into(), Value::Array(exceptions));
        }
        if !unexpected_assets.is_empty() {
            issues.insert("unexpected_assets".into(), Value::Array(unexpected_assets));
        }

        let mut payload = Map::new();
        payload.insert(
            "audit_type".into(),
            Value::String(match self.mode {
                AuditMode::FullZone => "full-zone".to_string(),
                AuditMode::SpotCheck => "spot-check".to_string(),
            }),
        );
        if let Some(zone) = &self.zone_info {
            payload.insert("zone_id".into(), json!(zone.id));
        }
        if !self.audit_name.trim().is_empty() {
            payload.insert("audit_name".into(), json!(self.audit_name.trim()));
        }
        payload.insert("started_by".into(), json!(user_id));
        payload.insert(
            "started_at".into(),
            json!(started_at.format("%Y-%m-%d %H:%M:%S").to_string()),
        );
        payload.insert(
            "completed_at".into(),
            json!(cancelled_at.format("%Y-%m-%d %H:%M:%S").to_string()),
        );
        payload.insert("status".into(), json!("cancelled"));
        if let Some(timeout) = self.timeout_minutes {
            payload.insert("timeout_minutes".into(), json!(timeout));
        }
        if issues.is_empty() {
            payload.insert("issues_found".into(), Value::Null);
        } else {
            payload.insert("issues_found".into(), Value::Object(issues));
        }
        payload.insert("assets_expected".into(), json!(required_total as i64));
        payload.insert("assets_found".into(), json!(found_count as i64));
        if !self.notes.trim().is_empty() {
            payload.insert("notes".into(), json!(self.notes.trim()));
        }
        let cancel_reason = if let Some(zone) = &self.zone_info {
            format!(
                "Audit cancelled for zone {} at {}",
                zone.zone_name,
                cancelled_at.format("%Y-%m-%d %H:%M:%S")
            )
        } else {
            format!(
                "Spot check cancelled at {}",
                cancelled_at.format("%Y-%m-%d %H:%M:%S")
            )
        };
        payload.insert("cancelled_reason".into(), json!(cancel_reason));

        let audit_insert = api_client.insert("physical_audits", Value::Object(payload))?;
        if !audit_insert.success {
            return Err(anyhow!(
                "Failed to cancel audit session: {}",
                audit_insert
                    .error
                    .unwrap_or_else(|| "unknown error".to_string())
            ));
        }
        let audit_id = audit_insert.data.unwrap_or(0) as i64;

        for asset in &self.expected_assets {
            if !asset.scanned {
                continue;
            }

            let mut log_payload = Map::new();
            log_payload.insert("physical_audit_id".into(), json!(audit_id));
            log_payload.insert("asset_id".into(), json!(asset.asset_id));
            log_payload.insert("status_found".into(), json!(asset.status_found));
            if let Some(task_id) = asset.audit_task_id {
                log_payload.insert("audit_task_id".into(), json!(task_id));
            }
            if let Some(responses) = &asset.task_responses {
                log_payload.insert("audit_task_responses".into(), responses.clone());
            }
            if let Some(exception) = &asset.exception_type {
                log_payload.insert("exception_type".into(), json!(exception));
            }
            if let Some(details) = &asset.exception_details {
                log_payload.insert("exception_details".into(), json!(details));
            }
            if let Some(zone) = &self.zone_info {
                log_payload.insert("found_in_zone_id".into(), json!(zone.id));
            }
            if !asset.notes.trim().is_empty() {
                log_payload.insert("notes".into(), json!(asset.notes.trim()));
            }
            let log_insert =
                api_client.insert("physical_audit_logs", Value::Object(log_payload))?;
            if !log_insert.success {
                return Err(anyhow!(
                    "Failed to record cancellation log for asset {}",
                    asset.asset_tag
                ));
            }
        }

        let completion = AuditCompletion {
            audit_id,
            status: "cancelled".to_string(),
        };
        self.completion_snapshot = Some(completion);
        self.has_recent_completion = true;
        self.reset_core_state();
        Ok(())
    }

    fn ensure_skip_assets_recorded(&mut self) {
        for asset in &mut self.expected_assets {
            if !asset.scanned && matches!(asset.scan_policy, AuditScanPolicy::Skip) {
                asset.scanned = true;
                asset.completed_at = Some(Utc::now());
            }
        }
    }

    fn render_scanning(&mut self, ui: &mut egui::Ui, ctx: &egui::Context, api_client: &ApiClient) {
        let required_total = self.required_total();
        let completed_total = self.completed_total();
        let progress = if required_total > 0 {
            completed_total as f32 / required_total as f32
        } else {
            0.0
        };
        let remaining_required = self.remaining_required();

        ui.horizontal(|ui| {
            ui.vertical(|ui| {
                if required_total > 0 {
                    ui.add(
                        egui::ProgressBar::new(progress)
                            .text(format!("{}/{} processed", completed_total, required_total))
                            .desired_width(320.0),
                    );
                } else {
                    ui.label("No required assets to scan");
                }

                if self.mode == AuditMode::FullZone && remaining_required > 0 {
                    ui.add_space(4.0);
                    ui.colored_label(
                        egui::Color32::from_rgb(255, 179, 0),
                        format!(
                            "{} required assets pending; finishing now marks them Missing.",
                            remaining_required
                        ),
                    );
                }
            });

            ui.with_layout(egui::Layout::right_to_left(egui::Align::Center), |ui| {
                let mut complete_button = ui.add(egui::Button::new("Complete Audit"));
                if self.mode == AuditMode::FullZone && remaining_required > 0 {
                    complete_button = complete_button.on_hover_text(format!(
                        "{} required assets pending. Completing now will mark them as Missing.",
                        remaining_required
                    ));
                }
                if complete_button.clicked() {
                    let needs_force = self.mode == AuditMode::FullZone && remaining_required > 0;
                    if self.trigger_pending_ask(PendingFinalizeIntent::FromButton { needs_force }) {
                        // Ask dialog opened; completion will resume after confirmations.
                    } else if needs_force {
                        let name = format!("{} pending items", remaining_required);
                        let detail = "Unscanned assets will be marked Missing upon completion.";
                        self.finalize_dialog.open(name, detail);
                    } else if let Err(err) = self.finalize_audit(api_client, false) {
                        self.last_error = Some(err.to_string());
                    }
                }
                if ui.button("Cancel Audit").clicked() {
                    if let Some(zone) = &self.zone_info {
                        self.cancel_dialog
                            .open(&zone.zone_name, zone.zone_code.clone().unwrap_or_default());
                    } else {
                        self.cancel_dialog.open("Spot Check", &self.audit_name);
                    }
                }
            });
        });

        ui.add_space(10.0);
        ui.horizontal(|ui| {
            let input = ui.add(
                egui::TextEdit::singleline(&mut self.scan_input)
                    .hint_text("Scan asset tag or numeric ID")
                    .desired_width(260.0),
            );
            let submitted = input.lost_focus() && ui.input(|i| i.key_pressed(egui::Key::Enter));
            if submitted {
                if let Err(err) = self.handle_scan(api_client) {
                    self.last_error = Some(err.to_string());
                }
                ctx.request_repaint();
            }
            if ui.button("Submit").clicked() {
                if let Err(err) = self.handle_scan(api_client) {
                    self.last_error = Some(err.to_string());
                }
                ctx.request_repaint();
            }
            if ui.button("Clear").clicked() {
                self.scan_input.clear();
            }
        });

        ui.add_space(12.0);
        ui.separator();
        ui.add_space(8.0);

        ui.columns(2, |columns| {
            let [left, right] = columns else {
                return;
            };

            left.set_min_width(320.0);
            left.set_max_width(360.0);

            left.heading("Assets");
            left.add_space(4.0);
            let mut selection_change = None;
            egui::ScrollArea::vertical()
                .id_salt("audit_assets_scroll")
                .auto_shrink([false; 2])
                .show(left, |ui| {
                    for idx in 0..self.expected_assets.len() {
                        let asset = &self.expected_assets[idx];
                        let selected = self.selected_asset == Some(idx);
                        let label = asset.display_label(self.mode);
                        let response = ui.selectable_label(selected, label);
                        let response = if !asset.scanned && asset.requires_scan() {
                            response.on_hover_text("Pending scan")
                        } else {
                            response
                        };
                        if response.clicked() {
                            selection_change = Some(idx);
                        }
                    }
                });
            if let Some(idx) = selection_change {
                self.selected_asset = Some(idx);
            }

            right.set_min_width(right.available_width().max(420.0));
            right.heading("Details");
            right.add_space(4.0);
            if let Some(idx) = self.selected_asset {
                let mut run_task_clicked = None;
                if let Some(asset) = self.expected_assets.get_mut(idx) {
                    right.label(format!("Asset Tag: {}", asset.asset_tag));
                    right.label(format!("Name: {}", asset.name));
                    if !asset.expected {
                        right.colored_label(
                            egui::Color32::from_rgb(255, 152, 0),
                            "Unexpected asset",
                        );
                    }
                    if let Some(policy_text) = match asset.scan_policy {
                        AuditScanPolicy::Required => None,
                        AuditScanPolicy::Ask => Some("Requires confirmation"),
                        AuditScanPolicy::Skip => Some("Auto-completed"),
                    } {
                        right.label(policy_text);
                    }

                    right.add_space(6.0);
                    let mut status_value = asset.status_found.clone();
                    egui::ComboBox::from_label("Status")
                        .selected_text(&status_value)
                        .show_ui(right, |ui| {
                            for option in STATUS_OPTIONS {
                                ui.selectable_value(&mut status_value, option.to_string(), *option);
                            }
                        });
                    if status_value != asset.status_found {
                        asset.set_status(&status_value, true);
                    }

                    right.add_space(6.0);
                    right.label("Notes");
                    right.add(
                        egui::TextEdit::multiline(&mut asset.notes)
                            .desired_rows(3)
                            .desired_width(right.available_width())
                            .hint_text("Optional notes for this asset"),
                    );

                    right.add_space(6.0);
                    right.horizontal(|ui| {
                        if ui.button("Mark Good").clicked() {
                            asset.set_status("Good", true);
                            asset.exception_type = None;
                            asset.exception_details = None;
                        }
                        if ui.button("Mark Missing").clicked() {
                            asset.set_status("Missing", true);
                            asset.exception_type = Some(EXCEPTION_OTHER.to_string());
                            asset.exception_details = Some(DEFAULT_MISSING_DETAIL.to_string());
                        }
                    });

                    if let Some(task_id) = asset.audit_task_id {
                        right.add_space(6.0);
                        if right.button("Run Audit Task").clicked() {
                            run_task_clicked = Some(task_id);
                        }
                    }

                    if asset.requires_scan() && !asset.scanned {
                        right.add_space(6.0);
                        if right.button("Mark Scanned").clicked() {
                            let current_status = asset.status_found.clone();
                            asset.set_status(&current_status, true);
                        }
                    }
                }

                if let Some(task_id) = run_task_clicked {
                    if let Err(err) = self.launch_task_runner(idx, task_id, api_client) {
                        self.last_error = Some(err.to_string());
                    }
                }
            } else {
                right.label("Select an asset to see details.");
            }

            right.add_space(12.0);
            right.separator();
            right.add_space(8.0);
            right.label("Audit Notes");
            right.add(
                egui::TextEdit::multiline(&mut self.notes)
                    .desired_rows(4)
                    .desired_width(right.available_width())
                    .hint_text("Optional notes for the entire audit"),
            );
        });
    }

    fn required_total(&self) -> usize {
        self.expected_assets
            .iter()
            .filter(|a| a.requires_scan())
            .count()
    }

    fn completed_total(&self) -> usize {
        self.expected_assets
            .iter()
            .filter(|a| a.requires_scan() && a.scanned)
            .count()
    }

    fn remaining_required(&self) -> usize {
        self.expected_assets
            .iter()
            .filter(|asset| asset.requires_scan() && !asset.scanned)
            .count()
    }

    fn next_unresolved_ask(&self) -> Option<usize> {
        self.expected_assets
            .iter()
            .enumerate()
            .find(|(_, asset)| matches!(asset.scan_policy, AuditScanPolicy::Ask) && !asset.scanned)
            .map(|(idx, _)| idx)
    }

    fn trigger_pending_ask(&mut self, intent: PendingFinalizeIntent) -> bool {
        if self.ask_dialog.show || self.pending_ask_index.is_some() {
            self.pending_finalize = Some(intent);
            return true;
        }

        if let Some(idx) = self.next_unresolved_ask() {
            if let Some(asset) = self.expected_assets.get(idx) {
                self.pending_finalize = Some(intent);
                self.pending_ask_index = Some(idx);
                self.ask_dialog
                    .open(asset.name.clone(), asset.asset_tag.clone());
                return true;
            }
        }

        false
    }

    fn handle_scan(&mut self, api_client: &ApiClient) -> Result<()> {
        let input = self.scan_input.trim();
        if input.is_empty() {
            return Ok(());
        }

        self.last_error = None;

        if let Some(idx) = self
            .expected_assets
            .iter()
            .position(|asset| asset.matches_identifier(input))
        {
            self.selected_asset = Some(idx);
            self.process_matched_asset(idx, api_client)?;
            self.scan_input.clear();
            return Ok(());
        }

        // Asset not in current list, try to fetch from the API
        if let Some(value) = find_asset_by_tag_or_numeric(api_client, input)? {
            let zone_id = value.get("zone_id").and_then(|v| v.as_i64());
            let mut state = AuditAssetState::from_value(
                value,
                self.zone_info.as_ref().map(|z| z.id),
                self.mode == AuditMode::FullZone && self.zone_info.is_some(),
            )?;

            if let Some(zone) = &self.zone_info {
                if zone_id != Some(zone.id) {
                    state.expected = false;
                    state.exception_type = Some(EXCEPTION_WRONG_ZONE.to_string());
                    state.exception_details = Some(format!(
                        "Asset assigned to zone {:?}, found in {}",
                        zone_id, zone.zone_name
                    ));
                } else if self.mode == AuditMode::FullZone {
                    state.expected = false;
                    state.exception_type = Some(EXCEPTION_UNEXPECTED_ASSET.to_string());
                    state.exception_details = Some("Asset not listed on zone roster".to_string());
                }
            } else {
                state.expected = false;
                state.exception_type = Some(EXCEPTION_UNEXPECTED_ASSET.to_string());
                state.exception_details = Some("Captured during spot check".to_string());
            }

            let idx = self.expected_assets.len();
            self.expected_assets.push(state);
            self.selected_asset = Some(idx);
            self.process_matched_asset(idx, api_client)?;
            self.scan_input.clear();
            return Ok(());
        }

        self.last_error = Some(format!("No asset found for '{}'.", input));
        self.scan_input.clear();
        Ok(())
    }

    fn process_matched_asset(&mut self, index: usize, api_client: &ApiClient) -> Result<()> {
        if index >= self.expected_assets.len() {
            return Ok(());
        }

        let (policy, already_scanned, task_id, name, tag, status_value) = {
            let asset = &self.expected_assets[index];
            (
                asset.scan_policy,
                asset.scanned,
                asset.audit_task_id,
                asset.name.clone(),
                asset.asset_tag.clone(),
                asset.status_found.clone(),
            )
        };

        if matches!(policy, AuditScanPolicy::Ask) && !already_scanned {
            self.pending_ask_index = Some(index);
            self.ask_dialog.open(name, tag);
            return Ok(());
        }

        if let Some(task_id) = task_id {
            if !already_scanned {
                self.launch_task_runner(index, task_id, api_client)?;
                return Ok(());
            }
        }

        if !already_scanned {
            self.expected_assets[index].set_status(&status_value, true);
        }

        Ok(())
    }

    fn launch_task_runner(
        &mut self,
        index: usize,
        task_id: i64,
        api_client: &ApiClient,
    ) -> Result<()> {
        if let Some(state) = &self.current_task_runner {
            if state.asset_index == index {
                return Ok(()); // already running for this asset
            }
        }

        let definition = if let Some(def) = self.cached_tasks.get(&task_id) {
            def.clone()
        } else {
            let task_value = get_audit_task_definition(api_client, task_id)?
                .ok_or_else(|| anyhow!("Audit task {} not found", task_id))?;
            let task_json = task_value
                .get("json_sequence")
                .cloned()
                .unwrap_or(Value::Null);
            let definition = AuditTaskDefinition::from_value(task_json)?;
            self.cached_tasks.insert(task_id, definition.clone());
            definition
        };

        let asset_label = self.expected_assets[index].name.clone();
        let runner = AuditTaskRunner::new(definition, asset_label);
        self.current_task_runner = Some(TaskRunnerState {
            asset_index: index,
            runner,
        });
        Ok(())
    }

    fn process_ask_dialog(&mut self, confirmed: bool, api_client: &ApiClient) {
        if let Some(idx) = self.pending_ask_index.take() {
            if idx < self.expected_assets.len() {
                let task_id = self.expected_assets[idx].audit_task_id;
                if confirmed {
                    if let Some(task_id) = task_id {
                        if let Err(err) = self.launch_task_runner(idx, task_id, api_client) {
                            self.last_error = Some(err.to_string());
                        }
                    } else {
                        let status_value = self.expected_assets[idx].status_found.clone();
                        self.expected_assets[idx].set_status(&status_value, true);
                    }
                } else {
                    self.expected_assets[idx].set_status("Missing", true);
                    self.expected_assets[idx].exception_type = Some(EXCEPTION_OTHER.to_string());
                    if self.expected_assets[idx].exception_details.is_none() {
                        self.expected_assets[idx].exception_details =
                            Some(DEFAULT_MISSING_DETAIL.to_string());
                    }
                }
            }
        }

        if let Some(intent) = self.pending_finalize.take() {
            if self.trigger_pending_ask(intent) {
                return;
            }

            match intent {
                PendingFinalizeIntent::FromButton { needs_force } => {
                    if needs_force {
                        let remaining = self.remaining_required();
                        if remaining > 0 {
                            let name = format!("{} pending items", remaining);
                            let detail = "Unscanned assets will be marked Missing upon completion.";
                            self.finalize_dialog.open(name, detail);
                        } else if let Err(err) = self.finalize_audit(api_client, false) {
                            self.last_error = Some(err.to_string());
                        }
                    } else if let Err(err) = self.finalize_audit(api_client, false) {
                        self.last_error = Some(err.to_string());
                    }
                }
                PendingFinalizeIntent::FromDialog { force_missing } => {
                    if let Err(err) = self.finalize_audit(api_client, force_missing) {
                        self.last_error = Some(err.to_string());
                    }
                }
            }
        }
    }

    fn apply_task_outcome(&mut self, index: usize, mut outcome: AuditTaskOutcome) {
        if let Some(asset) = self.expected_assets.get_mut(index) {
            if let Some(status) = outcome.status_override.take() {
                asset.set_status(&status, true);
            } else {
                let current_status = asset.status_found.clone();
                asset.set_status(&current_status, true);
            }

            if !outcome.additional_fields.is_empty() {
                asset.additional_fields = outcome.additional_fields.clone();
            }

            let mut payload = Map::new();
            payload.insert("responses".into(), outcome.responses);
            if !outcome.additional_fields.is_empty() {
                payload.insert(
                    "additional_fields".into(),
                    Value::Object(outcome.additional_fields.clone()),
                );
            }
            asset.task_responses = Some(Value::Object(payload));
        }
    }

    fn finalize_audit(&mut self, api_client: &ApiClient, force_missing: bool) -> Result<()> {
        let remaining = self.remaining_required();
        if remaining > 0 {
            if !force_missing {
                return Err(anyhow!(
                    "Cannot finalize audit. {} required assets still pending.",
                    remaining
                ));
            }

            for asset in &mut self.expected_assets {
                if asset.requires_scan() && !asset.scanned {
                    asset.set_status("Missing", true);
                    asset.exception_type = Some(EXCEPTION_OTHER.to_string());
                    if asset.exception_details.is_none() {
                        asset.exception_details = Some(DEFAULT_MISSING_DETAIL.to_string());
                    }
                }
            }
        }

        let user_id = self
            .user_id
            .ok_or_else(|| anyhow!("Missing current user id for audit session"))?;
        let started_at = self
            .started_at
            .unwrap_or_else(|| Utc::now() - chrono::Duration::minutes(1));
        let completed_at = Utc::now();

        let required_total = self.required_total();
        let mut found_count = 0;
        let mut missing_assets = Vec::new();
        let mut attention_assets = Vec::new();
        let mut exceptions = Vec::new();
        let mut unexpected_assets = Vec::new();

        for asset in &self.expected_assets {
            if asset.expected && asset.requires_scan() {
                if asset.status_found != "Missing" {
                    found_count += 1;
                } else {
                    missing_assets.push(json!({
                        "asset_id": asset.asset_id,
                        "asset_tag": asset.asset_tag,
                        "name": asset.name,
                    }));
                }

                if asset.status_found != "Good" && asset.status_found != "Missing" {
                    attention_assets.push(json!({
                        "asset_id": asset.asset_id,
                        "asset_tag": asset.asset_tag,
                        "name": asset.name,
                        "status": asset.status_found,
                    }));
                }
            }

            if let Some(ref exception) = asset.exception_type {
                exceptions.push(json!({
                    "asset_id": asset.asset_id,
                    "asset_tag": asset.asset_tag,
                    "name": asset.name,
                    "type": exception,
                    "details": asset.exception_details,
                }));
            }

            if !asset.expected {
                unexpected_assets.push(json!({
                    "asset_id": asset.asset_id,
                    "asset_tag": asset.asset_tag,
                    "name": asset.name,
                    "status": asset.status_found,
                }));
            }
        }

        let mut issues = Map::new();
        if !missing_assets.is_empty() {
            issues.insert("missing_assets".into(), Value::Array(missing_assets));
        }
        if !attention_assets.is_empty() {
            issues.insert("attention_assets".into(), Value::Array(attention_assets));
        }
        if !exceptions.is_empty() {
            issues.insert("exceptions".into(), Value::Array(exceptions));
        }
        if !unexpected_assets.is_empty() {
            issues.insert("unexpected_assets".into(), Value::Array(unexpected_assets));
        }

        let status = if issues.contains_key("missing_assets")
            || issues.contains_key("attention_assets")
        {
            "attention"
        } else if issues.contains_key("exceptions") || issues.contains_key("unexpected_assets") {
            "attention"
        } else {
            "all-good"
        };

        let mut payload = Map::new();
        payload.insert(
            "audit_type".into(),
            Value::String(match self.mode {
                AuditMode::FullZone => "full-zone".to_string(),
                AuditMode::SpotCheck => "spot-check".to_string(),
            }),
        );
        if let Some(zone) = &self.zone_info {
            payload.insert("zone_id".into(), json!(zone.id));
        }
        if !self.audit_name.trim().is_empty() {
            payload.insert("audit_name".into(), json!(self.audit_name.trim()));
        }
        payload.insert("started_by".into(), json!(user_id));
        payload.insert(
            "started_at".into(),
            json!(started_at.format("%Y-%m-%d %H:%M:%S").to_string()),
        );
        payload.insert(
            "completed_at".into(),
            json!(completed_at.format("%Y-%m-%d %H:%M:%S").to_string()),
        );
        payload.insert("status".into(), json!(status));
        if let Some(timeout) = self.timeout_minutes {
            payload.insert("timeout_minutes".into(), json!(timeout));
        }
        if issues.is_empty() {
            payload.insert("issues_found".into(), Value::Null);
        } else {
            payload.insert("issues_found".into(), Value::Object(issues));
        }
        payload.insert("assets_expected".into(), json!(required_total as i64));
        payload.insert("assets_found".into(), json!(found_count as i64));
        if !self.notes.trim().is_empty() {
            payload.insert("notes".into(), json!(self.notes.trim()));
        }

        let audit_insert = api_client.insert("physical_audits", Value::Object(payload))?;
        if !audit_insert.success {
            return Err(anyhow!(
                "Failed to create audit session: {}",
                audit_insert
                    .error
                    .unwrap_or_else(|| "unknown error".to_string())
            ));
        }
        let audit_id = audit_insert.data.unwrap_or(0) as i64;

        // Insert audit logs
        for asset in &self.expected_assets {
            let mut log_payload = Map::new();
            log_payload.insert("physical_audit_id".into(), json!(audit_id));
            log_payload.insert("asset_id".into(), json!(asset.asset_id));
            log_payload.insert("status_found".into(), json!(asset.status_found));
            if let Some(task_id) = asset.audit_task_id {
                log_payload.insert("audit_task_id".into(), json!(task_id));
            }
            if let Some(responses) = &asset.task_responses {
                log_payload.insert("audit_task_responses".into(), responses.clone());
            }
            if let Some(exception) = &asset.exception_type {
                log_payload.insert("exception_type".into(), json!(exception));
            }
            if let Some(details) = &asset.exception_details {
                log_payload.insert("exception_details".into(), json!(details));
            }
            if let Some(zone) = &self.zone_info {
                log_payload.insert("found_in_zone_id".into(), json!(zone.id));
            }
            if !asset.notes.trim().is_empty() {
                log_payload.insert("notes".into(), json!(asset.notes.trim()));
            }
            let log_insert =
                api_client.insert("physical_audit_logs", Value::Object(log_payload))?;
            if !log_insert.success {
                return Err(anyhow!(
                    "Failed to record audit log for asset {}",
                    asset.asset_tag
                ));
            }
        }

        let completion = AuditCompletion {
            audit_id,
            status: status.to_string(),
        };
        self.completion_snapshot = Some(completion);
        self.has_recent_completion = true;
        self.reset_core_state();
        Ok(())
    }
}

#[derive(Debug, Clone)]
struct AuditTaskDefinition {
    steps: Vec<AuditTaskStep>,
    index_by_step: HashMap<i64, usize>,
}

impl AuditTaskDefinition {
    fn from_value(value: Value) -> Result<Self> {
        let sequence_value = match value {
            Value::Object(ref obj) if obj.contains_key("json_sequence") => {
                obj.get("json_sequence").cloned().unwrap_or(Value::Null)
            }
            other => other,
        };

        let normalized_sequence = match sequence_value {
            Value::String(ref s) => {
                if let Ok(bytes) = BASE64_STANDARD.decode(s) {
                    serde_json::from_slice::<Value>(&bytes).map_err(|err| {
                        let raw_debug = String::from_utf8_lossy(&bytes).into_owned();
                        anyhow!(
                            "Invalid audit task JSON sequence: {}\nDecoded payload: {}",
                            err,
                            raw_debug
                        )
                    })?
                } else if let Ok(parsed) = serde_json::from_str::<Value>(s) {
                    parsed
                } else {
                    return Err(anyhow!(
                        "Invalid audit task JSON sequence: expected array but got string '{}'.",
                        s
                    ));
                }
            }
            other => other,
        };

        let raw_debug = serde_json::to_string_pretty(&normalized_sequence)
            .unwrap_or_else(|_| normalized_sequence.to_string());
        let steps: Vec<AuditTaskStep> = serde_json::from_value(normalized_sequence.clone())
            .map_err(|err| {
                anyhow!(
                    "Invalid audit task JSON sequence: {}\nSequence payload: {}",
                    err,
                    raw_debug
                )
            })?;
        if steps.is_empty() {
            return Err(anyhow!("Audit task contains no steps"));
        }
        let mut index_by_step = HashMap::new();
        for (idx, step) in steps.iter().enumerate() {
            index_by_step.insert(step.step, idx);
        }
        Ok(Self {
            steps,
            index_by_step,
        })
    }

    fn first_step(&self) -> i64 {
        self.steps.first().map(|s| s.step).unwrap_or(1)
    }

    fn get_step(&self, step_id: i64) -> Option<&AuditTaskStep> {
        self.index_by_step
            .get(&step_id)
            .and_then(|idx| self.steps.get(*idx))
    }

    fn next_step(&self, current_id: i64) -> Option<i64> {
        if let Some(idx) = self.index_by_step.get(&current_id) {
            self.steps.get(idx + 1).map(|s| s.step)
        } else {
            None
        }
    }
}

#[derive(Debug, Clone, Deserialize)]
struct AuditTaskStep {
    step: i64,
    question: String,
    #[serde(rename = "type")]
    question_type: AuditQuestionType,
    #[serde(default)]
    options: Vec<String>,
    #[serde(default)]
    actions: HashMap<String, AuditTaskAction>,
}

#[derive(Debug, Clone, Deserialize)]
#[serde(rename_all = "snake_case")]
enum AuditQuestionType {
    YesNo,
    MultipleChoice,
    TextInput,
}

#[derive(Debug, Clone, Deserialize, Default)]
struct AuditTaskAction {
    #[serde(default)]
    next_step: Option<i64>,
    #[serde(default)]
    set_status: Option<String>,
    #[serde(default)]
    set_additional_fields: Option<HashMap<String, String>>,
    #[serde(default)]
    end_audit: Option<bool>,
}

#[derive(Debug, Clone)]
struct AuditTaskRunner {
    definition: AuditTaskDefinition,
    current_step: i64,
    responses: Vec<TaskResponseEntry>,
    is_open: bool,
    user_input: String,
    asset_label: String,
    collected_fields: Map<String, Value>,
    status_override: Option<String>,
}

impl AuditTaskRunner {
    fn new(definition: AuditTaskDefinition, asset_label: String) -> Self {
        let first_step = definition.first_step();
        Self {
            definition,
            current_step: first_step,
            responses: Vec::new(),
            is_open: true,
            user_input: String::new(),
            asset_label,
            collected_fields: Map::new(),
            status_override: None,
        }
    }

    fn is_open(&self) -> bool {
        self.is_open
    }

    fn show(&mut self, ctx: &egui::Context) -> Option<AuditTaskOutcome> {
        if !self.is_open {
            return None;
        }

        let mut keep_open = self.is_open;
        let mut completed: Option<AuditTaskOutcome> = None;

        let title = format!("Audit Task – {}", self.asset_label);
        egui::Window::new(title)
            .id(egui::Id::new("audit_task_runner_window"))
            .collapsible(false)
            .resizable(false)
            .anchor(egui::Align2::CENTER_CENTER, [0.0, 0.0])
            .open(&mut keep_open)
            .show(ctx, |ui| {
                if let Some(step) = self.definition.get_step(self.current_step).cloned() {
                    ui.heading(&step.question);
                    ui.add_space(8.0);

                    match step.question_type {
                        AuditQuestionType::YesNo => {
                            ui.horizontal(|ui| {
                                if ui.button("Yes").clicked() {
                                    completed = self.handle_answer(
                                        &step,
                                        "yes",
                                        Value::String("Yes".to_string()),
                                        None,
                                    );
                                }
                                if ui.button("No").clicked() {
                                    completed = self.handle_answer(
                                        &step,
                                        "no",
                                        Value::String("No".to_string()),
                                        None,
                                    );
                                }
                            });
                        }
                        AuditQuestionType::MultipleChoice => {
                            for option in &step.options {
                                if ui.button(option).clicked() {
                                    completed = self.handle_answer(
                                        &step,
                                        option,
                                        Value::String(option.clone()),
                                        None,
                                    );
                                    if completed.is_some() {
                                        break;
                                    }
                                }
                            }
                        }
                        AuditQuestionType::TextInput => {
                            ui.label("Answer:");
                            ui.add(
                                egui::TextEdit::singleline(&mut self.user_input)
                                    .desired_width(280.0),
                            );
                            if ui.button("Submit").clicked() {
                                let answer_value = Value::String(self.user_input.clone());
                                completed = self.handle_answer(
                                    &step,
                                    "any",
                                    answer_value,
                                    Some(self.user_input.clone()),
                                );
                                self.user_input.clear();
                            }
                        }
                    }
                } else {
                    // No step found; close gracefully
                    completed = Some(self.finish());
                }
            });

        if !keep_open {
            self.is_open = false;
        }

        if let Some(result) = completed {
            self.is_open = false;
            Some(result)
        } else {
            None
        }
    }

    fn handle_answer(
        &mut self,
        step: &AuditTaskStep,
        answer_key: &str,
        answer_value: Value,
        user_input: Option<String>,
    ) -> Option<AuditTaskOutcome> {
        self.responses.push(TaskResponseEntry {
            step: step.step,
            question: step.question.clone(),
            answer: answer_value.clone(),
        });

        let key_lower = answer_key.to_lowercase();
        let action = step
            .actions
            .get(&key_lower)
            .or_else(|| step.actions.get(answer_key))
            .or_else(|| step.actions.get("any"));

        if let Some(act) = action {
            if let Some(ref status) = act.set_status {
                self.status_override = Some(status.clone());
            }
            if let Some(ref fields) = act.set_additional_fields {
                for (field, template) in fields {
                    let value = if let Some(ref input) = user_input {
                        template.replace("{user_input}", input)
                    } else {
                        template.clone()
                    };
                    self.collected_fields
                        .insert(field.clone(), Value::String(value));
                }
            }
            if act.end_audit.unwrap_or(false) {
                return Some(self.finish());
            }
            if let Some(next_step) = act.next_step {
                self.current_step = next_step;
                return None;
            }
        }

        if let Some(next) = self.definition.next_step(step.step) {
            self.current_step = next;
            None
        } else {
            Some(self.finish())
        }
    }

    fn finish(&mut self) -> AuditTaskOutcome {
        let responses = Value::Array(
            self.responses
                .iter()
                .map(|entry| {
                    json!({
                        "step": entry.step,
                        "question": entry.question,
                        "answer": entry.answer,
                    })
                })
                .collect(),
        );
        AuditTaskOutcome {
            status_override: self.status_override.clone(),
            additional_fields: self.collected_fields.clone(),
            responses,
        }
    }
}