import React, { useEffect, useState, useRef } from "react";
import { Editor, DiffEditor } from "@monaco-editor/react";
import { create } from "zustand";
import { io } from "socket.io-client";

const BASE_URL = "http://localhost:7860";

const useStore = create((set, get) => ({
  files: [],
  selectedFile: null,
  editorContent: "",
  dbSchema: { tables: [] },
  chat: [],
  consoleOutput: "",
  loading: false,
  fileProposals: {},
  showDiffForFile: null,
  fileContents: {},
  setFiles: (files) => set({ files }),
  setSelectedFile: (name) => set({ selectedFile: { name } }),
  setEditorContent: (editorContent) => set({ editorContent }),
  setDbSchema: (dbSchema) => set({ dbSchema }),
  pushChat: (msg) => set((s) => ({ chat: [...s.chat, msg] })),
  setConsoleOutput: (consoleOutput) => set({ consoleOutput }),
  setLoading: (loading) => set({ loading }),
  setFileProposal: (name, details) => set((s) => ({ fileProposals: { ...s.fileProposals, [name]: { ...details, active: true } } })),
  setShowDiffForFile: (fileName) => set({ showDiffForFile: fileName }),
  openFile: async (name) => {
    const currentState = get();
    if (currentState.selectedFile?.name && currentState.editorContent) {
      set({ fileContents: { ...currentState.fileContents, [currentState.selectedFile.name]: currentState.editorContent } });
    }
    set({ loading: true });
    let content = "";
    const state = get();
    const proposal = state.fileProposals[name];
    if (proposal?.isNew) {
      content = "";
    } else {
      content = state.fileContents[name];
      if (!content) {
        try {
          const res = await fetch(`${BASE_URL}/files/${encodeURIComponent(name)}`);
          if (!res.ok) throw new Error("Failed to load file");
          const data = await res.json();
          content = data.content;
          set({ fileContents: { ...get().fileContents, [name]: content } });
        } catch (e) {
          console.warn(e);
          content = `-- demo content for ${name}\nSELECT 1;`;
          set({ fileContents: { ...get().fileContents, [name]: content } });
        }
      }
    }
    set({ selectedFile: { name }, editorContent: content, loading: false });
    const finalState = get();
    const prop = finalState.fileProposals[name];
    if (prop?.active) {
      set({ showDiffForFile: name });
    } else {
      set({ showDiffForFile: null });
    }
  },
  acceptProposal: () => {
    const state = get();
    const file = state.selectedFile?.name;
    if (state.showDiffForFile !== file) return;
    const prop = state.fileProposals[file];
    if (!prop) return;
    if (prop.isNew) {
      const confirmedName = prompt("Save new file as (.sql required):", prop.suggestedName);
      if (!confirmedName?.trim() || !confirmedName.endsWith(".sql")) {
        return;
      }
      const finalName = confirmedName.trim();
      fetch(`${BASE_URL}/files`, {
        method: "POST",
        headers: { "Content-Type": "application/json" },
        body: JSON.stringify({ filename: finalName, content: prop.code }),
      })
        .then((res) => {
          if (!res.ok) throw new Error(`HTTP ${res.status}: ${res.statusText}`);
          return res.json();
        })
        .then(() => {
          const newProps = { ...state.fileProposals };
          delete newProps[file];
          const currentFiles = get().files;
          set({
            files: [...currentFiles, { name: finalName }],
            editorContent: prop.code,
            fileContents: { ...get().fileContents, [finalName]: prop.code },
            showDiffForFile: null,
            fileProposals: newProps,
          });
          if (finalName !== file) {
            set({ selectedFile: { name: finalName } });
          }
        })
        .catch((e) => console.error("Failed to create file:", e));
    } else {
      fetch(`${BASE_URL}/files/${encodeURIComponent(file)}`, {
        method: "PUT",
        headers: { "Content-Type": "application/json" },
        body: JSON.stringify({ content: prop.code }),
      })
        .then((res) => {
          if (!res.ok) throw new Error(`HTTP ${res.status}: ${res.statusText}`);
          const newProps = { ...state.fileProposals };
          delete newProps[file];
          set({
            editorContent: prop.code,
            fileContents: { ...state.fileContents, [file]: prop.code },
            showDiffForFile: null,
            fileProposals: newProps,
          });
        })
        .catch((e) => console.error("Failed to save after accept:", e));
    }
  },
  rejectProposal: () => {
    const state = get();
    const file = state.selectedFile?.name;
    if (state.showDiffForFile !== file) return;
    const newProps = { ...state.fileProposals };
    delete newProps[file];
    set({ showDiffForFile: null, fileProposals: newProps });
  },
  refreshFiles: async () => {
    try {
      const res = await fetch(`${BASE_URL}/files`);
      if (!res.ok) throw new Error("Failed to load files");
      let data = await res.json();
      data = data.filter((f) => f.name.endsWith(".sql"));
      set({ files: data || [] });
    } catch (e) {
      console.warn(e);
    }
  },
}));

