"use client";

import * as React from "react";
import { useState } from "react";
import { zodResolver } from "@hookform/resolvers/zod";
import { useForm } from "react-hook-form";
import * as z from "zod";

import { Button } from "@/components/ui/button";
import { Input } from "@/components/ui/input";
import {
  Form,
  FormControl,
  FormField,
  FormItem,
  FormLabel,
  FormMessage,
} from "@/components/ui/form";
import {
  Card,
  CardContent,
  CardDescription,
  CardHeader,
  CardTitle,
} from "@/components/ui/card";
import { ChartContainer, ChartTooltip } from "@/components/ui/chart";
import {
  Area,
  AreaChart,
  CartesianGrid,
  XAxis,
  YAxis,
  ReferenceLine,
} from "recharts";
import { Slider } from "@/components/ui/slider";
import assert from "assert";

// Schema for form validation
const formSchema = z.object({
  startingCapital: z.coerce.number(),
  monthlySavings: z.coerce
    .number()
    .min(0, "Monthly savings must be a non-negative number"),
  currentAge: z.coerce
    .number()
    .min(1, "Age must be at least 1")
    .max(100, "No point in starting this late"),
  cagr: z.coerce.number().min(0, "Growth rate must be a non-negative number"),
  desiredMonthlyAllowance: z.coerce
    .number()
    .min(0, "Monthly allowance must be a non-negative number"),
  inflationRate: z.coerce
    .number()
    .min(0, "Inflation rate must be a non-negative number"),
  lifeExpectancy: z.coerce
    .number()
    .min(40, "Be a bit more optimistic buddy :(")
    .max(100, "You should be more realistic..."),
  retirementAge: z.coerce
    .number()
    .min(18, "Retirement age must be at least 18")
    .max(100, "Retirement age must be at most 100"),
});

// Type for form values
type FormValues = z.infer<typeof formSchema>;

interface YearlyData {
  age: number;
  year: number;
  balance: number;
  phase: "accumulation" | "retirement";
  monthlyAllowance: number;
}

interface CalculationResult {
  fireNumber: number | null;
  yearlyData: YearlyData[];
  error?: string;
}

