From 3d24eb797200fd35073cbc6a0de6b5c1e6687aef Mon Sep 17 00:00:00 2001
From: yu-zhen <wlphg010688@gmail.com>
Date: Wed, 24 Jul 2024 22:36:25 +0900
Subject: [PATCH] fix: multiple ballot issues

---
 src/components/BallotOverview.tsx             | 34 +++++++++++++++----
 src/contexts/Ballot.tsx                       |  7 ++--
 src/features/ballot/types/index.ts            |  1 +
 src/features/projects/components/Projects.tsx |  4 +--
 src/layouts/DefaultLayout.tsx                 | 13 +++----
 src/pages/ballot/index.tsx                    | 17 ++++++++--
 6 files changed, 55 insertions(+), 21 deletions(-)

diff --git a/src/components/BallotOverview.tsx b/src/components/BallotOverview.tsx
index 31386a20..4ccb95cc 100644
--- a/src/components/BallotOverview.tsx
+++ b/src/components/BallotOverview.tsx
@@ -1,12 +1,32 @@
+import Link from "next/link";
+
+import { useBallot } from "~/contexts/Ballot";
+import { useAppState } from "~/utils/state";
+import { EAppState } from "~/utils/types";
+
 import { AddedProjects } from "./AddedProjects";
 import { VotingUsage } from "./VotingUsage";
 
-export const BallotOverview = (): JSX.Element => (
-  <div className="my-8 flex-col items-center gap-2 rounded-lg bg-white p-5 uppercase shadow-lg dark:bg-lightBlack dark:text-white">
-    <h3>My Ballot</h3>
+export const BallotOverview = (): JSX.Element => {
+  const { ballot } = useBallot();
+
+  const appState = useAppState();
+
+  return (
+    <Link
+      href={
+        ballot.published && (appState === EAppState.TALLYING || appState === EAppState.RESULTS)
+          ? "/ballot/confirmation"
+          : "/ballot"
+      }
+    >
+      <div className="my-8 flex-col items-center gap-2 rounded-lg bg-white p-5 uppercase shadow-lg dark:bg-lightBlack dark:text-white">
+        <h3>My Ballot</h3>
 
-    <AddedProjects />
+        <AddedProjects />
 
-    <VotingUsage />
-  </div>
-);
+        <VotingUsage />
+      </div>
+    </Link>
+  );
+};
diff --git a/src/contexts/Ballot.tsx b/src/contexts/Ballot.tsx
index 0ea08aa7..b629e299 100644
--- a/src/contexts/Ballot.tsx
+++ b/src/contexts/Ballot.tsx
@@ -8,7 +8,7 @@ import type { Ballot, Vote } from "~/features/ballot/types";
 
 export const BallotContext = createContext<BallotContextType | undefined>(undefined);
 