const Card = ({ children, className = "" }) => (
  <div className={`rounded-xl border bg-white shadow-sm ${className}`}>{children}</div>
);
const CardContent = ({ children, className = "" }) => (
  <div className={`p-3 ${className}`}>{children}</div>
);
const IconButton = ({ children, onClick, className = "" }) => (
  <button
    onClick={onClick}
    className={`px-3 py-1 rounded-md border hover:shadow-sm text-sm bg-blue-500 text-white ${className}`}
  >
    {children}
  </button>
);

function FileExplorer() {
  const { files, selectedFile, loading, fileProposals, fileContents, editorContent } = useStore();
  const openFile = useStore((s) => s.openFile);

  const loadFiles = async () => {
    useStore.setState({ loading: true });
    try {
      const res = await fetch(`${BASE_URL}/files`);
      if (!res.ok) throw new Error("Failed to load files");
      let data = await res.json();
      data = data.filter((f) => f.name.endsWith(".sql"));
      useStore.setState({ files: data || [] });
    } catch (e) {
      console.warn(e);
    } finally {
      useStore.setState({ loading: false });
    }
  };

  useEffect(() => {
    loadFiles();
  }, []);

  useEffect(() => {
    if (files.length > 0 && !selectedFile?.name) {
      openFile(files[0].name);
    }
  }, [files, selectedFile?.name, openFile]);

  async function createNewFile() {
    const name = prompt("New file name (must end with .sql):");
    if (!name || !name.endsWith(".sql")) return;
    try {
      await fetch(`${BASE_URL}/files`, {
        method: "POST",
        headers: { "Content-Type": "application/json" },
        body: JSON.stringify({ filename: name, content: "" }),
      });
      loadFiles();
    } catch (e) {
      console.warn(e);
    }
  }

  const proposedNews = Object.entries(fileProposals).filter(([n, p]) => p.isNew);

  return (
    <Card className="h-full flex flex-col">
      <CardContent className="flex items-center justify-between">
        <h3 className="text-sm font-semibold">Files</h3>
        <div className="flex gap-2">
          <IconButton onClick={createNewFile}>New</IconButton>
        </div>
      </CardContent>

      <div className="px-3 pb-3 overflow-auto">
        <ul className="space-y-1">
          {files.map((f) => {
            const hasProposal = fileProposals[f.name]?.active;
            const isDirty = f.name === selectedFile?.name && (!fileContents[f.name] || editorContent !== fileContents[f.name]);
            return (
              <li key={f.name}>
                <button
                  onClick={() => openFile(f.name)}
                  className={`w-full text-left px-2 py-1 rounded-md text-white hover:bg-slate-50 ${selectedFile?.name === f.name ? "bg-slate-100 font-medium" : ""}`}
                >
                  <div className="flex items-center justify-between">
                    <span>{f.name}{isDirty ? " *" : ""}</span>
                    {hasProposal && <div className="w-2 h-2 bg-green-500 rounded-full"></div>}
                  </div>
                </button>
              </li>
            );
          })}
          {proposedNews.map(([name]) => {
            const prop = fileProposals[name];
            const isDirty = name === selectedFile?.name && (!fileContents[name] || editorContent !== fileContents[name]);
            return (
              <li key={`p-${name}`}>
                <button
                  onClick={() => openFile(name)}
                  className={`w-full text-left px-2 py-1 rounded-md text-white hover:bg-slate-50 ${selectedFile?.name === name ? "bg-slate-100 font-medium" : ""}`}
                >
                  <div className="flex items-center justify-between">
                    <span className="italic text-blue-600">{name} (proposed){isDirty ? " *" : ""}</span>
                    <div className="w-2 h-2 bg-yellow-500 rounded-full"></div>
                  </div>
                </button>
              </li>
            );
          })}
        </ul>
      </div>

      <div className="mt-auto p-3 border-t">
        <div className="text-xs text-slate-500">App Status: {loading ? "loading" : "ready"}</div>
      </div>
    </Card>
  );
}

