frontend/src/components/Layout/AppLayout.tsx

import { useCallback, useRef, useEffect, useState } from 'react';
import {
  Avatar,
  Box,
  Drawer,
  Typography,
  IconButton,
  Alert,
  AlertTitle,
  useMediaQuery,
  useTheme,
  CircularProgress,
} from '@mui/material';
import MenuIcon from '@mui/icons-material/Menu';
import ChevronLeftIcon from '@mui/icons-material/ChevronLeft';
import DarkModeOutlinedIcon from '@mui/icons-material/DarkModeOutlined';
import LightModeOutlinedIcon from '@mui/icons-material/LightModeOutlined';
import { logger } from '@/utils/logger';

import { useSessionStore } from '@/store/sessionStore';
import { useAgentStore } from '@/store/agentStore';
import { useLayoutStore } from '@/store/layoutStore';
import { useAgentWebSocket } from '@/hooks/useAgentWebSocket';
import SessionSidebar from '@/components/SessionSidebar/SessionSidebar';
import CodePanel from '@/components/CodePanel/CodePanel';
import ChatInput from '@/components/Chat/ChatInput';
import MessageList from '@/components/Chat/MessageList';
import WelcomeScreen from '@/components/WelcomeScreen/WelcomeScreen';
import { apiFetch } from '@/utils/api';
import type { Message } from '@/types/agent';
import { triggerLogin } from '@/hooks/useAuth';

const DRAWER_WIDTH = 260;