-const defaultBallot = { votes: [], published: false };
+const defaultBallot = { votes: [], published: false, edited: false };
 
 export const BallotProvider: React.FC<BallotProviderProps> = ({ children }: BallotProviderProps) => {
   const [ballot, setBallot] = useState<Ballot>(defaultBallot);
@@ -38,6 +38,7 @@ export const BallotProvider: React.FC<BallotProviderProps> = ({ children }: Ball
     (addedVotes: Vote[], pollId: string) => ({
       ...ballot,
       pollId,
+      edited: true,
       votes: Object.values<Vote>({
         ...toObject("projectId", ballot.votes),
         ...toObject("projectId", addedVotes),
@@ -56,7 +57,7 @@ export const BallotProvider: React.FC<BallotProviderProps> = ({ children }: Ball
     (projectId: string) => {
       const votes = ballot.votes.filter((v) => v.projectId !== projectId);
 
-      setBallot({ ...ballot, votes, published: false });
+      setBallot({ ...ballot, votes });
     },
     [ballot, setBallot],
   );
@@ -81,7 +82,7 @@ export const BallotProvider: React.FC<BallotProviderProps> = ({ children }: Ball
 
   // set published to true
   const publishBallot = useCallback(() => {
-    setBallot({ ...ballot, published: true });
+    setBallot({ ...ballot, published: true, edited: false });
   }, [ballot, setBallot]);
 
   /// Read existing ballot in localStorage
diff --git a/src/features/ballot/types/index.ts b/src/features/ballot/types/index.ts
index 364dad69..0d664c91 100644
--- a/src/features/ballot/types/index.ts
+++ b/src/features/ballot/types/index.ts
@@ -8,6 +8,7 @@ export const VoteSchema = z.object({
 export const BallotSchema = z.object({
   votes: z.array(VoteSchema),
   published: z.boolean().default(false),
+  edited: z.boolean().default(false),
 });
 
 export type Vote = z.infer<typeof VoteSchema>;
diff --git a/src/features/projects/components/Projects.tsx b/src/features/projects/components/Projects.tsx
index 916b4dde..1ad974cd 100644
--- a/src/features/projects/components/Projects.tsx
+++ b/src/features/projects/components/Projects.tsx
@@ -54,10 +54,10 @@ export const Projects = (): JSX.Element => {
     if (!isRegistered) {
       return EProjectState.UNREGISTERED;
     }
-    if (ballotContains(projectId) && ballot.published) {
+    if (ballotContains(projectId) && ballot.published && !ballot.edited) {
       return EProjectState.SUBMITTED;
     }
-    if (ballotContains(projectId) && !ballot.published) {
+    if (ballotContains(projectId)) {
       return EProjectState.ADDED;
     }
     return EProjectState.DEFAULT;
diff --git a/src/layouts/DefaultLayout.tsx b/src/layouts/DefaultLayout.tsx
index 5a24e4f8..66503a47 100644
--- a/src/layouts/DefaultLayout.tsx
+++ b/src/layouts/DefaultLayout.tsx
@@ -35,17 +35,17 @@ export const Layout = ({ children = null, ...props }: ILayoutProps): JSX.Element
       },
     ];
 
-    if (ballot.published) {
+    if (appState === EAppState.VOTING && isRegistered) {
       links.push({
-        href: "/ballot/confirmation",
+        href: "/ballot",
         children: "My Ballot",
       });
     }
 
-    if (appState === EAppState.VOTING && !ballot.published && isRegistered) {
+    if ((appState === EAppState.TALLYING || appState === EAppState.RESULTS) && ballot.published) {
       links.push({
-        href: "/ballot",
-        children: "My Ballot",
+        href: "/ballot/confirmation",
+        children: "Submitted Ballot",
       });
     }
 
@@ -85,6 +85,7 @@ export const LayoutWithSidebar = ({ ...props }: ILayoutProps): JSX.Element => {
   const { isRegistered } = useMaci();
   const { address } = useAccount();
   const { ballot } = useBallot();
+  const appState = useAppState();
 
   const { showInfo, showBallot, showSubmitButton } = props;
 
@@ -95,7 +96,7 @@ export const LayoutWithSidebar = ({ ...props }: ILayoutProps): JSX.Element => {
         <div>
           {showInfo && <Info showVotingInfo size="sm" />}
 
-          {showBallot && address && isRegistered && <BallotOverview />}
+          {appState !== EAppState.APPLICATION && showBallot && address && isRegistered && <BallotOverview />}
 
           {showSubmitButton && ballot.votes.length > 0 && (
             <div className="flex flex-col gap-4">
diff --git a/src/pages/ballot/index.tsx b/src/pages/ballot/index.tsx
index 5dc0cffa..070d14d2 100644
--- a/src/pages/ballot/index.tsx
+++ b/src/pages/ballot/index.tsx
@@ -91,6 +91,12 @@ const BallotAllocationForm = (): JSX.Element => {
 
       <p className="my-4 text-gray-400">Once you have reviewed your vote allocation, you can submit your ballot.</p>
 
+      {ballot.published && (
+        <Link className="text-blue-400 hover:underline" href="/ballot/confirmation">
+          Check your submitted ballot
+        </Link>
+      )}
+
       <div className="mb-4 justify-end sm:flex">{ballot.votes.length ? <ClearBallot /> : null}</div>
 
       <div className="border-t border-gray-300">
@@ -116,6 +122,7 @@ const BallotPage = (): JSX.Element => {
   const { address, isConnecting } = useAccount();
   const { ballot, sumBallot } = useBallot();
   const router = useRouter();
+  const appState = useAppState();
 
   useEffect(() => {
     if (!address && !isConnecting) {
@@ -129,9 +136,13 @@ const BallotPage = (): JSX.Element => {
 
   return (
     <LayoutWithSidebar requireAuth requireRegistration showBallot showSubmitButton sidebar="right">
-      <Form defaultValues={ballot} schema={BallotSchema} values={ballot} onSubmit={handleSubmit}>
-        <BallotAllocationForm />
-      </Form>
+      {appState === EAppState.VOTING && (
+        <Form defaultValues={ballot} schema={BallotSchema} values={ballot} onSubmit={handleSubmit}>
+          <BallotAllocationForm />
+        </Form>
+      )}
+
+      {appState !== EAppState.VOTING && <div>You can only vote during the voting period.</div>}
     </LayoutWithSidebar>
   );
 };