function DBViewer() {
  const { dbSchema, setDbSchema, setConsoleOutput, setFileProposal, setShowDiffForFile} = useStore();

  useEffect(() => {
    async function loadSchema() {
      try {
        const tres = await fetch(`${BASE_URL}/tables`);
        if (!tres.ok) throw new Error("no tables");
        const tables = await tres.json();
        const tablesWithCols = await Promise.all(
          tables.map(async (t) => {
            const qres = await fetch(`${BASE_URL}/query`, {
              method: "POST",
              headers: { "Content-Type": "application/json" },
              body: JSON.stringify({ query: `PRAGMA table_info(${t})` }),
            });
            const colsData = await qres.json();
            console.log(colsData)
            const columns = colsData.map((c) => `${c.name} ${c.type}`);
            return { name: t, columns };
          })
        );
        setDbSchema({ tables: tables.map((t) => ({ name: t, columns: [] })) });
      } catch (e) {
      }
    }
    loadSchema();
  }, [setDbSchema]);

  async function inspectTable(table) {
    try {
      const res = await fetch(`${BASE_URL}/query`, {
        method: "POST",
        headers: { "Content-Type": "application/json" },
        body: JSON.stringify({ query: `SELECT * FROM ${table} LIMIT 50;` }),
      });
      const data = await res.json();
      setConsoleOutput(JSON.stringify(data, null, 2));
    } catch (e) {
      setConsoleOutput(`Failed to query table ${table}: ${e.message}`);
    }
  }

  return (
    <Card className="h-full">
      <CardContent>
        <div className="flex items-center justify-between">
          <h3 className="text-sm font-semibold">DB Viewer</h3>
          <div className="text-xs text-slate-500">SQLite</div>
        </div>

        <div className="mt-3 space-y-3 overflow-auto" style={{ maxHeight: "40vh" }}>
          {dbSchema.tables.map((t) => (
            <div key={t.name} className="p-2 rounded border hover:bg-slate-50">
              <div className="flex justify-between items-center">
                <div>
                  <div className="font-medium">{t.name}</div>
                  <div className="text-xs text-slate-500">{t.columns.length} columns</div>
                </div>
                <div className="flex gap-2">
                  <IconButton onClick={() => inspectTable(t.name)}>Preview</IconButton>
                </div>
              </div>
            </div>
          ))}
        </div>
      </CardContent>
    </Card>
  );
}

