src/components/tasks/TaskMatrixView.tsx

import React, { useState, useMemo } from 'react';
import { format, startOfMonth, endOfMonth, eachDayOfInterval, isSameDay, parseISO, isWithinInterval, startOfDay, endOfDay } from 'date-fns';
import { Task } from '@/services/tasksApi';
import { Employee } from '@/types';
import {
    Table,
    TableBody,
    TableCell,
    TableHead,
    TableHeader,
    TableRow,
} from "@/components/ui/table";
import { Card, CardContent, CardHeader, CardTitle } from "@/components/ui/card";
import { Button } from "@/components/ui/button";
import { ChevronLeft, ChevronRight, X, Calendar as CalendarIcon, User, ListFilter, Clock } from "lucide-react";
import { Badge } from "@/components/ui/badge";
import { useNavigate } from "react-router-dom";
import { cn } from "@/lib/utils";
import { CustomSheet } from "@/components/ui/custom-sheet";

interface TaskMatrixViewProps {
    tasks: Task[];
    employees: Employee[];
}

interface SelectionState {
    type: 'cell' | 'employee' | 'day';
    id: string; // Unique identifier for highlighting
    title: string;
    subTitle?: string;
    tasks: Task[];
}

export const TaskMatrixView: React.FC<TaskMatrixViewProps> = ({ tasks, employees }) => {
    const navigate = useNavigate();
    const [currentDate, setCurrentDate] = useState(new Date());
    const [selection, setSelection] = useState<SelectionState | null>(null);
    const [selectedTask, setSelectedTask] = useState<Task | null>(null);
    const [isSheetOpen, setIsSheetOpen] = useState(false);

    // Get all days in the current month
    const daysInMonth = useMemo(() => {
        const start = startOfMonth(currentDate);
        const end = endOfMonth(currentDate);
        return eachDayOfInterval({ start, end });
    }, [currentDate]);

    // Navigate months
    const nextMonth = () => {
        setCurrentDate(new Date(currentDate.getFullYear(), currentDate.getMonth() + 1, 1));
        setSelection(null);
    };

    const prevMonth = () => {
        setCurrentDate(new Date(currentDate.getFullYear(), currentDate.getMonth() - 1, 1));
        setSelection(null);
    };

    // Calculate matrix data (counts and tasks per cell)
    const matrixData = useMemo(() => {
        const data: Record<number, Record<string, Task[]>> = {};
        const dailyTotals: Record<string, number> = {};
        const dailyTasks: Record<string, Task[]> = {}; // Store all tasks for a day
        const employeeTasks: Record<number, Task[]> = {}; // Store all monthly tasks for an employee

        // Initialize daily totals
        daysInMonth.forEach(day => {
            const dayStr = format(day, 'yyyy-MM-dd');
            dailyTotals[dayStr] = 0;
            dailyTasks[dayStr] = [];
        });

        // Initialize employee rows
        employees.forEach(emp => {
            data[emp.id] = {};
            employeeTasks[emp.id] = [];
            daysInMonth.forEach(day => {
                const dayStr = format(day, 'yyyy-MM-dd');
                data[emp.id][dayStr] = [];
            });
        });

        // Loop through each task
        tasks.forEach(task => {
            if (!task.assignedTo) return;

            // Skip tasks for employees not in the current list
            if (!data[task.assignedTo]) return;

            const startDate = task.stateDate ? parseISO(task.stateDate) : (task.createdAt ? parseISO(task.createdAt) : new Date());
            let endDate = task.endDate ? parseISO(task.endDate) : startDate;

            if (endDate < startDate) endDate = startDate;

            const monthStart = startOfMonth(currentDate);
            const monthEnd = endOfMonth(currentDate);

            if (endDate < monthStart || startDate > monthEnd) return;

            // Track if we've added this task to the employee's monthly list already
            let addedToEmployee = false;

            daysInMonth.forEach(day => {
                const dayStr = format(day, 'yyyy-MM-dd');
                const currentDayStart = startOfDay(day);
                const rangeStart = startOfDay(startDate);
                const rangeEnd = endOfDay(endDate);

                if (currentDayStart >= rangeStart && currentDayStart <= rangeEnd) {
                    if (data[task.assignedTo] && data[task.assignedTo][dayStr] !== undefined) {
                        data[task.assignedTo][dayStr].push(task);
                        dailyTotals[dayStr] = (dailyTotals[dayStr] || 0) + 1;

                        // Add to day aggregate (ensure uniqueness if needed, but tasks are unique per iteration)
                        dailyTasks[dayStr].push(task);

                        // Add to employee aggregate
                        if (!addedToEmployee) {
                            employeeTasks[task.assignedTo].push(task);
                            addedToEmployee = true;
                        }
                    }
                }
            });
        });

        return { employeeData: data, dailyTotals, dailyTasks, employeeTasks };
    }, [tasks, employees, daysInMonth, currentDate]);

    // Handle selection interactions
    const handleCellClick = (employee: Employee, day: Date, tasks: Task[]) => {
        if (tasks.length > 0) {
            setSelection({
                type: 'cell',
                id: `${employee.id}-${format(day, 'yyyy-MM-dd')}`,
                title: `Tasks for ${employee.firstName} ${employee.lastName}`,
                subTitle: format(day, "EEEE, MMMM do, yyyy"),
                tasks: tasks
            });
        }
    };

    const handleEmployeeHeaderClick = (employee: Employee) => {
        const tasks = matrixData.employeeTasks[employee.id] || [];
        if (tasks.length > 0) {
            setSelection({
                type: 'employee',
                id: employee.id.toString(),
                title: `Tasks for ${employee.firstName} ${employee.lastName}`,
                subTitle: `Total for ${format(currentDate, 'MMMM yyyy')}`,
                tasks: tasks
            });
        }
    };

    const handleDayHeaderClick = (day: Date) => {
        const dayStr = format(day, 'yyyy-MM-dd');
        const tasks = matrixData.dailyTasks[dayStr] || [];
        if (tasks.length > 0) {
            setSelection({
                type: 'day',
                id: dayStr,
                title: `All Tasks for ${format(day, "MMMM do")}`,
                subTitle: format(day, "EEEE, yyyy"),
                tasks: tasks
            });
        }
    };

    const handleEditTask = (task: Task) => {
        setSelectedTask(task);
        setIsSheetOpen(true);
    };

    return (
        <div className="space-y-6">
            <Card className="w-full shadow-sm border-none">
                <CardHeader className="pb-4">
                    <div className="flex items-center justify-between">
                        <CardTitle className="text-xl font-bold">Monthly employee and day wise task report</CardTitle>
                        <div className="flex items-center space-x-2">
                            <Button variant="outline" size="icon" onClick={prevMonth}>
                                <ChevronLeft className="h-4 w-4" />
                            </Button>
                            <span className="min-w-[150px] text-center font-medium">
                                {format(currentDate, 'MMMM yyyy')}
                            </span>
                            <Button variant="outline" size="icon" onClick={nextMonth}>
                                <ChevronRight className="h-4 w-4" />
                            </Button>
                        </div>
                    </div>
                </CardHeader>
                <CardContent className="p-0">
                    {/* Key change: Added max-height and scrolling to the container */}
                    <div className="max-h-[600px] overflow-auto relative border rounded-md m-4">
                        <table className="w-full caption-bottom text-sm text-left">
                            <TableHeader className="bg-muted z-30 sticky top-0 shadow-sm">
                                <TableRow>
                                    <TableHead className="w-[200px] sticky left-0 top-0 z-50 bg-muted border-r shadow-[1px_0_0_0_rgba(0,0,0,0.1)]">
                                        Employee Name
                                    </TableHead>
                                    {daysInMonth.map((day) => {
                                        const dayStr = format(day, 'yyyy-MM-dd');
                                        const isSelected = selection?.type === 'day' && selection.id === dayStr;

                                        return (
                                            <TableHead
                                                key={day.toString()}
                                                className={cn(
                                                    "text-center min-w-[40px] px-1 border-r border-b font-bold cursor-pointer hover:bg-muted/80 transition-colors sticky top-0 z-30 bg-muted",
                                                    isSelected ? "bg-primary/20 text-primary" : "text-foreground"
                                                )}
                                                onClick={() => handleDayHeaderClick(day)}
                                            >
                                                <div className="flex flex-col items-center justify-center py-1">
                                                    <span className="text-[10px] font-normal uppercase text-muted-foreground">{format(day, 'EEE')}</span>
                                                    <span>{format(day, 'd')}</span>
                                                </div>
                                            </TableHead>
                                        );
                                    })}
                                </TableRow>
                            </TableHeader>
                            <TableBody>
                                {employees.map((employee) => {
                                    const isRowSelected = selection?.type === 'employee' && selection.id === employee.id.toString();

                                    return (
                                        <TableRow key={employee.id} className={isRowSelected ? "bg-primary/5" : ""}>
                                            <TableCell
                                                className={cn(
                                                    "font-medium sticky left-0 z-20 bg-background border-r border-b shadow-[1px_0_0_0_rgba(0,0,0,0.1)] cursor-pointer hover:bg-muted/50 transition-colors",
                                                    isRowSelected ? "bg-primary/10 text-primary" : ""
                                                )}
                                                onClick={() => handleEmployeeHeaderClick(employee)}
                                            >
                                                <div className="truncate w-[180px]" title={`${employee.firstName} ${employee.lastName}`}>
                                                    {employee.firstName} {employee.lastName}
                                                </div>
                                            </TableCell>

                                            {daysInMonth.map((day) => {
                                                const dayStr = format(day, 'yyyy-MM-dd');
                                                const cellTasks = matrixData.employeeData[employee.id]?.[dayStr] || [];
                                                const count = cellTasks.length;
                                                const cellId = `${employee.id}-${dayStr}`;

                                                // Determine highlighting based on selection type
                                                const isCellSelected = selection?.type === 'cell' && selection.id === cellId;
                                                const isColSelected = selection?.type === 'day' && selection.id === dayStr;

                                                // Base background color logic
                                                let bgClass = "";
                                                if (isCellSelected) bgClass = "bg-primary/20 ring-1 ring-primary ring-inset z-10";
                                                else if (isRowSelected || isColSelected) bgClass = "bg-primary/5";

                                                return (
                                                    <TableCell
                                                        key={dayStr}
                                                        className={cn(
                                                            "text-center border-r border-b p-0 relative transition-colors h-10 w-10",
                                                            count > 0 ? "cursor-pointer hover:bg-primary/10" : "",
                                                            bgClass
                                                        )}
                                                        onClick={() => count > 0 && handleCellClick(employee, day, cellTasks)}
                                                    >
                                                        <div className="h-full w-full flex items-center justify-center">
                                                            {count > 0 ? (
                                                                <span className={cn(
                                                                    "font-medium",
                                                                    isCellSelected ? "text-primary font-bold" : ""
                                                                )}>
                                                                    {count}
                                                                </span>
                                                            ) : (
                                                                <span className="text-muted-foreground/30 text-xs">-</span>
                                                            )}
                                                        </div>
                                                    </TableCell>
                                                );
                                            })}
                                        </TableRow>
                                    );
                                })}
                                {/* Totals Row */}
                                <TableRow className="font-bold bg-muted/30 sticky bottom-0 z-30 shadow-[0_-1px_0_0_rgba(0,0,0,0.1)]">
                                    <TableCell
                                        className="sticky left-0 z-40 bg-muted border-r border-t shadow-[1px_0_0_0_rgba(0,0,0,0.1)]"
                                    >
                                        Total Task
                                    </TableCell>
                                    {daysInMonth.map((day) => {
                                        const dayStr = format(day, 'yyyy-MM-dd');
                                        return (
                                            <TableCell key={dayStr} className="text-center border-t border-r bg-muted/30 p-2 text-xs">
                                                {matrixData.dailyTotals[dayStr] || 0}
                                            </TableCell>
                                        );
                                    })}
                                </TableRow>
                            </TableBody>
                        </table>
                    </div>
                </CardContent>
            </Card>

            {/* Selected Task List */}
            {selection && (
                <Card className="animate-in slide-in-from-top-4 duration-300 border-l-4 border-l-primary shadow-md">
                    <CardHeader className="pb-2 bg-muted/20">
                        <div className="flex justify-between items-center">
                            <div>
                                <CardTitle className="text-lg flex items-center gap-2">
                                    {selection.type === 'employee' && <User className="h-5 w-5 text-primary" />}
                                    {selection.type === 'day' && <CalendarIcon className="h-5 w-5 text-primary" />}
                                    {selection.type === 'cell' && <ListFilter className="h-5 w-5 text-primary" />}
                                    {selection.title}
                                </CardTitle>
                                {selection.subTitle && (
                                    <p className="text-sm text-muted-foreground mt-1 flex items-center">
                                        {selection.subTitle}
                                    </p>
                                )}
                            </div>
                            <Button variant="ghost" size="icon" onClick={() => setSelection(null)}>
                                <X className="h-4 w-4" />
                            </Button>
                        </div>
                    </CardHeader>
                    <CardContent className="pt-4 px-0">
                        <div className="overflow-x-auto">
                            <Table>
                                <TableHeader>
                                    <TableRow className="hover:bg-transparent">
                                        <TableHead className="pl-6">Title</TableHead>
                                        <TableHead>Status</TableHead>
                                        <TableHead>Start Date</TableHead>
                                        <TableHead className="pr-6">End Date</TableHead>
                                    </TableRow>
                                </TableHeader>
                                <TableBody>
                                    {selection.tasks.map((task) => (
                                        <TableRow
                                            key={`${task.id}-${task.assignedTo}`}
                                            className="cursor-pointer hover:bg-muted/50"
                                            onClick={() => handleEditTask(task)}
                                        >
                                            <TableCell className="font-medium pl-6 max-w-[300px] truncate" title={task.title}>{task.title}</TableCell>
                                            <TableCell>
                                                <Badge variant="outline" className={cn(
                                                    task.status === "Completed" ? "bg-green-100 text-green-700 border-green-200" :
                                                        task.status === "In Progress" ? "bg-blue-100 text-blue-700 border-blue-200" : ""
                                                )}>{task.status}</Badge>
                                            </TableCell>
                                            <TableCell className="text-muted-foreground text-sm">
                                                {task.stateDate ? format(parseISO(task.stateDate), 'MMM d, yyyy') : '-'}
                                            </TableCell>
                                            <TableCell className="text-muted-foreground text-sm pr-6">
                                                {task.endDate ? format(parseISO(task.endDate), 'MMM d, yyyy') : '-'}
                                            </TableCell>
                                        </TableRow>
                                    ))}
                                </TableBody>
                            </Table>
                        </div>
                    </CardContent>
                </Card>
            )}

            {/* Task Details Sheet - REPLACED WITH CUSTOM SHEET */}
            <CustomSheet
                open={isSheetOpen}
                onOpenChange={setIsSheetOpen}
                title={selectedTask?.title}
                description={selectedTask ? `Task ID: #${selectedTask.id}` : undefined}
                className="w-[400px] sm:w-[540px]"
            >
                {selectedTask && (
                    <div className="space-y-6">

                        <div className="space-y-6">
                            <div className="flex flex-wrap gap-2">
                                <Badge variant="outline" className={cn(
                                    "px-2 py-1 text-sm font-medium",
                                    selectedTask.status === "Completed" ? "bg-green-100 text-green-700 border-green-200" :
                                        selectedTask.status === "In Progress" ? "bg-blue-100 text-blue-700 border-blue-200" :
                                            "bg-gray-100 text-gray-700 border-gray-200"
                                )}>
                                    {selectedTask.status}
                                </Badge>
                                <Badge variant="secondary" className="px-2 py-1 text-sm">
                                    {selectedTask.type}
                                </Badge>
                                <Badge variant="outline" className="px-2 py-1 text-sm">
                                    Priority: {selectedTask.priority}
                                </Badge>
                            </div>

                            <div className="grid grid-cols-2 gap-4">
                                <div className="space-y-1">
                                    <h4 className="text-sm font-medium text-muted-foreground flex items-center gap-2">
                                        <CalendarIcon className="h-4 w-4" /> Start Date
                                    </h4>
                                    <p className="font-medium">
                                        {selectedTask.stateDate ? format(parseISO(selectedTask.stateDate), 'PPP') : 'Not set'}
                                    </p>
                                </div>
                                <div className="space-y-1">
                                    <h4 className="text-sm font-medium text-muted-foreground flex items-center gap-2">
                                        <CalendarIcon className="h-4 w-4" /> End Date
                                    </h4>
                                    <p className="font-medium">
                                        {selectedTask.endDate ? format(parseISO(selectedTask.endDate), 'PPP') : 'Not set'}
                                    </p>
                                </div>
                                <div className="space-y-1">
                                    <h4 className="text-sm font-medium text-muted-foreground flex items-center gap-2">
                                        <User className="h-4 w-4" /> Assignee
                                    </h4>
                                    <p className="font-medium">
                                        {employees.find(e => e.id === selectedTask.assignedTo)?.firstName} {employees.find(e => e.id === selectedTask.assignedTo)?.lastName || 'Unassigned'}
                                    </p>
                                </div>
                                <div className="space-y-1">
                                    <h4 className="text-sm font-medium text-muted-foreground flex items-center gap-2">
                                        <Clock className="h-4 w-4" /> Created At
                                    </h4>
                                    <p className="font-medium">
                                        {selectedTask.createdAt ? format(parseISO(selectedTask.createdAt), 'PPP') : 'Unknown'}
                                    </p>
                                </div>
                            </div>

                            <div className="space-y-2">
                                <h4 className="text-sm font-medium text-muted-foreground">Description</h4>
                                <div className="p-4 rounded-md bg-muted/50 text-sm whitespace-pre-wrap">
                                    {selectedTask.description || 'No description provided.'}
                                </div>
                            </div>

                            <div className="flex justify-end pt-4 border-t">
                                <Button onClick={() => navigate(`/tasks/${selectedTask.id}`)}>
                                    View Full Details
                                </Button>
                            </div>
                        </div>
                    </div>
                )}
            </CustomSheet>
        </div>
    );
};