import React, { useState, useEffect, useRef, useCallback } from 'react';
import { createPortal } from 'react-dom';
import { useNavigate, useSearchParams } from 'react-router-dom';
import axiosInstance from '../../../utils/axiosConfig';
import {
    FORMS_URL,
    SUBMISSIONS_URL,
    AUDITS_URL,
    QUESTION_TYPE_URL,
    RPN_LIKELIHOOD_URL,
    RPN_SEVERITY_URL,
    AI_FILL_FROM_IMAGE_URL
} from '../../../utils/config';
import { addCachedObservation } from '../../../utils/createAuditObservationsCache';
import SuccessPopup from '../../common/SuccessPopup';
import ErrorPopup from '../../common/ErrorPopup';
import { useScrollToError } from '../../../hooks/useScrollToError';
import { getUserIdFromToken } from '../../../utils/authUtils';
import './styles/FormSubmission.css';

const MAX_CANVAS_WIDTH = 640;
const MAX_CANVAS_HEIGHT = 480;

// Custom multi-select dropdown with chips and expandable panel
function MultiSelectDropdown({ options, selected, onChange, placeholder = 'Select options...' }) {
    const [isOpen, setIsOpen] = useState(false);
    const ref = useRef(null);

    useEffect(() => {
        const handleClickOutside = (e) => {
            if (ref.current && !ref.current.contains(e.target)) setIsOpen(false);
        };
        document.addEventListener('mousedown', handleClickOutside);
        return () => document.removeEventListener('mousedown', handleClickOutside);
    }, []);

    const toggleOption = (opt) => {
        const next = selected.includes(opt)
            ? selected.filter(o => o !== opt)
            : [...selected, opt];
        onChange(next.join(', '));
    };

    return (
        <div className="multi-select-dropdown" ref={ref}>
            <button
                type="button"
                className={`multi-select-trigger ${isOpen ? 'open' : ''}`}
                onClick={() => setIsOpen(!isOpen)}
                aria-expanded={isOpen}
                aria-haspopup="listbox"
                aria-label={placeholder}
            >
                <span className={`multi-select-value ${selected.length === 0 ? 'placeholder' : ''}`}>
                    {selected.length > 0
                        ? selected.map((s) => (
                            <span key={s} className="multi-select-chip">
                                {s}
                                <button
                                    type="button"
                                    className="multi-select-chip-remove"
                                    onClick={(e) => { e.stopPropagation(); toggleOption(s); }}
                                    aria-label={`Remove ${s}`}
                                >
                                    ×
                                </button>
                            </span>
                        ))
                        : placeholder}
                </span>
                <span className="multi-select-arrow" aria-hidden>▼</span>
            </button>
            {isOpen && (
                <div className="multi-select-panel" role="listbox">
                    {options.map((opt) => (
                        <div
                            key={opt}
                            role="option"
                            aria-selected={selected.includes(opt)}
                            className={`multi-select-option ${selected.includes(opt) ? 'selected' : ''}`}
                            onClick={() => toggleOption(opt)}
                            onKeyDown={(e) => {
                                if (e.key === 'Enter' || e.key === ' ') {
                                    e.preventDefault();
                                    toggleOption(opt);
                                }
                            }}
                            tabIndex={0}
                        >
                            <span className="multi-select-checkbox">{selected.includes(opt) ? '✓' : ''}</span>
                            <span>{opt}</span>
                        </div>
                    ))}
                </div>
            )}
        </div>
    );
}