export default function AppLayout() {
  const { sessions, activeSessionId, createSession, deleteSession, updateSessionTitle } = useSessionStore();
  const { isConnected, isProcessing, getMessages, addMessage, setProcessing, llmHealthError, setLlmHealthError, user, setPlan, setPanelContent } = useAgentStore();
  const { 
    isLeftSidebarOpen, 
    isRightPanelOpen, 
    rightPanelWidth,
    themeMode,
    setRightPanelWidth,
    setLeftSidebarOpen,
    toggleLeftSidebar, 
    toggleTheme,
  } = useLayoutStore();

  const theme = useTheme();
  const isMobile = useMediaQuery(theme.breakpoints.down('md'));

  const isResizing = useRef(false);
  const [isAutoCreating, setIsAutoCreating] = useState(false);
  const [autoCreateError, setAutoCreateError] = useState<string | null>(null);

  const handleMouseMove = useCallback((e: MouseEvent) => {
    if (!isResizing.current) return;
    const newWidth = window.innerWidth - e.clientX;
    const maxWidth = window.innerWidth * 0.6;
    const minWidth = 300;
    if (newWidth > minWidth && newWidth < maxWidth) {
      setRightPanelWidth(newWidth);
    }
  }, [setRightPanelWidth]);

  const stopResizing = useCallback(() => {
    isResizing.current = false;
    document.removeEventListener('mousemove', handleMouseMove);
    document.removeEventListener('mouseup', stopResizing);
    document.body.style.cursor = 'default';
  }, [handleMouseMove]);

  const startResizing = useCallback((e: React.MouseEvent) => {
    e.preventDefault();
    isResizing.current = true;
    document.addEventListener('mousemove', handleMouseMove);
    document.addEventListener('mouseup', stopResizing);
    document.body.style.cursor = 'col-resize';
  }, [handleMouseMove, stopResizing]);

  useEffect(() => {
    return () => {
      document.removeEventListener('mousemove', handleMouseMove);
      document.removeEventListener('mouseup', stopResizing);
    };
  }, [handleMouseMove, stopResizing]);

  // ── Auto-create session if none exist ───────────────────────────
  useEffect(() => {
    const hasAnySessions = sessions.length > 0;
    const isAuthenticated = user?.authenticated;
    const isDevUser = user?.username === 'dev';

    if (!hasAnySessions && (isAuthenticated || isDevUser) && !isAutoCreating) {
      (async () => {
        setIsAutoCreating(true);
        setAutoCreateError(null);
        try {
          const response = await apiFetch('/api/session', { method: 'POST' });
          if (response.status === 503) {
            const data = await response.json();
            setAutoCreateError(data.detail || 'Server is at capacity. Please try again later.');
            return;
          }
          if (response.status === 401) {
            triggerLogin();
            return;
          }
          if (!response.ok) {
            setAutoCreateError('Failed to create session. Please try again.');
            return;
          }
          const data = await response.json();
          createSession(data.session_id);
          setPlan([]);
          setPanelContent(null);
        } catch (e) {
          logger.error('Auto-create session failed:', e);
          setAutoCreateError('Failed to auto-create session.');
        } finally {
          setIsAutoCreating(false);
        }
      })();
    }
  }, [sessions.length, user, createSession, setPlan, setPanelContent, isAutoCreating]);

  // ── LLM health check on mount ───────────────────────────────────
  useEffect(() => {
    let cancelled = false;
    (async () => {
      try {
        const res = await apiFetch('/api/health/llm');
        const data = await res.json();
        if (!cancelled && data.status === 'error') {
          setLlmHealthError({
            error: data.error || 'Unknown LLM error',
            errorType: data.error_type || 'unknown',
            model: data.model,
          });
        } else if (!cancelled) {
          setLlmHealthError(null);
        }
      } catch {
        // Backend unreachable — not an LLM issue, ignore
      }
    })();
    return () => { cancelled = true; };
  }, []); // eslint-disable-line react-hooks/exhaustive-deps

  const messages = activeSessionId ? getMessages(activeSessionId) : [];
  const hasAnySessions = sessions.length > 0;

  useAgentWebSocket({
    sessionId: activeSessionId,
    onReady: () => logger.log('Agent ready'),
    onError: (error) => logger.error('Agent error:', error),
    onSessionDead: (deadSessionId) => {
      logger.log('Removing dead session:', deadSessionId);
      deleteSession(deadSessionId);
    },
  });

  const handleSendMessage = useCallback(
    async (text: string) => {
      if (!activeSessionId || !text.trim() || isProcessing) return;
      
      // Lock input immediately to prevent double-sends
      setProcessing(true);

      const userMsg: Message = {
        id: `user_${Date.now()}`,
        role: 'user',
        content: text.trim(),
        timestamp: new Date().toISOString(),
      };
      addMessage(activeSessionId, userMsg);

      // Auto-title the session from the first user message (async, non-blocking)
      const currentMessages = getMessages(activeSessionId);
      const isFirstMessage = currentMessages.filter((m) => m.role === 'user').length <= 1;
      if (isFirstMessage) {
        const sessionId = activeSessionId;
        apiFetch('/api/title', {
          method: 'POST',
          body: JSON.stringify({ session_id: sessionId, text: text.trim() }),
        })
          .then((res) => res.json())
          .then((data) => {
            if (data.title) updateSessionTitle(sessionId, data.title);
          })
          .catch(() => {
            const raw = text.trim();
            updateSessionTitle(sessionId, raw.length > 40 ? raw.slice(0, 40) + '…' : raw);
          });
      }

      try {
        await apiFetch('/api/submit', {
          method: 'POST',
          body: JSON.stringify({
            session_id: activeSessionId,
            text: text.trim(),
          }),
        });
      } catch (e) {
        logger.error('Send failed:', e);
      }
    },
    [activeSessionId, addMessage, getMessages, updateSessionTitle, isProcessing, setProcessing]
  );

  // Close sidebar on mobile after selecting a session
  const handleSidebarClose = useCallback(() => {
    if (isMobile) setLeftSidebarOpen(false);
  }, [isMobile, setLeftSidebarOpen]);

  // ── LLM error banner (shared) ─────────────────────────────────────
  const llmBanner = llmHealthError && (
    <Alert
      severity="error"
      variant="filled"
      onClose={() => setLlmHealthError(null)}
      sx={{ borderRadius: 0, flexShrink: 0, '& .MuiAlert-message': { flex: 1 } }}
    >
      <AlertTitle sx={{ fontWeight: 700, fontSize: '0.85rem' }}>
        {llmHealthError.errorType === 'credits'
          ? 'API Credits Exhausted'
          : llmHealthError.errorType === 'auth'
          ? 'Invalid API Key'
          : llmHealthError.errorType === 'rate_limit'
          ? 'Rate Limited'
          : llmHealthError.errorType === 'network'
          ? 'LLM Provider Unreachable'
          : 'LLM Error'}
      </AlertTitle>
      <Typography variant="body2" sx={{ fontSize: '0.8rem', opacity: 0.9 }}>
        Model: <strong>{llmHealthError.model}</strong> — {llmHealthError.error.slice(0, 200)}
      </Typography>
    </Alert>
  );

  // ── Welcome screen: no sessions at all ────────────────────────────
  if (!hasAnySessions) {
    return (
      <Box sx={{ width: '100%', height: '100%', display: 'flex', flexDirection: 'column' }}>
        {llmBanner}
        {isAutoCreating ? (
          <Box sx={{ flex: 1, display: 'flex', alignItems: 'center', justifyContent: 'center', background: 'var(--body-gradient)' }}>
            <CircularProgress color="primary" />
          </Box>
        ) : autoCreateError ? (
          <Box sx={{ flex: 1, display: 'flex', alignItems: 'center', justifyContent: 'center', background: 'var(--body-gradient)', p: 3 }}>
            <Alert severity="error" variant="outlined" sx={{ maxWidth: 400 }}>
              <AlertTitle>Auto-initialization Failed</AlertTitle>
              {autoCreateError}
              <Box sx={{ mt: 2 }}>
                <WelcomeScreen />
              </Box>
            </Alert>
          </Box>
        ) : (
          <WelcomeScreen />
        )}
      </Box>
    );
  }

  // ── Sidebar drawer ────────────────────────────────────────────────
  const sidebarDrawer = (
    <Drawer
      variant={isMobile ? 'temporary' : 'persistent'}
      anchor="left"
      open={isLeftSidebarOpen}
      onClose={() => setLeftSidebarOpen(false)}
      ModalProps={{ keepMounted: true }} // Better mobile perf
      sx={{
        '& .MuiDrawer-paper': {
          boxSizing: 'border-box',
          width: DRAWER_WIDTH,
          borderRight: '1px solid',
          borderColor: 'divider',
          top: 0,
          height: '100%',
          bgcolor: 'var(--panel)',
        },
      }}
    >
      <SessionSidebar onClose={handleSidebarClose} />
    </Drawer>
  );

  // ── Main chat interface ───────────────────────────────────────────
  return (
    <Box sx={{ display: 'flex', width: '100%', height: '100%' }}>
      {/* ── Left Sidebar ─────────────────────────────────────────── */}
      {isMobile ? (
        // Mobile: temporary overlay drawer (no reserved width)
        sidebarDrawer
      ) : (
        // Desktop: persistent drawer with reserved width
        <Box
          component="nav"
          sx={{
            width: isLeftSidebarOpen ? DRAWER_WIDTH : 0,
            flexShrink: 0,
            transition: isResizing.current ? 'none' : 'width 0.2s',
            overflow: 'hidden',
          }}
        >
          {sidebarDrawer}
        </Box>
      )}

      {/* ── Main Content (header + chat + code panel) ────────────── */}
      <Box
        sx={{
          flexGrow: 1,
          height: '100%',
          display: 'flex',
          flexDirection: 'column',
          transition: isResizing.current ? 'none' : 'width 0.2s',
          overflow: 'hidden',
          minWidth: 0,
        }}
      >
        {/* ── Top Header Bar ─────────────────────────────────────── */}
        <Box sx={{ 
          height: { xs: 52, md: 60 },
          px: { xs: 1, md: 2 }, 
          display: 'flex', 
          alignItems: 'center', 
          borderBottom: 1, 
          borderColor: 'divider',
          bgcolor: 'background.default',
          zIndex: 1200,
          flexShrink: 0,
        }}>
          <IconButton onClick={toggleLeftSidebar} size="small">
            {isLeftSidebarOpen && !isMobile ? <ChevronLeftIcon /> : <MenuIcon />}
          </IconButton>
          
          <Box sx={{ flex: 1, display: 'flex', justifyContent: 'center', alignItems: 'center', gap: 0.75 }}>
            <Box
              component="img"
              src="https://huggingface.co/front/assets/huggingface_logo-noborder.svg"
              alt="HF"
              sx={{ width: { xs: 20, md: 22 }, height: { xs: 20, md: 22 } }}
            />
            <Typography
              variant="subtitle1"
              sx={{
                fontWeight: 700,
                color: 'var(--text)',
                letterSpacing: '-0.01em',
                fontSize: { xs: '0.88rem', md: '0.95rem' },
              }}
            >
              ML Agent
            </Typography>
          </Box>

          <Box sx={{ display: 'flex', alignItems: 'center', gap: 0.5 }}>
            <IconButton
              onClick={toggleTheme}
              size="small"
              sx={{ color: 'var(--muted-text)' }}
            >
              {themeMode === 'dark' ? <LightModeOutlinedIcon fontSize="small" /> : <DarkModeOutlinedIcon fontSize="small" />}
            </IconButton>

            <Box sx={{ display: 'flex', alignItems: 'center', ml: 1, gap: 1 }}>
              {user?.username && (
                <Typography variant="caption" sx={{ color: 'var(--muted-text)', display: { xs: 'none', sm: 'block' } }}>
                  {user.username}
                </Typography>
              )}
              <Avatar
                src={user?.picture}
                sx={{
                  width: 28,
                  height: 28,
                  border: '1px solid',
                  borderColor: 'divider',
                }}
              />
            </Box>
          </Box>
        </Box>

        {/* ── Chat + Code Panel ──────────────────────────────────── */}
        <Box sx={{ 
          flexGrow: 1, 
          display: 'flex', 
          overflow: 'hidden', 
          position: 'relative',
          bgcolor: 'background.default'
        }}>
          {/* Chat Column */}
          <Box sx={{ 
            flexGrow: 1, 
            display: 'flex', 
            flexDirection: 'column',
            minWidth: 0,
            height: '100%',
            position: 'relative',
          }}>
            <Box sx={{ flexGrow: 1, overflow: 'visible' }}>
              <MessageList messages={messages} isProcessing={isProcessing} />
            </Box>
            <Box sx={{ p: { xs: 1.5, md: 2.5 }, pt: 0 }}>
              <ChatInput 
                onSend={handleSendMessage} 
                disabled={!isConnected || isProcessing}
              />
            </Box>
          </Box>

          {/* Resizable Code Panel (Desktop only) */}
          {!isMobile && isRightPanelOpen && (
            <>
              {/* Resize Handle */}
              <Box
                onMouseDown={startResizing}
                sx={{
                  width: '4px',
                  cursor: 'col-resize',
                  bgcolor: 'transparent',
                  '&:hover': { bgcolor: 'primary.main', opacity: 0.5 },
                  zIndex: 10,
                  position: 'relative',
                  '&::after': {
                    content: '""',
                    position: 'absolute',
                    left: '50%',
                    top: '50%',
                    transform: 'translate(-50%, -50%)',
                    height: '24px',
                    width: '1px',
                    bgcolor: 'divider',
                  }
                }}
              />
              <Box sx={{ 
                width: rightPanelWidth,
                flexShrink: 0,
                height: '100%',
                bgcolor: 'var(--panel)',
                borderLeft: '1px solid',
                borderColor: 'divider',
                display: 'flex',
                flexDirection: 'column',
                overflow: 'hidden',
              }}>
                <CodePanel />
              </Box>
            </>
          )}

          {/* Mobile Code Panel Overlay */}
          {isMobile && isRightPanelOpen && (
            <Box sx={{
              position: 'absolute',
              top: 0,
              left: 0,
              right: 0,
              bottom: 0,
              zIndex: 1300,
              bgcolor: 'background.default',
              display: 'flex',
              flexDirection: 'column',
            }}>
              <CodePanel />
            </Box>
          )}
        </Box>
      </Box>
    </Box>
  );
}