import "./WorkflowPipeline.css";

const STAGES = [
    { key: "EXTRACTION", label: "Extract" },
    { key: "CHUNKING", label: "Chunk" },
    { key: "EMBEDDING", label: "Embed" },
    { key: "CLASSIFICATION", label: "Classify" },
    { key: "KNOWLEDGE_EXTRACTION", label: "Knowledge" },
    { key: "RECONCILIATION", label: "Reconcile" },
    { key: "VALIDATION", label: "Validate" },
    { key: "DECISION", label: "Decide" },
    { key: "HUMAN_REVIEW", label: "Review" },
    { key: "COMPLETED", label: "Complete" },
];

// Map stage keys to metric stage names
const METRIC_MAP = {
    EXTRACTION: "extraction",
    CHUNKING: "chunking",
    EMBEDDING: "embedding",
    KNOWLEDGE_EXTRACTION: "knowledge_extraction",
    RECONCILIATION: "reconciliation",
    VALIDATION: "validation",
};

function getStageStatus(stage, currentNode, workflowStatus) {
    if (workflowStatus === "COMPLETED") return "done";
    if (workflowStatus === "FAILED") {
        const currentIdx = STAGES.findIndex((s) => s.key === currentNode);
        const stageIdx = STAGES.findIndex((s) => s.key === stage.key);
        if (stageIdx < currentIdx) return "done";
        if (stageIdx === currentIdx) return "failed";
        return "pending";
    }
    const currentIdx = STAGES.findIndex((s) => s.key === currentNode);
    const stageIdx = STAGES.findIndex((s) => s.key === stage.key);
    if (currentIdx < 0) return "pending";
    if (stageIdx < currentIdx) return "done";
    if (stageIdx === currentIdx) {
        if (workflowStatus === "WAITING_FOR_REVIEW" && stage.key === "HUMAN_REVIEW") return "active";
        if (workflowStatus === "RUNNING") return "active";
        return "done";
    }
    return "pending";
}

function formatMs(ms) {
    if (!ms || ms <= 0) return null;
    if (ms >= 1000) return `${(ms / 1000).toFixed(1)}s`;
    return `${Math.round(ms)}ms`;
}

/**
 * @param {object} props
 * @param {string} props.currentNode
 * @param {string} props.workflowStatus
 * @param {Array} [props.stageMetrics] - Array of {name, elapsed_ms, input_tokens, output_tokens}
 */
export function WorkflowPipeline({ currentNode, workflowStatus, stageMetrics }) {
    // Build metrics lookup
    const metricsMap = {};
    if (stageMetrics) {
        for (const m of stageMetrics) {
            metricsMap[m.name] = m;
        }
    }

    return (
        <div className="dw-pipeline">
            {STAGES.map((stage, i) => {
                const status = getStageStatus(stage, currentNode, workflowStatus);
                const metricName = METRIC_MAP[stage.key];
                const metric = metricName ? metricsMap[metricName] : null;
                const timeStr = metric ? formatMs(metric.elapsed_ms) : null;
                const tokens = metric ? (metric.input_tokens || 0) + (metric.output_tokens || 0) : 0;

                return (
                    <div key={stage.key} className="dw-pipeline__stage-wrapper">
                        <div className={`dw-pipeline__stage dw-pipeline__stage--${status}`}>
                            <span className="dw-pipeline__icon">
                                {status === "done" && "✓"}
                                {status === "active" && "●"}
                                {status === "failed" && "✗"}
                                {status === "pending" && "○"}
                            </span>
                            <span className="dw-pipeline__label">{stage.label}</span>
                            {timeStr && status === "done" && (
                                <span className="dw-pipeline__metric">{timeStr}</span>
                            )}
                            {tokens > 0 && status === "done" && (
                                <span className="dw-pipeline__metric dw-pipeline__metric--tokens">{tokens.toLocaleString()} tok</span>
                            )}
                        </div>
                        {i < STAGES.length - 1 && (
                            <div className={`dw-pipeline__connector dw-pipeline__connector--${status === "done" ? "done" : "pending"}`} />
                        )}
                    </div>
                );
            })}
        </div>
    );
}