import subprocess
import os
import json
import tempfile
from datetime import datetime, timedelta
from textual.app import App, ComposeResult
from textual.widgets import Header, Footer, DataTable, Button, Label, Static
from textual.screen import ModalScreen
from textual.containers import Grid
from textual.binding import Binding
from textual import work
from rich.text import Text

UID = os.getuid()
CACHE_FILE = os.path.join(tempfile.gettempdir(), f"scancel_cache_{UID}.json")
CACHE_TTL = timedelta(seconds=1)

def fetch_jobs():
    try:
        if os.path.exists(CACHE_FILE):
            mtime = datetime.fromtimestamp(os.path.getmtime(CACHE_FILE))
            if datetime.now() - mtime < CACHE_TTL:
                with open(CACHE_FILE, "r") as f:
                    return json.load(f)
    except Exception:
        pass

    user = os.getenv("USER") or str(UID)
    try:
        result = subprocess.run(
            ["squeue", "-u", user, "-h", "-o", "%A|%j|%M|%N"],
            text=True, stdout=subprocess.PIPE, stderr=subprocess.DEVNULL
        )
        jobs = []
        for line in result.stdout.strip().split("\n"):
            if line:
                parts = line.split("|")
                if len(parts) >= 4:
                    job_id, name, time, node = parts[0], parts[1], parts[2], parts[3]
                    jobs.append({"id": job_id, "name": name, "time": time, "node": node})
        
        with open(CACHE_FILE, "w") as f:
            json.dump(jobs, f)
        return jobs
    except Exception:
        return []

def run_scancel(job_ids):
    if not job_ids:
        return False
    try:
        subprocess.run(["scancel"] + job_ids, check=True)
        if os.path.exists(CACHE_FILE):
            os.remove(CACHE_FILE)
        return True
    except subprocess.CalledProcessError:
        return False

class ConfirmScreen(ModalScreen[bool]):
    CSS = """
    ConfirmScreen {
        align: center middle;
        background: rgba(0,0,0,0.5);
    }
    #dialog {
        grid-size: 2;
        grid-gutter: 1 2;
        grid-rows: 1fr 3;
        padding: 1 2;
        width: 60;
        height: 12;
        border: thick $background 80%;
        background: $surface;
    }
    #question {
        column-span: 2;
        height: 1fr;
        content-align: center middle;
        text-style: bold;
    }
    Button {
        width: 100%;
    }
    """

    BINDINGS = [
        Binding("escape", "cancel", "Cancel", show=False),
    ]

    def __init__(self, job_ids):
        super().__init__()
        self.job_ids = job_ids

    def compose(self) -> ComposeResult:
        yield Grid(
            Label(f"Are you sure you want to cancel\n{len(self.job_ids)} selected job(s)?", id="question"),
            Button("Yes, Cancel Jobs", variant="error", id="yes"),
            Button("No, Go Back", variant="primary", id="no"),
            id="dialog",
        )

    def on_mount(self) -> None:
        self.query_one("#no").focus()

    def key_left(self) -> None:
        self.focus_previous()
    def key_right(self) -> None:
        self.focus_next()
    def key_up(self) -> None:
        self.focus_previous()
    def key_down(self) -> None:
        self.focus_next()

    def on_button_pressed(self, event: Button.Pressed) -> None:
        if event.button.id == "yes":
            self.dismiss(True)
        else:
            self.dismiss(False)
    
    def action_cancel(self) -> None:
        self.dismiss(False)