// Modal for drawing on an image (mark color/lines)
function ImageAnnotationModal({ isOpen, onClose, imageSrc, onDone, defaultColor = '#e53935' }) {
    const canvasRef = useRef(null);
    const isDrawingRef = useRef(false);
    const lastPosRef = useRef({ x: 0, y: 0 });
    const imageObjRef = useRef(null);
    const [color, setColor] = useState(defaultColor);
    const [brushSize, setBrushSize] = useState(4);

    useEffect(() => {
        if (isOpen) setColor(defaultColor);
    }, [isOpen, defaultColor]);

    const drawImageToCanvas = useCallback((img) => {
        const canvas = canvasRef.current;
        if (!canvas || !img) return;
        const scale = Math.min(MAX_CANVAS_WIDTH / img.width, MAX_CANVAS_HEIGHT / img.height, 1);
        canvas.width = Math.round(img.width * scale);
        canvas.height = Math.round(img.height * scale);
        const ctx = canvas.getContext('2d');
        ctx.drawImage(img, 0, 0, canvas.width, canvas.height);
    }, []);

    useEffect(() => {
        if (!isOpen || !imageSrc) return;
        const img = new Image();
        img.crossOrigin = 'anonymous';
        img.onload = () => {
            imageObjRef.current = img;
            drawImageToCanvas(img);
        };
        img.src = imageSrc;
    }, [isOpen, imageSrc, drawImageToCanvas]);

    const getCanvasCoords = (e) => {
        const canvas = canvasRef.current;
        if (!canvas) return null;
        const rect = canvas.getBoundingClientRect();
        const scaleX = canvas.width / rect.width;
        const scaleY = canvas.height / rect.height;
        const clientX = e.touches ? e.touches[0].clientX : e.clientX;
        const clientY = e.touches ? e.touches[0].clientY : e.clientY;
        return {
            x: (clientX - rect.left) * scaleX,
            y: (clientY - rect.top) * scaleY
        };
    };

    const startDraw = (e) => {
        e.preventDefault();
        const coords = getCanvasCoords(e);
        if (!coords) return;
        isDrawingRef.current = true;
        lastPosRef.current = coords;
    };

    const moveDraw = (e) => {
        e.preventDefault();
        if (!isDrawingRef.current) return;
        const canvas = canvasRef.current;
        const coords = getCanvasCoords(e);
        if (!canvas || !coords) return;
        const ctx = canvas.getContext('2d');
        ctx.strokeStyle = color;
        ctx.lineWidth = brushSize;
        ctx.lineCap = 'round';
        ctx.lineJoin = 'round';
        ctx.beginPath();
        ctx.moveTo(lastPosRef.current.x, lastPosRef.current.y);
        ctx.lineTo(coords.x, coords.y);
        ctx.stroke();
        lastPosRef.current = coords;
    };

    const endDraw = (e) => {
        e.preventDefault();
        isDrawingRef.current = false;
    };

    const handleClear = () => {
        const img = imageObjRef.current;
        if (img) drawImageToCanvas(img);
    };

    const handleDone = () => {
        const canvas = canvasRef.current;
        if (!canvas) return;
        const dataUrl = canvas.toDataURL('image/png');
        onDone(dataUrl);
        onClose();
    };

    if (!isOpen) return null;

    const modalContent = (
        <div className="image-annotation-overlay" onClick={(e) => e.target === e.currentTarget && onClose()}>
            <div className="image-annotation-modal" onClick={(e) => e.stopPropagation()}>
                <div className="image-annotation-header">
                    <h3 className="image-annotation-title">Mark on image</h3>
                    <button type="button" className="image-annotation-close" onClick={onClose} aria-label="Close">×</button>
                </div>
                <div className="image-annotation-toolbar">
                    <label className="image-annotation-tool">
                        <span>Color</span>
                        <input type="color" value={color} onChange={(e) => setColor(e.target.value)} />
                    </label>
                    <label className="image-annotation-tool">
                        <span>Size</span>
                        <input type="range" min="1" max="24" value={brushSize} onChange={(e) => setBrushSize(Number(e.target.value))} />
                        <span className="image-annotation-size-value">{brushSize}px</span>
                    </label>
                    <button type="button" className="btn btn-secondary image-annotation-clear" onClick={handleClear}>Clear marks</button>
                </div>
                <div className="image-annotation-canvas-wrap">
                    <canvas
                        ref={canvasRef}
                        className="image-annotation-canvas"
                        onMouseDown={startDraw}
                        onMouseMove={moveDraw}
                        onMouseUp={endDraw}
                        onMouseLeave={endDraw}
                        onTouchStart={startDraw}
                        onTouchMove={moveDraw}
                        onTouchEnd={endDraw}
                    />
                </div>
                <div className="image-annotation-actions">
                    <button type="button" className="btn btn-secondary" onClick={onClose}>Cancel</button>
                    <button type="button" className="btn btn-primary" onClick={handleDone}>Done</button>
                </div>
            </div>
        </div>
    );

    return createPortal(modalContent, document.body);
}