export default function FireCalculatorForm() {
  const [result, setResult] = useState<CalculationResult | null>(null);
  const irlYear = new Date().getFullYear();

  // Initialize form with default values
  const form = useForm<FormValues>({
    resolver: zodResolver(formSchema),
    defaultValues: {
      startingCapital: 50000,
      monthlySavings: 1500,
      currentAge: 25,
      cagr: 7,
      desiredMonthlyAllowance: 3000,
      inflationRate: 2,
      lifeExpectancy: 84,
      retirementAge: 55,
    },
  });

  function onSubmit(values: FormValues) {
    setResult(null); // Reset previous results

    const startingCapital = values.startingCapital;
    const monthlySavings = values.monthlySavings;
    const age = values.currentAge;
    const annualGrowthRate = 1 + values.cagr / 100;
    const initialMonthlyAllowance = values.desiredMonthlyAllowance;
    const annualInflation = 1 + values.inflationRate / 100;
    const ageOfDeath = values.lifeExpectancy;
    const retirementAge = values.retirementAge;

    // Array to store yearly data for the chart
    const yearlyData: YearlyData[] = [];

    // Initial year data
    yearlyData.push({
      age: age,
      year: irlYear,
      balance: startingCapital,
      phase: "accumulation",
      monthlyAllowance: initialMonthlyAllowance,
    });

    // Calculate accumulation phase (before retirement)
    for (let year = irlYear + 1; year <= irlYear + (ageOfDeath - age); year++) {
      const currentAge = age + (year - irlYear);
      const previousYearData = yearlyData[yearlyData.length - 1];
      const inflatedAllowance =
        initialMonthlyAllowance * Math.pow(annualInflation, year - irlYear);

      const isRetirementYear = currentAge >= retirementAge;
      const phase = isRetirementYear ? "retirement" : "accumulation";

      assert(!!previousYearData);
      // Calculate balance based on phase
      let newBalance;
      if (phase === "accumulation") {
        // During accumulation: grow previous balance + add savings
        newBalance =
          previousYearData.balance * annualGrowthRate + monthlySavings * 12;
      } else {
        // During retirement: grow previous balance - withdraw allowance
        newBalance =
          previousYearData.balance * annualGrowthRate - inflatedAllowance * 12;
      }

      yearlyData.push({
        age: currentAge,
        year: year,
        balance: newBalance,
        phase: phase,
        monthlyAllowance: inflatedAllowance,
      });
    }

    // Calculate FIRE number at retirement
    const retirementYear = irlYear + (retirementAge - age);
    const retirementIndex = yearlyData.findIndex(
      (data) => data.year === retirementYear,
    );
    const retirementData = yearlyData[retirementIndex];

    if (retirementIndex === -1 || !retirementData) {
      setResult({
        fireNumber: null,
        error: "Could not calculate retirement data",
        yearlyData: yearlyData,
      });
    } else {
      // Set the result
      setResult({
        fireNumber: retirementData.balance,
        yearlyData: yearlyData,
      });
    }
  }

  // Helper function to format currency without specific symbols
  const formatNumber = (value: number | null) => {
    if (value === null) return "N/A";
    return new Intl.NumberFormat("en", {
      maximumFractionDigits: 0,
    }).format(value);
  };

  return (
    <>
      <Card className="mb-4">
        <CardHeader>
          <CardTitle className="text-2xl">FIRE Calculator</CardTitle>
          <CardDescription>
            Calculate your path to financial independence and retirement
          </CardDescription>
        </CardHeader>
        <CardContent>
          <Form {...form}>
            <form onSubmit={form.handleSubmit(onSubmit)} className="space-y-8">
              <div className="grid grid-cols-1 gap-6 md:grid-cols-2">
                <FormField
                  control={form.control}
                  name="startingCapital"
                  render={({ field }) => (
                    <FormItem>
                      <FormLabel>Starting Capital</FormLabel>
                      <FormControl>
                        <Input
                          placeholder="e.g., 10000"
                          type="number"
                          {...field}
                        />
                      </FormControl>
                      <FormMessage />
                    </FormItem>
                  )}
                />
                <FormField
                  control={form.control}
                  name="monthlySavings"
                  render={({ field }) => (
                    <FormItem>
                      <FormLabel>Monthly Savings</FormLabel>
                      <FormControl>
                        <Input
                          placeholder="e.g., 500"
                          type="number"
                          {...field}
                        />
                      </FormControl>
                      <FormMessage />
                    </FormItem>
                  )}
                />
                <FormField
                  control={form.control}
                  name="currentAge"
                  render={({ field }) => (
                    <FormItem>
                      <FormLabel>Current Age</FormLabel>
                      <FormControl>
                        <Input
                          placeholder="e.g., 30"
                          type="number"
                          {...field}
                        />
                      </FormControl>
                      <FormMessage />
                    </FormItem>
                  )}
                />
                <FormField
                  control={form.control}
                  name="lifeExpectancy"
                  render={({ field }) => (
                    <FormItem>
                      <FormLabel>Life Expectancy (Age)</FormLabel>
                      <FormControl>
                        <Input
                          placeholder="e.g., 90"
                          type="number"
                          {...field}
                        />
                      </FormControl>
                      <FormMessage />
                    </FormItem>
                  )}
                />
                <FormField
                  control={form.control}
                  name="cagr"
                  render={({ field }) => (
                    <FormItem>
                      <FormLabel>Expected Annual Growth Rate (%)</FormLabel>
                      <FormControl>
                        <Input
                          placeholder="e.g., 7"
                          type="number"
                          step="0.1"
                          {...field}
                        />
                      </FormControl>
                      <FormMessage />
                    </FormItem>
                  )}
                />
                <FormField
                  control={form.control}
                  name="inflationRate"
                  render={({ field }) => (
                    <FormItem>
                      <FormLabel>Annual Inflation Rate (%)</FormLabel>
                      <FormControl>
                        <Input
                          placeholder="e.g., 2"
                          type="number"
                          step="0.1"
                          {...field}
                        />
                      </FormControl>
                      <FormMessage />
                    </FormItem>
                  )}
                />
                <FormField
                  control={form.control}
                  name="desiredMonthlyAllowance"
                  render={({ field }) => (
                    <FormItem>
                      <FormLabel>
                        Desired Monthly Allowance (Today&apos;s Value)
                      </FormLabel>
                      <FormControl>
                        <Input
                          placeholder="e.g., 2000"
                          type="number"
                          {...field}
                        />
                      </FormControl>
                      <FormMessage />
                    </FormItem>
                  )}
                />

                {/* Retirement Age Slider */}
                <FormField
                  control={form.control}
                  name="retirementAge"
                  render={({ field }) => (
                    <FormItem>
                      <FormLabel>Retirement Age: {field.value}</FormLabel>
                      <FormControl>
                        <Slider
                          value={[field.value]}
                          min={form.getValues().currentAge + 1}
                          max={form.getValues().lifeExpectancy - 1}
                          step={1}
                          onValueChange={field.onChange}
                          className="py-4"
                        />
                      </FormControl>
                      <FormMessage />
                    </FormItem>
                  )}
                />
                <FormItem className="col-span-full"></FormItem>
              </div>

              <Button type="submit" className="w-full">
                Calculate
              </Button>
              {result?.yearlyData && (
                <Card className="rounded-md shadow-none">
                  <CardHeader>
                    <CardTitle>Financial Projection</CardTitle>
                    <CardDescription>
                      Projected balance growth with your selected retirement age
                    </CardDescription>
                  </CardHeader>
                  <CardContent>
                    <ChartContainer
                      className="aspect-auto h-80 w-full"
                      config={{
                        balance: {
                          label: "Balance",
                          color: "var(--chart-1)",
                        },
                        realBalance: {
                          label: "Real Balance",
                          color: "var(--chart-3)",
                        },
                      }}
                    >
                      <AreaChart
                        data={result.yearlyData}
                        margin={{ top: 20, right: 30, left: 20, bottom: 20 }}
                      >
                        <CartesianGrid strokeDasharray="3 3" />
                        <XAxis
                          dataKey="year"
                          label={{
                            value: "Year",
                            position: "insideBottom",
                            offset: -10,
                          }}
                        />
                        <YAxis
                          tickFormatter={(value: number) => {
                            if (value >= 1000000) {
                              return `${(value / 1000000).toPrecision(3)}M`;
                            } else if (value >= 1000) {
                              return `${(value / 1000).toPrecision(3)}K`;
                            } else if (value <= -1000000) {
                              return `${(value / 1000000).toPrecision(3)}M`;
                            } else if (value <= -1000) {
                              return `${(value / 1000).toPrecision(3)}K`;
                            }
                            return value.toString();
                          }}
                          width={25}
                        />
                        <ChartTooltip
                          content={({ active, payload }) => {
                            if (active && payload?.[0]?.payload) {
                              const data = payload[0]
                                .payload as (typeof result.yearlyData)[0];
                              return (
                                <div className="bg-background border p-2 shadow-sm">
                                  <p className="font-medium">{`Year: ${data.year.toString()} (Age: ${data.age.toString()})`}</p>
                                  <p className="text-chart-1">{`Balance: ${formatNumber(data.balance)}`}</p>
                                  <p className="text-chart-2">{`Monthly allowance: ${formatNumber(data.monthlyAllowance)}`}</p>
                                  <p>{`Phase: ${data.phase === "accumulation" ? "Accumulation" : "Retirement"}`}</p>
                                </div>
                              );
                            }
                            return null;
                          }}
                        />
                        <defs>
                          <linearGradient
                            id="fillBalance"
                            x1="0"
                            y1="0"
                            x2="0"
                            y2="1"
                          >
                            <stop
                              offset="5%"
                              stopColor="var(--chart-1)"
                              stopOpacity={0.8}
                            />
                            <stop
                              offset="95%"
                              stopColor="var(--chart-1)"
                              stopOpacity={0.1}
                            />
                          </linearGradient>
                          <linearGradient
                            id="fillAllowance"
                            x1="0"
                            y1="0"
                            x2="0"
                            y2="1"
                          >
                            <stop
                              offset="5%"
                              stopColor="var(--chart-2)"
                              stopOpacity={0.8}
                            />
                            <stop
                              offset="95%"
                              stopColor="var(--chart-2)"
                              stopOpacity={0.1}
                            />
                          </linearGradient>
                        </defs>
                        <Area
                          type="monotone"
                          dataKey="balance"
                          name="balance"
                          stroke="var(--chart-1)"
                          fill="url(#fillBalance)"
                          fillOpacity={0.4}
                          activeDot={{ r: 6 }}
                        />
                        <Area
                          type="monotone"
                          dataKey="monthlyAllowance"
                          name="allowance"
                          stroke="var(--chart-2)"
                          fill="url(#fillAllowance)"
                          fillOpacity={0.4}
                          activeDot={{ r: 6 }}
                        />
                        {result.fireNumber && (
                          <ReferenceLine
                            y={result.fireNumber}
                            stroke="var(--chart-3)"
                            strokeWidth={1}
                            strokeDasharray="2 2"
                            label={{
                              value: "FIRE Number",
                              position: "insideBottomRight",
                            }}
                          />
                        )}
                        <ReferenceLine
                          x={
                            irlYear +
                            (form.getValues().retirementAge -
                              form.getValues().currentAge)
                          }
                          stroke="var(--chart-2)"
                          strokeWidth={2}
                          label={{
                            value: "Retirement",
                            position: "insideTopRight",
                          }}
                        />
                      </AreaChart>
                    </ChartContainer>
                  </CardContent>
                </Card>
              )}
            </form>
          </Form>
        </CardContent>
      </Card>

      {result && (
        <div className="mb-4 grid grid-cols-1 gap-2 md:grid-cols-2">
          {result.error ? (
            <Card className="col-span-full">
              <CardContent className="pt-6">
                <p className="text-destructive">{result.error}</p>
              </CardContent>
            </Card>
          ) : (
            <>
              <Card>
                <CardHeader>
                  <CardTitle>FIRE Number</CardTitle>
                  <CardDescription className="text-xs">
                    Capital at retirement
                  </CardDescription>
                </CardHeader>
                <CardContent>
                  <p className="text-3xl font-bold">
                    {formatNumber(result.fireNumber)}
                  </p>
                </CardContent>
              </Card>

              <Card>
                <CardHeader>
                  <CardTitle>Retirement Duration</CardTitle>
                  <CardDescription className="text-xs">
                    Years to enjoy your financial independence
                  </CardDescription>
                </CardHeader>
                <CardContent>
                  <p className="text-3xl font-bold">
                    {form.getValues().lifeExpectancy -
                      form.getValues().retirementAge}
                  </p>
                </CardContent>
              </Card>
            </>
          )}
        </div>
      )}
    </>
  );
}