class SlurmManagerApp(App):
    CSS = """
    DataTable {
        height: 1fr;
        border: solid green;
    }
    .status-bar {
        dock: top;
        height: 1;
        background: $boost;
        color: $text;
        padding: 0 1;
    }
    """

    BINDINGS = [
        Binding("q", "quit", "Quit"),
        Binding("r", "refresh_data", "Refresh Now"),
        Binding("space", "toggle_selection", "Select/Unselect"),
        Binding("enter", "confirm_cancel", "Execute Cancel"),
    ]

    ICON_CHECKED = "✔"
    ICON_UNCHECKED = "☐"
    ICON_CANCELLING = "⏳"

    def __init__(self):
        super().__init__()
        self.selected_jobs = set()
        self.cancelling_ids = set()

    def compose(self) -> ComposeResult:
        yield Header()
        yield Static("Initializing...", id="status", classes="status-bar")
        yield DataTable(cursor_type="row")
        yield Footer()

    def on_mount(self) -> None:
        table = self.query_one(DataTable)
        screen_width = self.app.size.width
        fixed_overhead = 49
        name_width = max(20, screen_width - fixed_overhead)

        table.add_column("Sel", width=5, key="Select")
        table.add_column("Job ID", width=10)
        table.add_column("Name", width=name_width)
        table.add_column("Time", width=10)
        table.add_column("Node", width=12)
        
        self.action_refresh_data()
        self.set_interval(2.0, self.action_refresh_data)

    @work(exclusive=True)
    async def action_refresh_data(self) -> None:
        status = self.query_one("#status", Static)
        jobs = fetch_jobs()
        table = self.query_one(DataTable)
        
        saved_cursor_index = table.cursor_row
        saved_row_key = None
        try:
            if table.cursor_coordinate:
                saved_row_key = table.coordinate_to_cell_key(table.cursor_coordinate).row_key
        except Exception:
            pass

        table.clear()
        
        current_job_ids = {j["id"] for j in jobs}
        self.cancelling_ids = {jid for jid in self.cancelling_ids if jid in current_job_ids}
        self.selected_jobs = {jid for jid in self.selected_jobs if jid in current_job_ids}

        for job in jobs:
            jid = job["id"]
            
            if jid in self.cancelling_ids:
                marker = Text(self.ICON_CANCELLING, style="dim")
                t_id = Text(jid, style="dim strike")
                t_name = Text(job["name"], style="dim")
                t_time = Text(job["time"], style="dim")
                t_node = Text(job["node"], style="dim")
            else:
                if jid in self.selected_jobs:
                    marker = Text(self.ICON_CHECKED, style="bold green")
                else:
                    marker = self.ICON_UNCHECKED
                
                t_id = jid
                t_name = job["name"]
                t_time = job["time"]
                t_node = job["node"]
            
            table.add_row(
                marker, 
                t_id, 
                t_name, 
                t_time, 
                t_node, 
                key=jid 
            )

        if table.row_count > 0:
            target_index = 0
            
            if saved_row_key:
                try:
                    target_index = table.get_row_index(saved_row_key)
                except Exception:
                    target_index = saved_cursor_index

            target_index = min(target_index, table.row_count - 1)
            target_index = max(0, target_index)
            
            table.move_cursor(row=target_index, animate=False)

        self.update_status()

    def action_toggle_selection(self) -> None:
        table = self.query_one(DataTable)
        try:
            row_key = table.coordinate_to_cell_key(table.cursor_coordinate).row_key
        except Exception:
            return

        if not row_key:
            return

        job_id = row_key.value
        
        if job_id in self.cancelling_ids:
            self.notify("This job is being cancelled.", severity="warning")
            return

        if job_id in self.selected_jobs:
            self.selected_jobs.remove(job_id)
            table.update_cell(row_key, "Select", self.ICON_UNCHECKED)
        else:
            self.selected_jobs.add(job_id)
            table.update_cell(row_key, "Select", Text(self.ICON_CHECKED, style="bold green"))
            
        self.update_status()

    def update_status(self):
        status = self.query_one("#status", Static)
        total = 0
        try:
            total = self.query_one(DataTable).row_count
        except:
            pass
        status.update(f"Active: {total} | Selected: {len(self.selected_jobs)} | Cancelling: {len(self.cancelling_ids)}")

    def on_data_table_row_selected(self, event: DataTable.RowSelected) -> None:
        self.action_confirm_cancel()

    def action_confirm_cancel(self) -> None:
        if not self.selected_jobs:
            self.notify("No jobs selected.", severity="warning")
            return

        def check_confirm(should_cancel: bool) -> None:
            if should_cancel:
                job_list = list(self.selected_jobs)
                
                for jid in job_list:
                    self.cancelling_ids.add(jid)
                
                self.selected_jobs.clear()
                self.action_refresh_data()
                
                success = run_scancel(job_list)
                
                if success:
                    self.notify(f"Cancellation signal sent for {len(job_list)} jobs.", severity="information")
                else:
                    self.notify("Error executing scancel.", severity="error")
            else:
                pass

        self.push_screen(ConfirmScreen(list(self.selected_jobs)), check_confirm)

def main():
    app = SlurmManagerApp()
    app.run()

if __name__ == "__main__":
    main()