import React, { useState, useEffect, useRef } from "react";
import axios from "axios";
import {
  Box,
  Button,
  CircularProgress,
  InputLabel,
  MenuItem,
  Select,
  Typography,
  FormControl,
  Grid,
  Paper,
  Snackbar,
  Alert,
  Dialog,
  DialogTitle,
  DialogContent,
  DialogActions,
  useMediaQuery,
  Menu,
  MenuItem as ContextMenuItem,
} from "@mui/material";
import { useTheme } from "@mui/material/styles";
import CloudUploadIcon from "@mui/icons-material/CloudUpload";
import DownloadIcon from "@mui/icons-material/Download";
import FolderIcon from "@mui/icons-material/Folder";
import InsertDriveFileIcon from "@mui/icons-material/InsertDriveFile";
import { SimpleTreeView, TreeItem } from "@mui/x-tree-view";
import Prism from './prism-setup';
import { detectLanguage } from './utils/languageDetect';

const FileUploadForm = () => {
  const [file, setFile] = useState(null);
  const [frontend, setFrontend] = useState("React");
  const [backend, setBackend] = useState("Flask");
  const [database, setDatabase] = useState("PostgreSQL");
  const [loading, setLoading] = useState(false);
  const [showSnackbar, setShowSnackbar] = useState(false);
  const [showModal, setShowModal] = useState(false);
  const [zipBlob, setZipBlob] = useState(null);
  const [fileListInfo, setFileListInfo] = useState({});
  const [selectedFile, setSelectedFile] = useState(null);
  const [expandedIds, setExpandedIds] = useState([]); // ✅ array, not Set

  const [contextAnchor, setContextAnchor] = useState(null);
  const [contextFile, setContextFile] = useState(null);

  const codeRef = useRef(null);

 


  

  const theme = useTheme();
  const isMobile = useMediaQuery(theme.breakpoints.down("sm"));

  const buildTree = (files) => {
    const root = {};
    for (const path in files) {
      const parts = path.split("/");
      let current = root;
      for (let i = 0; i < parts.length; i++) {
        const part = parts[i];
        if (!current[part]) {
          current[part] = i === parts.length - 1 ? files[path] : {};
        }
        current = current[part];
      }
    }
    return root;
  };

  const renderTree = (node, path = "") => {
    if (!node || typeof node !== "object") return null;

    return Object.entries(node).map(([key, value]) => {
      const fullPath = path ? `${path}/${key}` : key;
      const itemId = fullPath.replaceAll("/", "-");
      const isFile = value && typeof value.preview !== "undefined";

      return (
        <TreeItem
          key={fullPath}
          itemId={itemId}
          label={
            <Box
              display="flex"
              alignItems="center"
              onContextMenu={(e) => handleContextMenu(e, fullPath, value)}
            >
              {isFile ? <InsertDriveFileIcon sx={{ mr: 1 }} /> : <FolderIcon sx={{ mr: 1 }} />}
              <Typography variant="body2">
                {key}
                {isFile ? ` (${(value.size / 1024).toFixed(1)} KB)` : ""}
              </Typography>
            </Box>
          }
          onClick={() => {
            if (isFile) {
        setSelectedFile({ name: fullPath, content: value.preview });
        Prism.highlightAll();
      } else {
        setExpandedIds((prev) =>
          prev.includes(itemId)
            ? prev.filter((id) => id !== itemId) // collapse
            : [...prev, itemId] // expand
        );
      }
          }}
        >
          {!isFile && renderTree(value, fullPath)}
        </TreeItem>
      );
    });
  };

  const handleDownload = async (e) => {
    e.preventDefault();
    if (!file) return alert("Please select a requirement document!");

    const formData = new FormData();
    formData.append("file", file);
    formData.append("frontend", frontend);
    formData.append("backend", backend);
    formData.append("database", database);

    try {
      setLoading(true);
      const response = await axios.post(
        "https://fredericksundeep-aichatmatedev.hf.space/chat-stream-doc",
        formData,
        { responseType: "blob" }
      );

      const blob = new Blob([response.data], { type: "application/zip" });
      setZipBlob(blob);

      const JSZip = (await import("jszip")).default;
      const zip = await JSZip.loadAsync(blob);
      const fileInfo = {};
      await Promise.all(
        Object.entries(zip.files).map(async ([name, file]) => {
          if (!file.dir) {
            const content = await file.async("string");
            fileInfo[name] = {
              size: file._data.uncompressedSize,
              preview: content.slice(0, 1000),
            };
          }
        })
      );
      setFileListInfo(fileInfo);
      setExpandedIds([]); // reset tree state
      setShowModal(true);
    } catch (error) {
      console.error("Error:", error);
      alert("Failed to generate project.");
    } finally {
      setLoading(false);
    }
  };

  const handleZipDownload = () => {
    const url = window.URL.createObjectURL(zipBlob);
    const a = document.createElement("a");
    a.href = url;
    a.download = "generated_project.zip";
    document.body.appendChild(a);
    a.click();
    window.URL.revokeObjectURL(url);
    setShowSnackbar(true);
    setShowModal(false);
  };

  const handleContextMenu = (event, filename, value) => {
    event.preventDefault();
    if (typeof value.preview === "undefined") return; // not a file
    setContextFile({ name: filename, content: value.preview });
    setContextAnchor(event.currentTarget);
  };

  const handleDownloadFile = () => {
    if (!contextFile) return;
    const blob = new Blob([contextFile.content], { type: "text/plain" });
    const a = document.createElement("a");
    a.href = URL.createObjectURL(blob);
    a.download = contextFile.name.split("/").pop();
    a.click();
    setContextAnchor(null);
  };


// ✅ Expand/collapse logic
const expandAll = () => {
  const allPaths = Object.keys(fileListInfo); // files only
  const allIds = new Set();

  allPaths.forEach((filePath) => {
    const parts = filePath.split("/");
    for (let i = 1; i <= parts.length; i++) {
      const partialPath = parts.slice(0, i).join("/");
      allIds.add(partialPath.replaceAll("/", "-")); // folder or file
    }
  });

  setExpandedIds(Array.from(allIds));
};


const collapseAll = () => {
  setExpandedIds([]); // ✅ empty array
};




const { lang, code } = React.useMemo(() => {
  if (selectedFile && selectedFile.content) {
    const lines = selectedFile.content.split('\n');
    const firstLine = lines[0].trim();
    const hasLang = /^[a-zA-Z]+$/.test(firstLine);
    const lang = hasLang ? firstLine.toLowerCase() : detectLanguage(selectedFile.content) || 'markdown';
    const code = hasLang ? lines.slice(1).join('\n') : selectedFile.content;
    return { lang, code };
  }
  return { lang: 'markdown', code: '' };
}, [selectedFile]);


useEffect(() => {
    if (codeRef.current) {
      Prism.highlightElement(codeRef.current);
    }
  }, [code, lang]);
  return (
    <Box sx={{
      backgroundColor: "#121212",
      minHeight: "100vh",
      color: "#fff",
      padding: 4,
      display: "flex",
      justifyContent: "center",
      alignItems: "center",
      flexDirection: "column",
    }}>
      <Paper elevation={6} sx={{ padding: 4, backgroundColor: "#1e1e1e", maxWidth: 600, width: "100%" }}>
        <Box textAlign="center" my={4}>
          <Typography variant="h4" component="h1">🧠 AI - Full-Stack Project Generator</Typography>
        </Box>

        <form onSubmit={handleDownload}>
          <Grid container spacing={2} direction="column">
            <Grid item>
              <FormControl fullWidth>
                <Button
                  variant="contained"
                  component="label"
                  startIcon={<CloudUploadIcon />}
                  fullWidth
                >
                  {file ? file.name : "Upload Requirement Document"}
                  <input type="file" hidden required onChange={(e) => setFile(e.target.files[0])} accept=".pdf,.txt" />
                </Button>
              </FormControl>
            </Grid>

            {["Frontend", "Backend", "Database"].map((label, idx) => (
              <Grid item key={label}>
                <FormControl fullWidth>
                  <InputLabel>{label}</InputLabel>
                  <Select
                    value={label === "Frontend" ? frontend : label === "Backend" ? backend : database}
                    onChange={(e) => {
                      if (label === "Frontend") setFrontend(e.target.value);
                      else if (label === "Backend") setBackend(e.target.value);
                      else setDatabase(e.target.value);
                    }}
                    label={label}
                  >
                    {label === "Frontend" && ["React", "Angular", "Vue", "Next.js", "Svelte", "HTML/CSS/JS"].map(opt => <MenuItem key={opt} value={opt}>{opt}</MenuItem>)}
                    {label === "Backend" && ["Flask", "Node.js", "Django", "Express.js", "Spring Boot", "FastAPI"].map(opt => <MenuItem key={opt} value={opt}>{opt}</MenuItem>)}
                    {label === "Database" && ["PostgreSQL", "MongoDB", "MySQL", "SQLite", "Firebase", "Supabase"].map(opt => <MenuItem key={opt} value={opt}>{opt}</MenuItem>)}
                  </Select>
                </FormControl>
              </Grid>
            ))}

            <Grid item>
              <Button
                type="submit"
                variant="contained"
                color="secondary"
                startIcon={loading ? <CircularProgress size={20} /> : <DownloadIcon />}
                fullWidth
                disabled={loading}
              >
                {loading ? "Generating..." : "Generate Project ZIP"}
              </Button>
            </Grid>
          </Grid>
        </form>
      </Paper>

      {/* 📦 Preview Dialog */}
      <Dialog open={showModal} onClose={() => setShowModal(false)} maxWidth="md" fullWidth>
        <DialogTitle>📦 Generated Files</DialogTitle>
        <DialogContent dividers>
          <Grid container spacing={2}>
            <Grid item xs={12} md={5}>
              <Box mb={1} display="flex" gap={1}>
                <Button size="small" variant="outlined" onClick={expandAll}>Expand All</Button>
                <Button size="small" variant="outlined" onClick={collapseAll}>Collapse All</Button>
              </Box>
              <SimpleTreeView
  expandedItems={expandedIds}
  onExpandedItemsChange={(newItems) => {
    if (Array.isArray(newItems)) {
      setExpandedIds(newItems); // ✅ directly set array
    } else {
      console.warn("Unexpected expandedItems:", newItems);
      setExpandedIds([]);
    }
  }}
   onItemToggle={(itemId, isExpanded) => {
    setExpandedIds((prev) =>
      isExpanded ? [...prev, itemId] : prev.filter((id) => id !== itemId)
    );
  }}
>
  {renderTree(buildTree(fileListInfo))}
</SimpleTreeView>
            </Grid>

            <Grid item xs={12} md={7}>
              {selectedFile ? (
                <>
                  <Typography variant="subtitle1" gutterBottom>{selectedFile.name}</Typography>
                  <Paper variant="outlined" sx={{ p: 2, maxHeight: 400, overflowY: "auto", backgroundColor: "#1e1e1e" }}>
                    <pre className={`line-numbers language-${lang}`} style={{ borderRadius: '8px', overflowX: 'auto' }}>
  <code ref={codeRef} className={`language-${lang}`}>
    {code}
  </code>
</pre>
                  </Paper>
                </>
              ) : (
                <Typography color="text.secondary">Select a file to preview its content</Typography>
              )}
            </Grid>
          </Grid>
        </DialogContent>
        <DialogActions>
          <Button onClick={handleZipDownload} variant="contained">Download ZIP</Button>
          <Button onClick={() => setShowModal(false)}>Close</Button>
        </DialogActions>
      </Dialog>

      {/* 🎯 Context Menu for Download */}
      <Menu
        anchorEl={contextAnchor}
        open={Boolean(contextAnchor)}
        onClose={() => setContextAnchor(null)}
      >
        <ContextMenuItem onClick={handleDownloadFile}>📥 Download File</ContextMenuItem>
      </Menu>

      {/* ✅ Success Snackbar */}
      <Snackbar
        open={showSnackbar}
        autoHideDuration={4000}
        onClose={() => setShowSnackbar(false)}
        anchorOrigin={{ vertical: "bottom", horizontal: "center" }}
      >
        <Alert severity="success" sx={{ width: "100%" }} onClose={() => setShowSnackbar(false)}>
          ✅ ZIP downloaded successfully!
        </Alert>
      </Snackbar>
    </Box>
  );
};

export default FileUploadForm;