function WSListener() {
  const { setFileProposal, setShowDiffForFile, selectedFile, pushChat } = useStore();
  const [isWaiting, setIsWaiting] = useState(true);

  useEffect(() => {
    const socket = io("http://localhost:7860");

    socket.on("connect", () => {
      console.log("Connected to WebSocket server");
      setIsWaiting(false);
    });

    socket.on("event", (data) => {
      console.log("WS message:", data);
      if (data.kind === "code_change") {
        setFileProposal(data.filename || selectedFile?.name, {
          code: data.proposedCode,
          isNew: data.isNew,
          suggestedName: data.newFileName,
        });
        setShowDiffForFile(data.filename || selectedFile?.name);
      }
    });

    socket.on("log", (data) => {
      console.log("WS message:", data);
      pushChat({ role: "agent", text: data.msg.replace(/\u001b\[[0-9;]*m/g, '') });
    });

    socket.on("disconnect", () => {
      console.log("Disconnected from WebSocket server");
      setIsWaiting(true);
    });

    return () => {
      socket.disconnect();
    };
  }, []);

  return isWaiting ? <div>Waiting for agent response...</div> : null;
}

function AgentChat() {
  const { chat, pushChat, setFileProposal, setShowDiffForFile, selectedFile, openFile } = useStore();
  const [input, setInput] = useState("");
  const [sending, setSending] = useState(false);

  async function sendMessage() {
    if (!input.trim() || sending) return;
    pushChat({ role: "user", text: input });
    setInput("");
    setSending(true);

    try {
      const res = await fetch(`${BASE_URL}/run_crew`, {
        method: "POST",
        headers: { "Content-Type": "application/json" },
        body: JSON.stringify({ instructions: input, file: selectedFile?.name }),
      });
      const data = await res.json();
      pushChat({ role: "agent", text: data.status || "Crew is Working..." });
    } catch (e) {
      pushChat({ role: "agent", text: `Error: ${e.message}` });
    } finally {
      setSending(false);
    }
  }

  return (
   <Card className="flex flex-col h-full">
    <CardContent className="flex-1 flex flex-col">
      <div className="flex items-center justify-between mb-2">
        <h3 className="text-sm font-semibold">Agent</h3>
      </div>

      <div className="h-64 overflow-y-auto border rounded p-2 bg-slate-50">
        {chat.map((m, i) => (
          <div key={i} className={`mb-2 ${m.role === "user" ? "text-right" : ""}`}>
            <div
              className={`${
                m.role === "user" ? "inline-block bg-blue-600 text-white" : "inline-block bg-white text-slate-800"
              } px-3 py-1 rounded-lg`}
            >
              {m.text}
              {m.proposedCode && (
                <div className="mt-2 p-2 bg-gray-100 rounded">
                  <pre className="text-xs overflow-auto">{m.proposedCode}</pre>
                </div>
              )}
            </div>
          </div>
        ))}
      </div>

      <div className="mt-2 flex gap-2">
        <input
          className="flex-1 border rounded px-3 py-2"
          value={input}
          onChange={(e) => setInput(e.target.value)}
          placeholder="Instruct SQL agent"
          onKeyDown={(e) => {
            if (e.key === "Enter") sendMessage();
          }}
        />
        <IconButton onClick={sendMessage} className="bg-blue-600 text-white">
          {sending ? "..." : "Send"}
        </IconButton>
      </div>
    </CardContent>
  </Card>
  );
}

function OutputConsole() {
  const { consoleOutput } = useStore();
  return (
    <Card>
      <CardContent className="p-0">
        <div className="bg-black text-green-300 font-mono text-sm p-3 h-40 overflow-auto">
          {consoleOutput || "Console output will appear here."}
        </div>
      </CardContent>
    </Card>
  );
}

function SQLMonacoEditor() {
  const {
    editorContent,
    setEditorContent,
    selectedFile,
    showDiffForFile,
    fileProposals,
    fileContents,
    acceptProposal,
    rejectProposal,
  } = useStore();
  const proposal = fileProposals[selectedFile?.name];
  const showDiff = !!proposal?.active && showDiffForFile === selectedFile?.name;
  const proposedContent = proposal?.code || "";
  const isDirty = selectedFile?.name && (!fileContents[selectedFile.name] || editorContent !== fileContents[selectedFile.name]);

  const saveFile = async () => {
    let fileName = selectedFile?.name;
    if (!fileName) {
      const name = prompt("Save as (must end with .sql):");
      if (!name || !name.endsWith(".sql")) {
        alert("Invalid filename");
        return;
      }
      fileName = name;
      try {
        const res = await fetch(`${BASE_URL}/files`, {
          method: "POST",
          headers: { "Content-Type": "application/json" },
          body: JSON.stringify({ filename: name, content: editorContent }),
        });
        if (!res.ok) throw new Error("Create failed");
        const currentFiles = useStore.getState().files;
        useStore.setState({ files: [...currentFiles, { name }] });
        useStore.setState({ selectedFile: { name }, editorContent, fileContents: { ...useStore.getState().fileContents, [name]: editorContent } });
        alert("File created and saved");
        return;
      } catch (e) {
        alert("Failed to create file: " + e.message);
        return;
      }
    }
    try {
      await fetch(`${BASE_URL}/files/${encodeURIComponent(fileName)}`, {
        method: "PUT",
        headers: { "Content-Type": "application/json" },
        body: JSON.stringify({ content: editorContent }),
      });
      useStore.setState({ fileContents: { ...useStore.getState().fileContents, [fileName]: editorContent } });
      alert("Saved");
    } catch (e) {
      alert("Save failed: " + e.message);
    }
  };

  const runQuery = async () => {
    let fileName = selectedFile?.name;
    if (!fileName) {
      const name = prompt("Save query as (must end with .sql):");
      if (!name || !name.endsWith(".sql")) {
        alert("Invalid filename");
        return;
      }
      fileName = name;
      try {
        const res = await fetch(`${BASE_URL}/files`, {
          method: "POST",
          headers: { "Content-Type": "application/json" },
          body: JSON.stringify({ filename: name, content: editorContent }),
        });
        if (!res.ok) throw new Error("Create failed");
        const currentFiles = useStore.getState().files;
        useStore.setState({ files: [...currentFiles, { name }], selectedFile: { name }, editorContent, fileContents: { ...useStore.getState().fileContents, [name]: editorContent } });
      } catch (e) {
        alert("Failed to create file: " + e.message);
        return;
      }
    }
    try {
      const res = await fetch(`${BASE_URL}/query`, {
        method: "POST",
        headers: { "Content-Type": "application/json" },
        body: JSON.stringify({ query: editorContent }),
      });
      if (!res.ok) {
        let err;
        try {
          err = await res.json();
        } catch {
          err = "Query failed";
        }
        throw new Error(err.error || JSON.stringify(err) || err);
      }
      const data = await res.json();
      useStore.setState({ consoleOutput: JSON.stringify(data, null, 2) });
    } catch (e) {
      useStore.setState({ consoleOutput: "Run failed: " + e.message });
    }
  };

  const commonOptions = { minimap: { enabled: false }, fontSize: 13, wordWrap: "on" };
  const editorKey = `${selectedFile?.name}-${showDiff ? "diff" : "normal"}`;

  return (
    <Card className="h-full flex flex-col">
      <CardContent className="flex-1 flex flex-col p-0">
        <div className="p-3 border-b flex items-center justify-between">
          <div>
            <div className="text-sm font-semibold">{selectedFile?.name || "No file selected"}{isDirty ? " *" : ""}</div>
            <div className="text-xs text-slate-500">SQL Editor {showDiff ? "(Diff View)" : ""}</div>
          </div>
          <div className="flex gap-2">
            {showDiff ? (
              <>
                <IconButton onClick={acceptProposal}>Accept</IconButton>
                <IconButton onClick={rejectProposal} className="bg-red-500">
                  Reject
                </IconButton>
              </>
            ) : (
              <>
                <IconButton onClick={saveFile}>Save</IconButton>
                <IconButton onClick={runQuery}>Run</IconButton>
              </>
            )}
          </div>
        </div>

        <div key={editorKey} className="flex-1">
          {showDiff ? (
            <DiffEditor
              height="100%"
              language="sql"
              original={editorContent}
              modified={proposedContent}
              options={{ ...commonOptions, renderSideBySide: false, diffCodeLens: false }}
            />
          ) : (
            <Editor
              height="100%"
              defaultLanguage="sql"
              value={editorContent}
              onChange={(v) => setEditorContent(v || "")}
              options={commonOptions}
            />
          )}
        </div>
      </CardContent>
    </Card>
  );
}

export default function App() {
  return (
    <div className="h-screen w-screen bg-slate-50 p-4 box-border">
      <div className="grid grid-cols-12 gap-4 h-full">
        <div className="col-span-2 flex flex-col gap-4 h-full">
          <FileExplorer />
          <div className="h-1/3"></div>
        </div>

        <div className="col-span-7 flex flex-col gap-4 h-full">
          <SQLMonacoEditor />
        </div>

        <div className="col-span-3 flex flex-col gap-4 h-full">
          <div className="flex-1">
            <AgentChat />
          </div>
          <div className="h-1/3">
            <DBViewer />
          </div>
          <div className="h-1/3">
            <OutputConsole />
          </div>
        </div>
      </div>
      <WSListener />
    </div>
  );
}