"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 { Label } from "@/components/ui/label";
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";

// Schema for form validation
const formSchema = z.object({
  startingCapital: z.coerce
    .number()
    .min(0, "Starting capital must be a non-negative number"),
  monthlySavings: z.coerce
    .number()
    .min(0, "Monthly savings must be a non-negative number"),
  currentAge: z.coerce.number().min(18, "Age must be at least 18"),
  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(50, "Life expectancy must be at least 50"),
});

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

interface CalculationResult {
  fireNumber: number | null;
  retirementAge: number | null;
  inflationAdjustedAllowance: number | null;
  retirementYears: number | null;
  error?: string;
  yearlyData?: Array<{
    age: number;
    year: number;
    balance: number;
    phase: "accumulation" | "retirement";
  }>;
}

export default function FireCalculatorForm() {
  const [result, setResult] = useState<CalculationResult | null>(null);
  const currentYear = 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: 2000,
      inflationRate: 2,
      lifeExpectancy: 84,
    },
  });

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

    const startingCapital = values.startingCapital;
    const monthlySavings = values.monthlySavings;
    const currentAge = values.currentAge;
    const annualGrowthRate = values.cagr / 100;
    const initialMonthlyAllowance = values.desiredMonthlyAllowance;
    const annualInflation = values.inflationRate / 100;
    const lifeExpectancy = values.lifeExpectancy;

    const monthlyGrowthRate = Math.pow(1 + annualGrowthRate, 1 / 12) - 1;
    const monthlyInflationRate = Math.pow(1 + annualInflation, 1 / 12) - 1;
    const maxIterations = 1000; // Safety limit for iterations

    // Binary search for the required retirement capital
    let low = initialMonthlyAllowance * 12; // Minimum: one year of expenses
    let high = initialMonthlyAllowance * 12 * 100; // Maximum: hundred years of expenses
    let requiredCapital = 0;
    let retirementAge = 0;
    let finalInflationAdjustedAllowance = 0;

    // First, find when retirement is possible with accumulation phase
    let canRetire = false;
    let currentCapital = startingCapital;
    let age = currentAge;
    let monthlyAllowance = initialMonthlyAllowance;
    let iterations = 0;

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

    // Add starting point
    yearlyData.push({
      age: currentAge,
      year: currentYear,
      balance: startingCapital,
      phase: "accumulation",
    });

    // Accumulation phase simulation
    while (age < lifeExpectancy && iterations < maxIterations) {
      // Simulate one year of saving and growth
      for (let month = 0; month < 12; month++) {
        currentCapital += monthlySavings;
        currentCapital *= 1 + monthlyGrowthRate;
        // Update allowance for inflation
        monthlyAllowance *= 1 + monthlyInflationRate;
      }
      age++;
      iterations++;

      // Record yearly data
      yearlyData.push({
        age: age,
        year: currentYear + (age - currentAge),
        balance: Math.round(currentCapital),
        phase: "accumulation",
      });

      // Check each possible retirement capital target through binary search
      const mid = (low + high) / 2;
      if (high - low < 1) {
        // Binary search converged
        requiredCapital = mid;
        break;
      }

      // Test if this retirement capital is sufficient
      let testCapital = mid;
      let testAge = age;
      let testAllowance = monthlyAllowance;
      let isSufficient = true;

      // Simulate retirement phase with this capital
      while (testAge < lifeExpectancy) {
        for (let month = 0; month < 12; month++) {
          // Withdraw inflation-adjusted allowance
          testCapital -= testAllowance;
          // Grow remaining capital
          testCapital *= 1 + monthlyGrowthRate;
          // Adjust allowance for inflation
          testAllowance *= 1 + monthlyInflationRate;
        }
        testAge++;

        // Check if we've depleted capital before life expectancy
        if (testCapital <= 0) {
          isSufficient = false;
          break;
        }
      }

      if (isSufficient) {
        high = mid; // This capital or less might be enough
        if (currentCapital >= mid) {
          // We can retire now with this capital
          canRetire = true;
          retirementAge = age;
          requiredCapital = mid;
          finalInflationAdjustedAllowance = monthlyAllowance;
          break;
        }
      } else {
        low = mid; // We need more capital
      }
    }

    // If we didn't find retirement possible in the loop
    if (!canRetire && iterations < maxIterations) {
      // Continue accumulation phase until we reach sufficient capital
      while (age < lifeExpectancy && iterations < maxIterations) {
        // Simulate one year
        for (let month = 0; month < 12; month++) {
          currentCapital += monthlySavings;
          currentCapital *= 1 + monthlyGrowthRate;
          monthlyAllowance *= 1 + monthlyInflationRate;
        }
        age++;
        iterations++;

        // Record yearly data
        yearlyData.push({
          age: age,
          year: currentYear + (age - currentAge),
          balance: Math.round(currentCapital),
          phase: "accumulation",
        });

        // Test with current capital
        let testCapital = currentCapital;
        let testAge = age;
        let testAllowance = monthlyAllowance;
        let isSufficient = true;

        // Simulate retirement with current capital
        while (testAge < lifeExpectancy) {
          for (let month = 0; month < 12; month++) {
            testCapital -= testAllowance;
            testCapital *= 1 + monthlyGrowthRate;
            testAllowance *= 1 + monthlyInflationRate;
          }
          testAge++;

          if (testCapital <= 0) {
            isSufficient = false;
            break;
          }
        }

        if (isSufficient) {
          canRetire = true;
          retirementAge = age;
          requiredCapital = currentCapital;
          finalInflationAdjustedAllowance = monthlyAllowance;
          break;
        }
      }
    }

    // If retirement is possible, simulate the retirement phase for the chart
    if (canRetire) {
      // Update the phase for all years after retirement
      yearlyData.forEach((data) => {
        if (data.age >= retirementAge) {
          data.phase = "retirement";
        }
      });

      // Continue simulation for retirement phase if needed
      let simulationCapital = currentCapital;
      let simulationAllowance = monthlyAllowance;
      let simulationAge = age;

      // If we haven't simulated up to life expectancy, continue
      while (simulationAge < lifeExpectancy) {
        for (let month = 0; month < 12; month++) {
          simulationCapital -= simulationAllowance;
          simulationCapital *= 1 + monthlyGrowthRate;
          simulationAllowance *= 1 + monthlyInflationRate;
        }
        simulationAge++;

        // Record yearly data
        yearlyData.push({
          age: simulationAge,
          year: currentYear + (simulationAge - currentAge),
          balance: Math.round(simulationCapital),
          phase: "retirement",
        });
      }
    }

    if (canRetire) {
      setResult({
        fireNumber: requiredCapital,
        retirementAge: retirementAge,
        inflationAdjustedAllowance: finalInflationAdjustedAllowance,
        retirementYears: lifeExpectancy - retirementAge,
        yearlyData: yearlyData,
        error: undefined,
      });
    } else {
      setResult({
        fireNumber: null,
        retirementAge: null,
        inflationAdjustedAllowance: null,
        retirementYears: null,
        yearlyData: yearlyData,
        error:
          iterations >= maxIterations
            ? "Calculation exceeded maximum iterations."
            : "Cannot reach FIRE goal before life expectancy with current parameters.",
      });
    }
  }

  // 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 (
    <div className="w-full max-w-3xl">
      <Card className="mb-8">
        <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="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="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>
                  )}
                />
                <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="lifeExpectancy"
                  render={({ field }) => (
                    <FormItem>
                      <FormLabel>Life Expectancy (Age)</FormLabel>
                      <FormControl>
                        <Input
                          placeholder="e.g., 90"
                          type="number"
                          {...field}
                        />
                      </FormControl>
                      <FormMessage />
                    </FormItem>
                  )}
                />
              </div>

              <Button type="submit" className="w-full">
                Calculate
              </Button>
            </form>
          </Form>
        </CardContent>
      </Card>

      {result && (
        <>
          <Card className="mb-8">
            <CardHeader>
              <CardTitle>Results</CardTitle>
            </CardHeader>
            <CardContent>
              {result.error ? (
                <p className="text-destructive">{result.error}</p>
              ) : (
                <div className="space-y-4">
                  <div>
                    <Label>FIRE Number (Required Capital)</Label>
                    <p className="text-2xl font-bold">
                      {formatNumber(result.fireNumber)}
                    </p>
                  </div>
                  <div>
                    <Label>Estimated Retirement Age</Label>
                    <p className="text-2xl font-bold">
                      {result.retirementAge ?? "N/A"}
                    </p>
                  </div>
                  {result.inflationAdjustedAllowance && (
                    <div>
                      <Label>
                        Monthly Allowance at Retirement (Inflation Adjusted)
                      </Label>
                      <p className="text-2xl font-bold">
                        {formatNumber(result.inflationAdjustedAllowance)}
                      </p>
                    </div>
                  )}
                  {result.retirementYears && (
                    <div>
                      <Label>Retirement Duration (Years)</Label>
                      <p className="text-2xl font-bold">
                        {result.retirementYears}
                      </p>
                    </div>
                  )}
                </div>
              )}
            </CardContent>
          </Card>

          {result.yearlyData && result.yearlyData.length > 0 && (
            <Card>
              <CardHeader>
                <CardTitle>Financial Projection</CardTitle>
                <CardDescription>
                  Projected balance growth and FIRE number threshold
                </CardDescription>
              </CardHeader>
              <CardContent>
                <ChartContainer
                  className="h-80"
                  config={{
                    balance: {
                      label: "Balance",
                      color: "var(--chart-1)",
                    },
                    fireNumber: {
                      label: "FIRE Number",
                      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).toFixed(1)}M`;
                        } else if (value >= 1000) {
                          return `${(value / 1000).toFixed(0)}K`;
                        }
                        return value.toString();
                      }}
                      width={80}
                    />
                    <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-primary">{`Balance: ${formatNumber(data.balance)}`}</p>
                              {result.fireNumber && (
                                <p className="text-destructive">{`FIRE Number: ${formatNumber(result.fireNumber)}`}</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>
                    </defs>
                    <Area
                      type="monotone"
                      dataKey="balance"
                      name="balance"
                      stroke="var(--chart-1)"
                      fill="url(#fillBalance)"
                      fillOpacity={0.4}
                      activeDot={{ r: 6 }}
                    />
                    {result.fireNumber && (
                      <ReferenceLine
                        y={result.fireNumber}
                        stroke="var(--chart-3)"
                        strokeWidth={2}
                        strokeDasharray="5 5"
                        label={{
                          value: "FIRE Number",
                          position: "insideBottomRight",
                        }}
                      />
                    )}
                    {result.retirementAge && (
                      <ReferenceLine
                        x={
                          currentYear +
                          (result.retirementAge - form.getValues().currentAge)
                        }
                        stroke="var(--chart-2)"
                        strokeWidth={2}
                        label={{
                          value: "Retirement",
                          position: "insideTopRight",
                        }}
                      />
                    )}
                  </AreaChart>
                </ChartContainer>
              </CardContent>
            </Card>
          )}
        </>
      )}
    </div>
  );
}