const FormSubmission = () => {
    const navigate = useNavigate();
    const [searchParams] = useSearchParams();
    const formId = searchParams.get('form_id');
    const auditId = searchParams.get('audit_id');
    const openingSubIdFromUrl = searchParams.get('opening_sub_id');
    const returnToUrl = searchParams.get('return_to');
    const fromCreateAudit = searchParams.get('from_create_audit') === '1';

    const [formData, setFormData] = useState(null);
    const [questionTypes, setQuestionTypes] = useState([]);
    const [answers, setAnswers] = useState({});
    const [mediaFiles, setMediaFiles] = useState({});
    const [loading, setLoading] = useState(false);
    const [fetchingData, setFetchingData] = useState(true);
    const [error, setError] = useState(null);
    useScrollToError(error);
    const [allQuestions, setAllQuestions] = useState([]);

    // RPN related state
    const [rpnLikelihoodOptions, setRpnLikelihoodOptions] = useState([]);
    const [rpnSeverityOptions, setRpnSeverityOptions] = useState([]);
    const [likelihoodQuestionId, setLikelihoodQuestionId] = useState(null);
    const [severityQuestionId, setSeverityQuestionId] = useState(null);
    const [rpnQuestionId, setRpnQuestionId] = useState(null);

    // Success popup after form submission
    const [showPopup, setShowPopup] = useState(false);
    const [popupMessage, setPopupMessage] = useState('');
    const [lastCreatedSubmissionId, setLastCreatedSubmissionId] = useState(null);

    // For Closing forms: link this submission to an opening observation
    const [openingSubmissions, setOpeningSubmissions] = useState([]);
    const [selectedOpeningSubId, setSelectedOpeningSubId] = useState('');

    // Image annotation modal: mark/draw on attached images
    const [annotationModal, setAnnotationModal] = useState({ open: false, questionId: null, fileIndex: null });

    // AI fill from image (Ollama)
    const [aiFillLoading, setAiFillLoading] = useState(false);

    useEffect(() => {
        const fetchData = async () => {
            if (!formId) {
                setError('No form ID provided');
                setFetchingData(false);
                return;
            }

            try {
                setFetchingData(true);
                const [
                    formRes,
                    questionTypesRes,
                    rpnLikelihoodRes,
                    rpnSeverityRes
                ] = await Promise.all([
                    axiosInstance.get(`${FORMS_URL}/${formId}`),
                    axiosInstance.get(QUESTION_TYPE_URL),
                    axiosInstance.get(RPN_LIKELIHOOD_URL),
                    axiosInstance.get(RPN_SEVERITY_URL)
                ]);

                if (formRes.data.success && formRes.data.data) {
                    const form = formRes.data.data;
                    setFormData(form);
                    // For Closing forms (form_type_id 2), load all opening submissions for this audit so user can link
                    if (form.form_type_id === 2 && auditId) {
                        try {
                            const subRes = await axiosInstance.get(`${SUBMISSIONS_URL}?audit_id=${auditId}`);
                            if (subRes.data?.success && Array.isArray(subRes.data.data)) {
                                const openingOnly = subRes.data.data.filter(
                                    (sub) => Number(sub.form_type_id) === 1
                                ).map((s) => ({ submission_id: s.submission_id, form_title: s.form_title }));
                                setOpeningSubmissions(openingOnly);
                                // Pre-select opening observation when coming from "Close" flow (opening_sub_id in URL)
                                if (openingSubIdFromUrl && openingOnly.some((s) => String(s.submission_id) === String(openingSubIdFromUrl))) {
                                    setSelectedOpeningSubId(String(openingSubIdFromUrl));
                                }
                            }
                        } catch (_) {
                            // ignore
                        }
                    }
                    const questions = form.sat_questions || form.questions || [];
                    setAllQuestions(questions);

                    // Identify special RPN-related questions by their text
                    const toLower = (str) => (str || '').toString().toLowerCase();
                    const likelihoodQuestion = questions.find(q =>
                        toLower(q.question).includes('likelihood of harm')
                    );
                    const severityQuestion = questions.find(q =>
                        toLower(q.question).includes('severity of harm')
                    );
                    const rpnQuestion = questions.find(q =>
                        toLower(q.question).includes('rpn')
                    );

                    if (likelihoodQuestion) {
                        setLikelihoodQuestionId(likelihoodQuestion.question_id);
                    }
                    if (severityQuestion) {
                        setSeverityQuestionId(severityQuestion.question_id);
                    }
                    if (rpnQuestion) {
                        setRpnQuestionId(rpnQuestion.question_id);
                    }
                    
                    // Initialize answers object
                    const initialAnswers = {};
                    questions.forEach(q => {
                        initialAnswers[q.question_id] = {
                            question_id: q.question_id,
                            is_media: false,
                            answer_value: '',
                            media: []
                        };
                    });
                    setAnswers(initialAnswers);
                } else {
                    setError('Failed to load form details.');
                }

                setQuestionTypes(questionTypesRes.data.data || []);
                setRpnLikelihoodOptions(rpnLikelihoodRes.data.data || []);
                setRpnSeverityOptions(rpnSeverityRes.data.data || []);
            } catch (error) {
                console.error('Error fetching data:', error);
                setError('Failed to load form data. Please refresh the page.');
            } finally {
                setFetchingData(false);
            }
        };

        fetchData();
    }, [formId, auditId, openingSubIdFromUrl]);

    // Build ordered list for display: main questions (with their child questions grouped under them) and standalone questions
    const questionById = React.useMemo(() => {
        const m = new Map();
        (allQuestions || []).forEach((q) => m.set(q.question_id, q));
        return m;
    }, [allQuestions]);

    const displayQuestions = React.useMemo(() => {
        const list = [];
        (allQuestions || []).forEach((q) => {
            if (q.main_question_id != null) return; // child: rendered under its main
            if ((q.child_question_ids || []).length > 0) {
                list.push({ question: q, isChild: false });
                (q.child_question_ids || []).forEach((cid) => {
                    const cq = questionById.get(cid);
                    if (cq) list.push({ question: cq, isChild: true });
                });
            } else {
                list.push({ question: q, isChild: false });
            }
        });
        return list;
    }, [allQuestions, questionById]);

    const getQuestionTypeName = (questionTypeId) => {
        const type = questionTypes.find(t => String(t.question_type_id) === String(questionTypeId));
        return type ? type.question_type : 'Text';
    };

    /** For Single Selection Dropdown / Multiple Dropdown: first line = prompt, rest = options. */
    const getPromptAndOptions = (questionText) => {
        const lines = (questionText || '').split(/\r?\n/).map(s => s.trim()).filter(Boolean);
        const prompt = lines[0] || '';
        const options = lines.slice(1);
        return { prompt, options };
    };

    const getQuestionLabel = (question) => {
        const typeName = getQuestionTypeName(question.question_type_id);
        if (typeName === 'MMCQ' || typeName === 'Single Selection Dropdown' || typeName === 'Multiple Dropdown Selection') {
            const { prompt } = getPromptAndOptions(question.question);
            return prompt || question.question;
        }
        return question.question;
    };

    const computeRpnValue = (likelihoodVal, severityVal) => {
        const l = Number(likelihoodVal);
        const s = Number(severityVal);
        const hasBoth = !isNaN(l) && !isNaN(s);
        return hasBoth ? String(l * s) : '';
    };

    /** RPN colour: 1 < x <= 8 green, 8 < x <= 32 yellow, x > 32 red */
    const getRpnColorClass = (rpnValue) => {
        const x = Number(rpnValue);
        if (rpnValue === '' || rpnValue == null || isNaN(x)) return '';
        if (x < 8) return 'rpn-green';
        if (x < 32) return 'rpn-yellow';
        return 'rpn-red';
    };

    const handleAnswerChange = (questionId, value) => {
        setAnswers(prev => ({
            ...prev,
            [questionId]: {
                ...prev[questionId],
                answer_value: value,
                // Do NOT force is_media to false here; keep existing value
                is_media: prev[questionId]?.is_media || false
            }
        }));
    };

    const handleLikelihoodChange = (value) => {
        if (!likelihoodQuestionId) return;

        setAnswers(prev => {
            const updated = {
                ...prev,
                [likelihoodQuestionId]: {
                    ...(prev[likelihoodQuestionId] || {
                        question_id: likelihoodQuestionId,
                        is_media: false,
                        media: []
                    }),
                    is_media: false,
                    answer_value: value
                }
            };
            const severityVal = updated[severityQuestionId]?.answer_value;
            const rpnValue = rpnQuestionId ? computeRpnValue(value, severityVal) : null;

            if (rpnQuestionId && rpnValue !== null) {
                updated[rpnQuestionId] = {
                    ...(prev[rpnQuestionId] || {
                        question_id: rpnQuestionId,
                        is_media: false,
                        media: []
                    }),
                    is_media: false,
                    answer_value: rpnValue
                };
            }
            return updated;
        });
    };

    const handleSeverityChange = (value) => {
        if (!severityQuestionId) return;

        setAnswers(prev => {
            const updated = {
                ...prev,
                [severityQuestionId]: {
                    ...(prev[severityQuestionId] || {
                        question_id: severityQuestionId,
                        is_media: false,
                        media: []
                    }),
                    is_media: false,
                    answer_value: value
                }
            };
            const likelihoodVal = updated[likelihoodQuestionId]?.answer_value;
            const rpnValue = rpnQuestionId ? computeRpnValue(likelihoodVal, value) : null;

            if (rpnQuestionId && rpnValue !== null) {
                updated[rpnQuestionId] = {
                    ...(prev[rpnQuestionId] || {
                        question_id: rpnQuestionId,
                        is_media: false,
                        media: []
                    }),
                    is_media: false,
                    answer_value: rpnValue
                };
            }
            return updated;
        });
    };

    const handleMediaChange = async (questionId, files) => {
        const fileArray = Array.from(files);
        const base64Promises = fileArray.map(file => {
            return new Promise((resolve, reject) => {
                const reader = new FileReader();
                reader.onload = () => resolve(reader.result);
                reader.onerror = reject;
                reader.readAsDataURL(file);
            });
        });

        try {
            const base64Array = await Promise.all(base64Promises);
            setMediaFiles(prev => ({
                ...prev,
                [questionId]: fileArray
            }));
            setAnswers(prev => ({
                ...prev,
                [questionId]: {
                    ...prev[questionId],
                    is_media: true,
                    answer_value: prev[questionId]?.answer_value ?? null,
                    media: base64Array
                }
            }));
        } catch (error) {
            console.error('Error converting files to base64:', error);
            setError('Failed to process media files.');
        }
    };

    const removeMedia = (questionId) => {
        setMediaFiles(prev => {
            const newFiles = { ...prev };
            delete newFiles[questionId];
            return newFiles;
        });
        setAnswers(prev => ({
            ...prev,
            [questionId]: {
                ...prev[questionId],
                is_media: false,
                answer_value: prev[questionId]?.answer_value ?? '',
                media: []
            }
        }));
    };

    const removeMediaAtIndex = (questionId, index) => {
        const files = mediaFiles[questionId];
        const media = answers[questionId]?.media || [];
        if (!files || !media.length) return;
        if (files.length <= 1 && media.length <= 1) {
            removeMedia(questionId);
            return;
        }
        setMediaFiles(prev => ({
            ...prev,
            [questionId]: prev[questionId].filter((_, i) => i !== index)
        }));
        setAnswers(prev => {
            const newMedia = (prev[questionId]?.media || []).filter((_, i) => i !== index);
            return {
                ...prev,
                [questionId]: {
                    ...prev[questionId],
                    is_media: newMedia.length > 0,
                    answer_value: prev[questionId]?.answer_value ?? '',
                    media: newMedia
                }
            };
        });
    };

    const openAnnotationModal = (questionId, fileIndex) => {
        setAnnotationModal({ open: true, questionId, fileIndex });
    };

    const applyAnnotatedImage = (newDataUrl) => {
        const { questionId, fileIndex } = annotationModal;
        if (questionId == null || fileIndex == null) return;
        setAnswers(prev => ({
            ...prev,
            [questionId]: {
                ...prev[questionId],
                media: (prev[questionId]?.media || []).map((m, i) => (i === fileIndex ? newDataUrl : m))
            }
        }));
        setAnnotationModal({ open: false, questionId: null, fileIndex: null });
    };

    /**
     * Get the first uploaded image from the form state (in-memory).
     * The image is not in the database until the user submits the form;
     * we read it from answers[].media set by handleMediaChange.
     */
    const getFirstUploadedImage = () => {
        for (const entry of Object.entries(answers)) {
            const [, a] = entry;
            if (a?.is_media && a?.media?.length > 0 && typeof a.media[0] === 'string') {
                return a.media[0];
            }
        }
        return null;
    };

    /** Resize/compress image to reduce payload and speed up AI. Returns data URL (JPEG). */
    const compressImageForAi = (dataUrl, maxSize = 1024, quality = 0.82) => {
        return new Promise((resolve, reject) => {
            const img = new Image();
            img.crossOrigin = 'anonymous';
            img.onload = () => {
                const w = img.width;
                const h = img.height;
                const scale = Math.min(1, maxSize / Math.max(w, h));
                const cw = Math.round(w * scale);
                const ch = Math.round(h * scale);
                const canvas = document.createElement('canvas');
                canvas.width = cw;
                canvas.height = ch;
                const ctx = canvas.getContext('2d');
                ctx.drawImage(img, 0, 0, cw, ch);
                try {
                    const jpeg = canvas.toDataURL('image/jpeg', quality);
                    resolve(jpeg);
                } catch (e) {
                    resolve(dataUrl);
                }
            };
            img.onerror = () => reject(new Error('Failed to load image'));
            img.src = dataUrl;
        });
    };

    /** Question IDs for AI fill: those with ai_fill === true, or all question IDs if none marked. */
    const allQuestionIds = (formData?.sat_questions || formData?.questions)?.map((q) => q.question_id) || [];
    const aiFillQuestionIds = (formData?.sat_questions || formData?.questions)?.filter((q) => q.ai_fill === true).map((q) => q.question_id) || [];
    const questionIdsToFill = aiFillQuestionIds.length > 0 ? aiFillQuestionIds : allQuestionIds;

    const handleFillFromAi = async () => {
        const imageFromState = getFirstUploadedImage();
        if (!imageFromState) {
            setError('Attach an image first (e.g. use “Attach Media” on a question). The image is taken from your current form and is not sent to the server until you submit.');
            return;
        }
        if (questionIdsToFill.length === 0) {
            setError('No questions in this form to fill.');
            return;
        }
        setAiFillLoading(true);
        setError(null);
        try {
            const imageBase64 = await compressImageForAi(imageFromState);
            const response = await axiosInstance.post(
                AI_FILL_FROM_IMAGE_URL,
                {
                    form_id: parseInt(formId, 10),
                    image_base64: imageBase64,
                    question_ids: questionIdsToFill
                },
                { timeout: 180000 }
            );
            if (!response.data?.success || !Array.isArray(response.data.answers)) {
                setError(response.data?.message || 'AI fill failed.');
                return;
            }
            const answersFromAi = response.data.answers;
            if (answersFromAi.length === 0) {
                setError('AI returned no answers for this image. Try a clearer image or check server logs.');
                return;
            }
            setAnswers((prev) => {
                const next = { ...prev };
                for (const { question_id, answer_value } of answersFromAi) {
                    const qId = Number(question_id);
                    if (Number.isNaN(qId) || next[qId] == null) continue;
                    next[qId] = {
                        ...next[qId],
                        answer_value: answer_value ?? ''
                    };
                }
                if (likelihoodQuestionId && severityQuestionId && rpnQuestionId && next[rpnQuestionId]) {
                    const lVal = next[likelihoodQuestionId]?.answer_value;
                    const sVal = next[severityQuestionId]?.answer_value;
                    const rpnValue = computeRpnValue(lVal, sVal);
                    if (rpnValue !== '') {
                        next[rpnQuestionId] = {
                            ...next[rpnQuestionId],
                            answer_value: rpnValue
                        };
                    }
                }
                return next;
            });
        } catch (err) {
            let msg = err.response?.data?.message || err.message || 'Fill from AI failed.';
            if (err.code === 'ECONNABORTED' || (err.message && err.message.toLowerCase().includes('timeout'))) {
                msg = 'Request timed out. Ollama can take 1–2 minutes (especially on first run). Ensure Ollama is running (ollama serve) and llava is pulled (ollama pull llava), then try again with a smaller image.';
            } else if (!err.response?.data?.message) {
                msg = 'Fill from AI failed. Is Ollama running (ollama serve) and llava pulled (ollama pull llava)?';
            }
            setError(msg);
        } finally {
            setAiFillLoading(false);
        }
    };

    const handleSubmit = async (e) => {
        e.preventDefault();
        setLoading(true);
        setError(null);

        const userId = getUserIdFromToken();
        if (!userId) {
            setError('User not authenticated. Please login again.');
            setLoading(false);
            return;
        }

        // Create-audit draft flow: save to cache instead of API
        if (fromCreateAudit) {
            if (formData.form_type_id !== 1) {
                setError('Only opening forms can be added during audit creation.');
                setLoading(false);
                return;
            }
            const answersArray = Object.values(answers).filter(answer => {
                return answer.answer_value || (answer.is_media && answer.media && answer.media.length > 0);
            });
            addCachedObservation({
                form_id: parseInt(formId, 10),
                form_title: formData.title || `Form ${formId}`,
                submitter_id: userId,
                sat_answers: answersArray
            });
            setPopupMessage('Observation added to draft. Add at least 2 observations, then create the audit.');
            setShowPopup(true);
            setLoading(false);
            return;
        }

        if (!auditId) {
            setError('Audit ID is required. Please go back and select an audit.');
            setLoading(false);
            return;
        }

        try {
            const answersArray = Object.values(answers).filter(answer => {
                // Include answer if it has a value or media
                return answer.answer_value || (answer.is_media && answer.media && answer.media.length > 0);
            });

            const formTypeId = Number(formData.form_type_id) || 1;
            const observation = {
                submitter_id: userId,
                form_id: parseInt(formId, 10),
                submission_type_id: formTypeId,
                sat_answers: answersArray
            };
            const payload = {
                observations: [observation]
            };

            const openingToLink = selectedOpeningSubId || openingSubIdFromUrl;
            if (formTypeId === 2 && openingToLink) {
                payload.opening_sub_id = parseInt(openingToLink, 10);
            }

            console.log('Submitting observation via audit endpoint with payload:', payload);

            const response = await axiosInstance.post(`${AUDITS_URL}/${auditId}/observations`, payload);

            console.log('Observation submission response:', response.data);

            if (response.data.success) {
                const submissionId = response.data.data?.submission_id ?? response.data.submission_id;
                if (submissionId != null) setLastCreatedSubmissionId(submissionId);
                setPopupMessage('Form submitted successfully!');
                setShowPopup(true);
            } else {
                setError(response.data.message || 'Failed to submit observation. Please try again.');
            }
        } catch (error) {
            console.error('Error submitting form:', error);
            console.error('Error response:', error.response);
            
            let errorMessage = 'Failed to submit observation. Please try again.';
            
            if (error.response) {
                errorMessage = error.response.data?.message 
                    || error.response.data?.error 
                    || `Server error: ${error.response.status} ${error.response.statusText}`;
            } else if (error.request) {
                errorMessage = 'No response from server. Please check your connection.';
            } else {
                errorMessage = error.message || 'An error occurred while submitting the form.';
            }
            
            setError(errorMessage);
        } finally {
            setLoading(false);
        }
    };

    const renderQuestionInput = (question) => {
        const questionType = getQuestionTypeName(question.question_type_id);
        const questionTypeId = question.question_type_id;
        const answer = answers[question.question_id] || {};

        // Custom rendering for RPN-related questions
        if (likelihoodQuestionId && question.question_id === likelihoodQuestionId) {
            return (
                <select
                    className="form-input"
                    value={answer.answer_value || ''}
                    onChange={(e) => handleLikelihoodChange(e.target.value)}
                >
                    <option value="">Select likelihood of harm</option>
                    {rpnLikelihoodOptions.map((opt) => (
                        <option key={opt.likelihood_of_harm} value={opt.value}>
                            {opt.likelihood_of_harm} ({opt.value})
                        </option>
                    ))}
                </select>
            );
        }

        if (severityQuestionId && question.question_id === severityQuestionId) {
            return (
                <select
                    className="form-input"
                    value={answer.answer_value || ''}
                    onChange={(e) => handleSeverityChange(e.target.value)}
                >
                    <option value="">Select severity of harm</option>
                    {rpnSeverityOptions.map((opt) => (
                        <option key={opt.severity_of_harm} value={opt.value}>
                            {opt.severity_of_harm} ({opt.value})
                        </option>
                    ))}
                </select>
            );
        }

        if (rpnQuestionId && question.question_id === rpnQuestionId) {
            const rpnColorClass = getRpnColorClass(answer.answer_value);
            return (
                <div className="rpn-input-wrapper">
                    <input
                        type="number"
                        className={`form-input rpn-value-input ${rpnColorClass}`}
                        value={answer.answer_value || ''}
                        readOnly
                        placeholder="RPN value (auto-calculated)"
                    />
                    {answer.answer_value && rpnColorClass && (
                        <span className={`rpn-badge ${rpnColorClass}`}>
                            {rpnColorClass === 'rpn-green' && 'Low risk'}
                            {rpnColorClass === 'rpn-yellow' && 'Medium risk'}
                            {rpnColorClass === 'rpn-red' && 'High risk'}
                        </span>
                    )}
                </div>
            );
        }

        // Handle Single Selection Dropdown by question_type_id (4) - was MMCQ, now single-select dropdown
        if (Number(questionTypeId) === 4) {
            const options = Array.isArray(question.options) && question.options.length > 0
                ? question.options
                : getPromptAndOptions(question.question).options;
            const selectedValue = answer.answer_value || '';
            if (options.length === 0) {
                return (
                    <input
                        type="text"
                        className="form-input"
                        value={selectedValue}
                        onChange={(e) => handleAnswerChange(question.question_id, e.target.value)}
                        placeholder="Add options to this question to enable dropdown selection"
                    />
                );
            }
            return (
                <select
                    className="form-input"
                    value={selectedValue}
                    onChange={(e) => handleAnswerChange(question.question_id, e.target.value)}
                >
                    <option value="">Select an option</option>
                    {options.map((opt) => (
                        <option key={opt} value={opt}>{opt}</option>
                    ))}
                </select>
            );
        }

        // Handle Multiple Dropdown Selection by question_type_id (5) - ensures dropdown appears even if type name lookup fails
        if (Number(questionTypeId) === 5) {
            const options = Array.isArray(question.options) && question.options.length > 0
                ? question.options
                : getPromptAndOptions(question.question).options;
            const selected = (answer.answer_value || '').split(',').map(s => s.trim()).filter(Boolean);
            if (options.length === 0) {
                return (
                    <input
                        type="text"
                        className="form-input"
                        value={answer.answer_value || ''}
                        onChange={(e) => handleAnswerChange(question.question_id, e.target.value)}
                        placeholder="Add options to this question to enable dropdown selection"
                    />
                );
            }
            return (
                <MultiSelectDropdown
                    options={options}
                    selected={selected}
                    onChange={(value) => handleAnswerChange(question.question_id, value)}
                    placeholder="Click to select one or more options"
                />
            );
        }

        // Default rendering for all other questions
        switch (questionType) {
            case 'Text':
                return (
                    <textarea
                        className="form-textarea"
                        value={answer.answer_value || ''}
                        onChange={(e) => handleAnswerChange(question.question_id, e.target.value)}
                        placeholder="Enter your answer"
                        rows="4"
                    />
                );

            case 'Number':
                return (
                    <input
                        type="number"
                        className="form-input"
                        value={answer.answer_value || ''}
                        onChange={(e) => handleAnswerChange(question.question_id, e.target.value)}
                        placeholder="Enter a number"
                    />
                );

            case 'Yes/No':
                return (
                    <div className="yes-no-options">
                        <label className="radio-option">
                            <input
                                type="radio"
                                name={`question-${question.question_id}`}
                                value="Yes"
                                checked={answer.answer_value === 'Yes'}
                                onChange={(e) => handleAnswerChange(question.question_id, e.target.value)}
                            />
                            <span>Yes</span>
                        </label>
                        <label className="radio-option">
                            <input
                                type="radio"
                                name={`question-${question.question_id}`}
                                value="No"
                                checked={answer.answer_value === 'No'}
                                onChange={(e) => handleAnswerChange(question.question_id, e.target.value)}
                            />
                            <span>No</span>
                        </label>
                    </div>
                );

            case 'MMCQ': {
                const options = Array.isArray(question.options) && question.options.length > 0
                    ? question.options
                    : getPromptAndOptions(question.question).options;
                const selected = (answer.answer_value || '').split(',').map(s => s.trim()).filter(Boolean);
                const toggleOption = (opt) => {
                    const has = selected.includes(opt);
                    const next = has ? selected.filter(o => o !== opt) : [...selected, opt];
                    handleAnswerChange(question.question_id, next.join(', '));
                };
                if (options.length === 0) {
                    return (
                        <input
                            type="text"
                            className="form-input"
                            value={answer.answer_value || ''}
                            onChange={(e) => handleAnswerChange(question.question_id, e.target.value)}
                            placeholder="Add options to this question to enable multiple choice selection"
                        />
                    );
                }
                return (
                    <div className="mmcq-options">
                        {options.map((opt) => (
                            <label key={opt} className="mmcq-option">
                                <input
                                    type="checkbox"
                                    checked={selected.includes(opt)}
                                    onChange={() => toggleOption(opt)}
                                />
                                <span>{opt}</span>
                            </label>
                        ))}
                    </div>
                );
            }

            case 'Single Selection Dropdown': {
                const options = Array.isArray(question.options) && question.options.length > 0
                    ? question.options
                    : getPromptAndOptions(question.question).options;
                const selectedValue = answer.answer_value || '';
                if (options.length === 0) {
                    return (
                        <input
                            type="text"
                            className="form-input"
                            value={selectedValue}
                            onChange={(e) => handleAnswerChange(question.question_id, e.target.value)}
                            placeholder="Add options to this question to enable dropdown selection"
                        />
                    );
                }
                return (
                    <select
                        className="form-input"
                        value={selectedValue}
                        onChange={(e) => handleAnswerChange(question.question_id, e.target.value)}
                    >
                        <option value="">Select an option</option>
                        {options.map((opt) => (
                            <option key={opt} value={opt}>{opt}</option>
                        ))}
                    </select>
                );
            }

            case 'Multiple Dropdown Selection': {
                const options = Array.isArray(question.options) && question.options.length > 0
                    ? question.options
                    : getPromptAndOptions(question.question).options;
                const selected = (answer.answer_value || '').split(',').map(s => s.trim()).filter(Boolean);
                if (options.length === 0) {
                    return (
                        <input
                            type="text"
                            className="form-input"
                            value={answer.answer_value || ''}
                            onChange={(e) => handleAnswerChange(question.question_id, e.target.value)}
                            placeholder="Add options to this question to enable dropdown selection"
                        />
                    );
                }
                return (
                    <MultiSelectDropdown
                        options={options}
                        selected={selected}
                        onChange={(value) => handleAnswerChange(question.question_id, value)}
                        placeholder="Click to select one or more options"
                        questionId={question.question_id}
                    />
                );
            }

            case 'Image Mark/Attach':
                return (
                    <p className="image-type-hint">
                        Attach image(s) below. You can mark or draw on them using &quot;Mark on image&quot; after attaching.
                    </p>
                );

            default:
                return (
                    <input
                        type="text"
                        className="form-input"
                        value={answer.answer_value || ''}
                        onChange={(e) => handleAnswerChange(question.question_id, e.target.value)}
                        placeholder="Enter your answer"
                    />
                );
        }
    };

    if (fetchingData) {
        return (
            <div className="form-submission-container">
                <div style={{ textAlign: 'center', padding: '2rem' }}>
                    <p>Loading form...</p>
                </div>
            </div>
        );
    }

    if (error && !formData) {
        return (
            <div className="form-submission-container">
                <div className="error-message">
                    {error}
                </div>
                <button
                    className="btn btn-secondary"
                    onClick={() => navigate(fromCreateAudit ? (returnToUrl ? decodeURIComponent(returnToUrl) : '/create-audit-observations') : '/create-audit')}
                >
                    Go Back
                </button>
            </div>
        );
    }

    if (!formData) {
        return null;
    }

    const linkedOpeningId = selectedOpeningSubId || openingSubIdFromUrl;
    const linkedOpening = openingSubmissions.find((s) => String(s.submission_id) === String(linkedOpeningId));
    const isClosingWithFixedOpening = formData.form_type_id === 2 && linkedOpeningId;

    return (
        <div className="form-submission-container">
            <div className="form-submission-header">
                <div>
                    <h1 className="form-submission-title">{formData.title}</h1>
                    {isClosingWithFixedOpening && (
                        <p className="form-submission-opening-heading" style={{ margin: '0.35rem 0 0 0', fontSize: '0.9375rem', color: '#4d4d4f', fontWeight: 600 }}>
                            Closing for: {linkedOpening ? `${linkedOpening.form_title || 'Opening'} (#${linkedOpening.submission_id})` : `Observation #${linkedOpeningId}`}
                        </p>
                    )}
                    {formData.description && (
                        <p className="form-submission-description">{formData.description}</p>
                    )}
                </div>
                <button
                    className="btn btn-secondary"
                    onClick={() => navigate(fromCreateAudit ? (returnToUrl ? decodeURIComponent(returnToUrl) : '/create-audit-observations') : '/create-audit')}
                >
                    Cancel
                </button>
            </div>

            {error && (
                <div className="error-message">
                    {error}
                </div>
            )}
            <ErrorPopup
                show={!!error}
                message={error || ''}
                onClose={() => setError(null)}
            />

            <form onSubmit={handleSubmit}>
                {formData.form_type_id === 2 && openingSubmissions.length > 0 && !openingSubIdFromUrl && (
                    <div className="form-submission-opening-link" style={{ marginBottom: '1.5rem', padding: '1rem', background: '#e5e5e5', border: '1px solid #e5e5e5', borderRadius: '8px' }}>
                        <label htmlFor="opening-sub-id" style={{ display: 'block', fontWeight: 600, marginBottom: '0.5rem', color: '#4d4d4f' }}>
                            Link to opening observation (for every opening there should be a closing)
                        </label>
                        <select
                            id="opening-sub-id"
                            className="form-input"
                            value={selectedOpeningSubId}
                            onChange={(e) => setSelectedOpeningSubId(e.target.value)}
                            style={{ maxWidth: '320px' }}
                        >
                            <option value="">— Select opening observation (optional) —</option>
                            {openingSubmissions.map((op) => (
                                <option key={op.submission_id} value={op.submission_id}>
                                    #{op.submission_id} – {op.form_title || 'Opening'}
                                </option>
                            ))}
                        </select>
                    </div>
                )}
                {/* Fill from AI: show when form has at least one question */}
                {allQuestionIds.length > 0 && (
                    <div className="form-submission-ai-fill" style={{ marginBottom: '1.5rem', padding: '1rem', background: '#f0f9ff', border: '1px solid #bae6fd', borderRadius: '8px' }}>
                        <button
                            type="button"
                            className="btn btn-secondary"
                            onClick={handleFillFromAi}
                            disabled={aiFillLoading || !getFirstUploadedImage()}
                            title={!getFirstUploadedImage() ? 'Upload an image first (e.g. attach media on any question)' : 'Fill answers from the uploaded image using AI'}
                        >
                            {aiFillLoading ? 'Filling from AI…' : 'Fill from AI'}
                        </button>
                        {!getFirstUploadedImage() && (
                            <span style={{ marginLeft: '0.75rem', fontSize: '0.875rem', color: '#64748b' }}>
                                Upload an image on any question to enable.
                            </span>
                        )}
                    </div>
                )}
                {/* Questions */}
                {displayQuestions.length === 0 ? (
                    <div style={{ padding: '2rem', textAlign: 'center', color: '#4d4d4f' }}>
                        No questions for this form.
                    </div>
                ) : (
                    <div className="questions-section">
                        {displayQuestions.map(({ question, isChild }, index) => (
                                    <div
                                        key={question.question_id}
                                        className={`question-item ${isChild ? 'question-item-child' : ''}`}
                                    >
                                        <label className="question-label">
                                            {index + 1}. {getQuestionLabel(question)}
                                        </label>
                                        <div className="question-input-wrapper">
                                            {renderQuestionInput(question)}
                                        </div>
                                        
                                        {/* Media upload when question allows media or is Image Mark/Attach type */}
                                        {(!!question.allows_media || getQuestionTypeName(question.question_type_id) === 'Image Mark/Attach') && (
                                            <div className="media-upload-section">
                                                <label className="media-upload-label">
                                                    <input
                                                        type="file"
                                                        accept={getQuestionTypeName(question.question_type_id) === 'Image Mark/Attach' ? 'image/*' : 'image/*,video/*'}
                                                        multiple
                                                        onChange={(e) => handleMediaChange(question.question_id, e.target.files)}
                                                        style={{ display: 'none' }}
                                                    />
                                                    <span className="media-upload-button">
                                                        📎 Attach Media
                                                    </span>
                                                </label>
                                                {mediaFiles[question.question_id] && mediaFiles[question.question_id].length > 0 && (
                                                    <div className="media-preview">
                                                        {mediaFiles[question.question_id].map((file, idx) => {
                                                            const isImage = file.type && file.type.startsWith('image/');
                                                            const mediaDataUrl = answers[question.question_id]?.media?.[idx];
                                                            return (
                                                                <div key={idx} className="media-item media-item-with-thumb">
                                                                    {isImage && mediaDataUrl ? (
                                                                        <div className="media-thumb-wrap">
                                                                            <img src={mediaDataUrl} alt="" className="media-thumb" />
                                                                            <button
                                                                                type="button"
                                                                                className="media-mark-btn"
                                                                                onClick={() => openAnnotationModal(question.question_id, idx)}
                                                                            >
                                                                                ✏️ Mark on image
                                                                            </button>
                                                                        </div>
                                                                    ) : (
                                                                        <span className="media-file-name">{file.name}</span>
                                                                    )}
                                                                    <button
                                                                        type="button"
                                                                        className="btn-remove-small"
                                                                        onClick={() => removeMediaAtIndex(question.question_id, idx)}
                                                                        aria-label="Remove"
                                                                    >
                                                                        ×
                                                                    </button>
                                                                </div>
                                                            );
                                                        })}
                                                    </div>
                                                )}
                                            </div>
                                        )}
                                    </div>
                                ))}
                    </div>
                )}

                <div className="form-actions">
                    <button
                        type="button"
                        className="btn btn-secondary"
                        onClick={() => navigate(fromCreateAudit ? (returnToUrl ? decodeURIComponent(returnToUrl) : '/create-audit-observations') : '/create-audit')}
                    >
                        Cancel
                    </button>
                    <button
                        type="submit"
                        className="btn btn-primary"
                        disabled={loading}
                    >
                        {loading ? 'Submitting...' : 'Submit Observation'}
                    </button>
                </div>
            </form>

            <SuccessPopup
                show={showPopup}
                message={popupMessage}
                onClose={() => {
                    setShowPopup(false);
                    if (lastCreatedSubmissionId != null) {
                        navigate(`/view-submission/${lastCreatedSubmissionId}`);
                        setLastCreatedSubmissionId(null);
                    } else {
                        const target = fromCreateAudit
                            ? (returnToUrl ? decodeURIComponent(returnToUrl) : '/create-audit-observations')
                            : (returnToUrl ? decodeURIComponent(returnToUrl) : '/audit-status');
                        navigate(target);
                    }
                }}
                onOk={() => {
                    setShowPopup(false);
                    if (lastCreatedSubmissionId != null) {
                        navigate(`/view-submission/${lastCreatedSubmissionId}`);
                        setLastCreatedSubmissionId(null);
                    } else {
                        const target = fromCreateAudit
                            ? (returnToUrl ? decodeURIComponent(returnToUrl) : '/create-audit-observations')
                            : (returnToUrl ? decodeURIComponent(returnToUrl) : '/audit-status');
                        navigate(target);
                    }
                }}
            />

            <ImageAnnotationModal
                isOpen={annotationModal.open}
                onClose={() => setAnnotationModal({ open: false, questionId: null, fileIndex: null })}
                imageSrc={annotationModal.questionId != null && annotationModal.fileIndex != null
                    ? (answers[annotationModal.questionId]?.media?.[annotationModal.fileIndex] || null)
                    : null}
                onDone={applyAnnotatedImage}
                defaultColor={formData?.form_type_id === 2 ? '#16a34a' : '#e53935'}
            />
        </div>
    );
};

export default FormSubmission;
