diff --git a/src/pages/Prototype/index.tsx b/src/pages/Prototype/index.tsx
index 6421e2f..2d58e0a 100644
--- a/src/pages/Prototype/index.tsx
+++ b/src/pages/Prototype/index.tsx
@@ -1,9 +1,21 @@
-import { useEffect, useState } from 'react';
+import ReactFlow, {
+  removeElements,
+  addEdge,
+  MiniMap,
+  Controls,
+  ControlButton,
+  FlowElement,
+  Elements,
+  Connection,
+  Edge,
+  ArrowHeadType,
+} from 'react-flow-renderer';
+import { useEffect, useState, useRef } from 'react';
 import { useHistory, useParams } from 'react-router-dom';
-import { useLazyQuery, useReactiveVar } from '@apollo/client';
+import { useLazyQuery, useMutation, useReactiveVar } from '@apollo/client';
 import { Redirect } from 'react-router';
 import { roleVar } from '../../graphql/state';
-import { Empty, ArrowLeft, Edit, Check } from '../../assets';
+import { Empty, ArrowLeft, Edit, CheckCircle } from '../../assets';
 import {
   Box,
   Text,
@@ -11,22 +23,39 @@ import {
   Spinner,
   FrontendFeatureCard,
   BackendFeatureCard,
+  Alert,
 } from '../../components';
 import { Wrapper } from './styles';
 import {
+  AddPrototypeMutation,
+  AddPrototypeMutationVariables,
+  GetPrototypeByIdQuery,
+  GetPrototypeByIdQueryVariables,
   GetTemplateByIdQuery,
   GetTemplateByIdQueryVariables,
+  ProtoTypeOutput,
   TemplateOutput,
+  UpdatePrototypeMutation,
+  UpdatePrototypeMutationVariables,
 } from '../../graphql/types';
 import { GET_TEMPLATE_BY_ID } from '../../graphql/template.api';
-import { theme } from '../../themes';
+import {
+  ADD_PROTOTYPE,
+  GET_PROTOTYPE_BY_ID,
+  UPDATE_PROTOTYPE,
+} from '../../graphql/prototype.api';
 
 const Prototype = () => {
   const role = useReactiveVar(roleVar);
   const history = useHistory();
   const { id } = useParams<{ id: string }>();
   const [template, setTemplate] = useState<TemplateOutput>();
+  const [prototype, setPrototype] = useState<Array<ProtoTypeOutput>>();
+  const [elements, setElements] = useState<Elements>([]);
   const [editing, setEditing] = useState<boolean>(false);
+  const [error, setError] = useState<string>('');
+  const [success, setSuccess] = useState<boolean>(false);
+  const diagramParentRef = useRef<HTMLDivElement>(null);
 
   const [getTemplate, { loading: templateLoading }] = useLazyQuery<
     GetTemplateByIdQuery,
@@ -38,17 +67,157 @@ const Prototype = () => {
     fetchPolicy: 'network-only',
   });
 
+  const [getPrototype, { loading: prototypeLoading }] = useLazyQuery<
+    GetPrototypeByIdQuery,
+    GetPrototypeByIdQueryVariables
+  >(GET_PROTOTYPE_BY_ID, {
+    onCompleted({ getPrototypeById }) {
+      setPrototype(getPrototypeById.prototype);
+    },
+  });
+
+  const [addPrototype] = useMutation<
+    AddPrototypeMutation,
+    AddPrototypeMutationVariables
+  >(ADD_PROTOTYPE, {
+    onCompleted({ addPrototype: addPrototypeResult }) {
+      setPrototype(addPrototypeResult.prototype);
+      setSuccess(true);
+      setTimeout(() => setSuccess(false), 3000);
+    },
+    onError({ graphQLErrors }) {
+      setError(graphQLErrors[0]?.extensions?.info);
+      setTimeout(() => setError(''), 3000);
+    },
+  });
+
+  const [updatePrototype] = useMutation<
+    UpdatePrototypeMutation,
+    UpdatePrototypeMutationVariables
+  >(UPDATE_PROTOTYPE, {
+    onCompleted({ updatePrototype: updatePrototypeResult }) {
+      setPrototype(updatePrototypeResult.prototype);
+      setSuccess(true);
+      setTimeout(() => setSuccess(false), 3000);
+    },
+  });
+
   useEffect(() => {
     if (id) {
       getTemplate({ variables: { id } });
+      getPrototype({ variables: { id } });
     }
 
     // eslint-disable-next-line
   }, [id]);
 
+  useEffect(() => {
+    if (template && template.features) {
+      const initialElements = template.features.map((feature, index) => {
+        if (['frontend', 'fullstack'].includes(feature.featureType)) {
+          return {
+            id: feature.id,
+            type: 'default',
+            data: {
+              label: <FrontendFeatureCard feature={feature} />,
+            },
+            position: { x: index * 100, y: index * 200 },
+            style: {
+              width: 'auto',
+            },
+            connectable: role === 'developer' && editing,
+          } as FlowElement;
+        }
+        return {} as FlowElement;
+      });
+
+      if (initialElements) setElements(initialElements);
+    }
+
+    if (prototype) {
+      const initialElements: Array<Edge> = [];
+      prototype.forEach((link) => {
+        link.connections.forEach((connection) => {
+          initialElements.push({
+            id: `edge-${link.feature.id}`,
+            source: link.feature.id,
+            target: connection.to,
+            arrowHeadType: ArrowHeadType.ArrowClosed,
+            className: 'normal-edge',
+          });
+        });
+      });
+
+      if (initialElements) setElements((els) => [...els, ...initialElements]);
+    }
+
+    // eslint-disable-next-line
+  }, [template, prototype, editing]);
+
+  const onElementsRemove = (elementsToRemove: Elements<any>) =>
+    setElements((els) => removeElements(elementsToRemove, els));
+  const onConnect = (params: Edge<any> | Connection) =>
+    setElements((els) =>
+      addEdge({ ...params, arrowHeadType: ArrowHeadType.ArrowClosed }, els)
+    );
+
+  const handleEditPrototype = () => {
+    if (editing) {
+      const prototypeInput = elements
+        // @ts-ignore
+        .filter((element) => element.source || element.target)
+        .map((element) => {
+          if (
+            element.hasOwnProperty('source') ||
+            element.hasOwnProperty('target')
+          ) {
+            return {
+              // @ts-ignore
+              featureId: element.source,
+              connections: [
+                {
+                  // @ts-ignore
+                  to: element.target,
+                  releations: { back: false, forword: true },
+                },
+              ],
+            };
+          }
+          return {};
+        });
+      if (prototypeInput && prototypeInput.length > 0) {
+        if (prototype) {
+          updatePrototype({
+            variables: {
+              prototype: {
+                templateId: id,
+                // @ts-ignore
+                prototype: prototypeInput,
+              },
+            },
+          });
+        } else {
+          addPrototype({
+            variables: {
+              prototype: {
+                templateId: id,
+                // @ts-ignore
+                prototype: prototypeInput,
+              },
+            },
+          });
+        }
+      }
+
+      setEditing(false);
+    } else {
+      setEditing(true);
+    }
+  };
+
   return role === 'productOwner' || role === 'developer' ? (
     <>
-      {!templateLoading ? (
+      {!templateLoading && !prototypeLoading ? (
         <>
           {template ? (
             <Wrapper color={role}>
@@ -59,7 +228,7 @@ const Prototype = () => {
                   alignItems='center'
                   marginBottom='20px'
                 >
-                  <Box>
+                  <Box marginRight='50px'>
                     <Button
                       text='Back'
                       color={role || 'client'}
@@ -71,6 +240,13 @@ const Prototype = () => {
                       Prototype
                     </Text>
                   </Box>
+                  {success && (
+                    <Alert
+                      color='success'
+                      text='Prototype updated successfully'
+                    />
+                  )}
+                  {error && <Alert color='error' text={error} />}
                 </Box>
                 {template.features && (
                   <>
@@ -86,55 +262,45 @@ const Prototype = () => {
                       </Box>
                       <Box
                         display='grid'
-                        gridTemplateColumns='repeat(2, auto)'
-                        rowGap='100px'
-                        alignItems='stretch'
                         background='#F9FAFA'
                         boxShadow='1px 1px 10px rgba(50, 59, 105, 0.25)'
                         borderRadius='10px'
-                        padding='30px 120px'
-                        position='relative'
+                        width='100%'
+                        height='400px'
+                        ref={diagramParentRef}
                       >
-                        {template.features.map((feature, index) => {
-                          if (
-                            feature.featureType === 'frontend' ||
-                            feature.featureType === 'fullstack'
-                          ) {
-                            return (
-                              <FrontendFeatureCard
-                                feature={feature}
-                                key={feature.id}
-                                className={
-                                  index === 0 || index % 2 === 0
-                                    ? 'frontend-feature-even'
-                                    : 'frontend-feature-odd'
-                                }
-                              />
-                            );
-                          }
-                          return null;
-                        })}
                         <Box
-                          position='absolute'
-                          top='30px'
-                          right='30px'
-                          width='35px'
-                          height='35px'
-                          padding='10px'
-                          borderRadius='10px'
-                          background={
-                            !editing
-                              ? theme.colors.white.main
-                              : theme.colors[role].main
+                          width={
+                            diagramParentRef.current
+                              ? `${
+                                  getComputedStyle(diagramParentRef.current)
+                                    ?.width
+                                }}px`
+                              : '100%'
                           }
-                          boxShadow='1px 1px 5px rgba(50, 59, 105, 0.25)'
-                          display='flex'
-                          alignItems='center'
-                          justifyContent='center'
-                          cursor='pointer'
-                          onClick={() => setEditing(!editing)}
+                          height='auto'
                         >
-                          {!editing ? <Edit /> : <Check />}
+                          <ReactFlow
+                            elements={elements}
+                            onElementsRemove={onElementsRemove}
+                            onConnect={onConnect}
+                            deleteKeyCode={46}
+                            edgeTypes={{ arrowHeadType: 'arrow' }}
+                          >
+                            {role === 'developer' && (
+                              <>
+                                <MiniMap />
+                                <Controls
+                                  showInteractive={false}
+                                  showFitView={false}
+                                >
+                                  <ControlButton onClick={handleEditPrototype}>
+                                    {!editing ? <Edit /> : <CheckCircle />}
+                                  </ControlButton>
+                                </Controls>
+                              </>
+                            )}
+                          </ReactFlow>
                         </Box>
                       </Box>
                     </Box>
diff --git a/src/pages/Prototype/styles.ts b/src/pages/Prototype/styles.ts
index 9ccd12b..bf7ee3d 100644
--- a/src/pages/Prototype/styles.ts
+++ b/src/pages/Prototype/styles.ts
@@ -9,21 +9,4 @@ export const Wrapper = styled.div<WrapperProps>`
     fill: ${({ theme, color }) =>
       color ? theme.colors[color].main : theme.colors.client.main};
   }
-
-  .frontend-feature-odd {
-    justify-self: flex-end;
-  }
-
-  .frontend-feature-even {
-    justify-self: flex-start;
-  }
-
-  .frontend-feature-even,
-  .frontend-feature-odd {
-    &:hover {
-      border: 2px solid
-        ${({ theme, color }) =>
-          color ? theme.colors[color].main : theme.colors.client.main};
-    }
-  }
 `;