"""Patch ChatInput.tsx - Fixed import for DatasetUploadResponse."""

CHAT_INPUT = "/source/frontend/src/components/Chat/ChatInput.tsx"

def patch():
    import os, re
    if not os.path.exists(CHAT_INPUT):
        print(f"SKIP: {CHAT_INPUT} not found")
        return
        
    with open(CHAT_INPUT, "r", encoding="utf-8") as f:
        content = f.read()
    
    # 1. Add Button to MUI imports
    if "Button" not in content and "@mui/material" in content:
        content = content.replace(
            "import { Box, IconButton, Stack, Tooltip } from '@mui/material';",
            "import { Box, IconButton, Stack, Tooltip, Button } from '@mui/material';"
        )
    
    # 2. Fix DatasetUploadResponse import - check existing imports
    if "DatasetUploadResponse" not in content:
        # Find existing import from types/agent
        if "from '@/types/agent'" in content:
            content = content.replace(
                "from '@/types/agent'",
                "from '@/types/agent'"
            )
            # Add DatasetUploadResponse to existing import
            content = content.replace(
                "import type { ",
                "import type { DatasetUploadResponse, "
            )
        else:
            content = content.replace(
                "import { apiFetch } from '@/utils/api';",
                "import { apiFetch } from '@/utils/api';\nimport type { DatasetUploadResponse } from '@/types/agent';"
            )
    
    # 3. Add props to interface
    if "showAutoContinue" not in content:
        props_code = '''
interface ChatInputProps {
  sessionId: string;
  initialModelPath: string | null | undefined;
  onSend: (text: string) => Promise<void>;
  onStop: () => void;
  onDatasetUploaded: () => Promise<boolean>;
  isProcessing: boolean;
  disabled: boolean;
  placeholder?: string;
  showAutoContinue?: boolean;
  taskIncompleteInfo?: { incompletePlan: Array<{ id: string; content: string; status: string }> } | null;
  startAutoContinue?: () => void;
  cancelAutoContinue?: () => void;
}
'''
        match = re.search(r'interface ChatInputProps[^}]*[}][^}]*[}]', content, re.DOTALL)
        if match:
            content = content.replace(match.group(0), props_code)
        else:
            content = content.replace(
                "export default function ChatInput",
                props_code + "\nexport default function ChatInput"
            )
    
    # 4. Add auto-continue buttons before the flex box
    buttons_code = '''
      {/* Auto-continue controls for free models */}
      {showAutoContinue && taskIncompleteInfo && (
        <Stack direction="row" spacing={1} sx={{ mb: 1, px: 1 }}>
          <Button
            variant="contained"
            color="primary"
            size="small"
            onClick={startAutoContinue}
            sx={{ textTransform: 'none', fontSize: '0.75rem' }}
          >
            Tự động tiếp tục (10s)
          </Button>
          <Button
            variant="outlined"
            color="secondary"
            size="small"
            onClick={cancelAutoContinue}
            sx={{ textTransform: 'none', fontSize: '0.75rem' }}
          >
            Tạm dừng
          </Button>
        </Stack>
      )}
'''
    
    if "Tự động tiếp tục" not in content:
        content = content.replace(
            "<Box sx={{ flex: 1 ",
            buttons_code + "        <Box sx={{ flex: 1 "
        )
    
    with open(CHAT_INPUT, "w", encoding="utf-8") as f:
        f.write(content)
    print(f"OK: Patched {CHAT_INPUT}")

if __name__ == "__main__":
    patch()