import logging
import os
import sys
import dotenv
from pathlib import Path

from textual.app import App, ComposeResult
from textual.screen import Screen
from textual.widgets import (
    Header,
    Tabs,
    Tab,
    Static,
    Footer,
    DirectoryTree,
    Button,
)
from textual.containers import Vertical, Horizontal, Grid
from textual.reactive import reactive

from dotenv import set_key

from utils.configmanager import load_env
from utils.utils import open_directory  # we will actually use this now

from utils.setup import get_base_directory, load_user_config

from flows.otp import otp_activities_by_agent, otp_generate, otp_revoke
from flows.prepPolicy import menu_policy_enforce
from flows.quietAgent import findQuietAgents

from services.agenthandler import findAgents, moveAgents, toggleEnforcement
from services.policyhandler import confirmUpdateAfromE
from services.API import AirlockAPIWrapper

dotenv.load_dotenv()

# ---------------------------------------------------------------------------
# GLOBAL STASH
# ---------------------------------------------------------------------------

_PENDING_JOB = None

logger = logging.getLogger(__name__)


ASCII_ART = r"""
   _____  .__       .__                 __     ___________           .__          
  /  _  \ |__|______|  |   ____   ____ |  | __ \__    ___/___   ____ |  |   ______ 
 /  /_\  \|  \_  __ \  |  /  _ \_/ ___\|  |/ /   |    | /  _ \ /  _ \|  |  /  ___/ 
/    |    \  ||  | \/  |_(  <_> )  \___|    <    |    |(  <_> |  <_> )  |__\___ \  
\____|__  /__||__|  |____/\____/ \___  >__|_ \   |____| \____/ \____/|____/____  >
        \/                           \/     \/                                 \/
        """


# ---------------------------------------------------------------------------
# helper to persist TEXTUAL_THEME to *user* config and mirror to .env
# ---------------------------------------------------------------------------
def _persist_user_theme(theme_name: str) -> None:
    """
    Store the chosen Textual theme in the user's config:
      <base>/config/user_config.json
    and also mirror to <base>/.env so load_env(...) sees it.
    """
    base_dir = get_base_directory()
    config_dir = base_dir / "config"
    user_config_path = config_dir / "user_config.json"
    env_path = base_dir / ".env"

    # ensure dirs / files exist similarly to setup()
    config_dir.mkdir(parents=True, exist_ok=True)
    if not user_config_path.exists():
        # minimal default like your load_user_config does
        user_config_path.write_text('{"URL": "", "LOG_LEVEL": "INFO"}\n', encoding="utf-8")

    # load existing user config
    user_conf = load_user_config(config_dir)
    user_conf["TEXTUAL_THEME"] = theme_name

    # write it back
    user_config_path.write_text(
        # pretty print so it stays human-readable
        __import__("json").dumps(user_conf, indent=4),
        encoding="utf-8",
    )
    logger.debug("Updated user_config.json with TEXTUAL_THEME=%s", theme_name)

    # mirror to .env (like setup.write_config_to_env does)
    env_path.parent.mkdir(parents=True, exist_ok=True)
    if not env_path.exists():
        env_path.touch()
    try:
        set_key(str(env_path), "TEXTUAL_THEME", theme_name)
    except Exception as exc:  # keep going even if .env write fails
        logger.warning("Failed to mirror TEXTUAL_THEME to .env: %s", exc)

    # reload so load_env(...) sees the new value right now
    dotenv.load_dotenv(dotenv_path=env_path, override=True)
    logger.debug("Reloaded .env from %s", env_path)


# ---------------------------------------------------------------------------
# 1) SCREEN
# ---------------------------------------------------------------------------
class MainMenuScreen(Screen):
    current_tab = reactive("")

    BUTTON_DEFS = {
        "find": [
            ("🔍 - Device Search", "find_device_button"),
            ("🔇 - Find Quiet Hosts", "find_quiet_button"),
        ],
        "move": [
            ("✅ - Move to local approval", "move_local_button"),
            ("🔄 - Move to Audit/Enforcement", "move_audit_button"),
            ("🔀 - Move - Other", "move_other_button"),
        ],
        "otp": [
            ("🔐 - Generate OTPs", "otp_generate_button"),
            ("📊 - OTP Activities By Agent", "otp_activities_button"),
            ("❌ - Revoke OTPs", "otp_revoke_button"),
        ],
        "policy": [
            ("🔒 - Prepare Policy For Enforcement", "policy_prep_button"),
            ("🔄 - Update Audit Policies", "policy_audit_update_button"),
        ],
    }

    # textual themes to expose
    THEME_BUTTONS = [
        ("textual-dark", "textual-dark"),
        ("textual-light", "textual-light"),
        ("nord", "nord"),
        ("gruvbox", "gruvbox"),
        ("catppuccin-mocha", "catppuccin-mocha"),
        ("dracula", "dracula"),
        ("tokyo-night", "tokyo-night"),
        ("monokai", "monokai"),
        ("flexoki", "flexoki"),
        ("catppuccin-latte", "catppuccin-latte"),
        ("solarized-light", "solarized-light"),
    ]

    def __init__(self) -> None:
        super().__init__()
        self.extras = load_env("EXTRAS")
        wd = load_env("WORKING_DIR") or os.getcwd()
        if not os.path.isdir(wd):
            wd = os.getcwd()
        self.working_dir = wd


    def _make_buttons_for(self, tab_id: str) -> Vertical:
        defs = self.BUTTON_DEFS.get(tab_id, [])
        buttons = []
        for label, btn_id in defs:
            btn = Button(label, id=btn_id)
            btn.styles.width = "100%"  # Make button span full width of parent
            buttons.append(btn)
        return Vertical(*buttons)




    def compose(self) -> ComposeResult:
        yield Header()
        yield Static(ASCII_ART, id="logo")

        tabs = [
            Tab("Device Search", id="find"),
            Tab("Move Agent", id="move"),
            Tab("OTP", id="otp"),
            Tab("Directory", id="dir"),
            Tab("Settings", id="settings"),
        ]

        if self.extras == "POLICYPREP":
            tabs.insert(3, Tab("Policy Prep", id="policy"))

        yield Tabs(*tabs, id="tabs")
        yield Vertical(id="content")
        yield Footer()

    def on_mount(self) -> None:
        self.switch_tab("find")

    # focus helpers
    def _get_content_buttons(self) -> list[Button]:
        content = self.query_one("#content", Vertical)
        return list(content.query(Button))

    def _focus_first_button(self) -> None:
        buttons = self._get_content_buttons()
        if buttons:
            buttons[0].focus()

    def _focus_tabs(self) -> None:
        tabs = self.query_one("#tabs", Tabs)
        tabs.focus()

    def _focus_nearby_button(self, direction: int) -> None:
        buttons = self._get_content_buttons()
        if not buttons:
            return

        try:
            current = next(i for i, b in enumerate(buttons) if b.has_focus)
        except StopIteration:
            if direction > 0:
                buttons[0].focus()
            else:
                buttons[-1].focus()
            return

        if direction < 0 and current == 0:
            self._focus_tabs()
            return

        new_index = current + direction
        if 0 <= new_index < len(buttons):
            buttons[new_index].focus()
    
    
    def switch_tab(self, tab_id: str) -> None:
        self.current_tab = tab_id
        content = self.query_one("#content", Vertical)
        content.remove_children()

        if tab_id in self.BUTTON_DEFS:
            content.mount(self._make_buttons_for(tab_id))
            self.call_later(self._focus_first_button)
        elif tab_id == "dir":
            content.mount(DirectoryTree(self.working_dir, id="dir_tree"))
        elif tab_id == "settings":
            # Create and mount the horizontal container
            horizontal_container = Horizontal(id="settings_grid")
            horizontal_container.styles.layout = "horizontal"
            horizontal_container.styles.height = "auto"
            content.mount(Static("Theme Options"))
            content.mount(horizontal_container)  # Mount the horizontal container first

            # Create 3 columns
            for i in range(1):
                column = Vertical()
                column.styles.width = "1fr"
                column.styles.height = "auto"
                horizontal_container.mount(column)  # Mount each column

                for j in range(i, len(self.THEME_BUTTONS), 1):
                    if j < len(self.THEME_BUTTONS):
                        label, btn_id = self.THEME_BUTTONS[j]
                        button = Button(label, id=f"set_theme_{btn_id}", compact=True)
                        #button.styles.width = "100%"
                        column.mount(button)  # Mount each button

        else:
            content.mount(Static(f"Unknown tab: {tab_id}"))

    def on_tabs_tab_activated(self, event: Tabs.TabActivated) -> None:
        self.switch_tab(event.tab.id)

    def on_key(self, event) -> None:
        key = event.key
        logger.debug("KEY: %r", key)

        if self.current_tab == "dir":
            return

        if key in ("down", "j"):
            self._focus_nearby_button(+1)
            event.stop()
            return
        if key in ("up", "k"):
            self._focus_nearby_button(-1)
            event.stop()
            return

        if key in ("left", "right"):
            tabs = self.query_one("#tabs", Tabs)
            if not tabs.has_focus:
                tabs.focus()
                if key == "left":
                    tabs.action_previous_tab()
                else:
                    tabs.action_next_tab()
                event.stop()
                return
            return

    def on_directory_tree_file_selected(self, event: DirectoryTree.FileSelected) -> None:
        path = event.path
        logger.debug("Directory file selected: %s", path)
        try:
            open_directory(str(path))
        except Exception as exc:
            logger.error("Failed to open %s: %s", path, exc)
            self.app.bell()

    def on_button_pressed(self, event: Button.Pressed) -> None:
        global _PENDING_JOB
        button_id = event.button.id
        logger.debug("Button pressed: %s", button_id)

        # theme selection → user config
        if button_id.startswith("set_theme_"):
            theme_name = button_id.replace("set_theme_", "")
            _persist_user_theme(theme_name)
            _PENDING_JOB = ("restart",)
            self.app.exit()
            return

        match button_id:
            case "find_device_button":
                _PENDING_JOB = ("legacy", findAgents, (self.app.api, False), {})
            case "find_quiet_button":
                _PENDING_JOB = ("legacy", findQuietAgents, (self.app.api,), {})
            case "move_local_button":
                _PENDING_JOB = (
                    "legacy",
                    print,
                    ("Move to local approval (placeholder)",),
                    {},
                )
            case "move_audit_button":
                _PENDING_JOB = ("legacy", toggleEnforcement, (self.app.api,), {})
            case "move_other_button":
                _PENDING_JOB = ("legacy", moveAgents, (self.app.api,), {})
            case "otp_generate_button":
                _PENDING_JOB = ("legacy", otp_generate, (self.app.api,), {})
            case "otp_activities_button":
                _PENDING_JOB = ("legacy", otp_activities_by_agent, (self.app.api,), {})
            case "otp_revoke_button":
                _PENDING_JOB = ("legacy", otp_revoke, (self.app.api,), {})
            case "policy_prep_button":
                _PENDING_JOB = ("legacy", menu_policy_enforce, (self.app.api,), {})
            case "policy_audit_update_button":
                _PENDING_JOB = ("legacy", confirmUpdateAfromE, (self.app.api,), {})
            case _:
                self.app.bell()
                logger.warning("Unknown button pressed: %s", button_id)
                return

        logger.debug("Set _PENDING_JOB = %r", _PENDING_JOB)
        self.app.exit()






# ---------------------------------------------------------------------------
# 2) APP
# ---------------------------------------------------------------------------
class AirlockTools(App):
    CSS = """
    #logo {
        width: 100%;
        content-align: center middle;
        text-align: center;
    }
    """

    BINDINGS = [
        ("q", "quit", "Quit"),
        ("d", "open_dir", "Open Directory"),
    ]

    def __init__(self, api: AirlockAPIWrapper):
        self._textual_theme = load_env("TEXTUAL_THEME") or "nord"
        super().__init__()
        self.api = api
        wd = load_env("WORKING_DIR") or os.getcwd()
        if not os.path.isdir(wd):
            wd = os.getcwd()
        self.working_dir = wd

    def on_mount(self) -> None:
        self.theme = self._textual_theme
        self.push_screen(MainMenuScreen())

    def action_quit(self) -> None:
        global _PENDING_JOB
        _PENDING_JOB = None
        self.exit()

    def action_open_dir(self) -> None:
        screen = self.screen_stack[-1]
        if isinstance(screen, MainMenuScreen):
            if screen.current_tab != "dir":
                screen.switch_tab("dir")



# ---------------------------------------------------------------------------
# 3) TERMINAL + LEGACY
# ---------------------------------------------------------------------------
def _restore_terminal_for_legacy() -> None:
    sys.stdout.write("\033[?1049l")
    sys.stdout.write("\033[?25h")
    sys.stdout.write("\033[0m")
    sys.stdout.write("\033[?1000l\033[?1002l\033[?1003l\033[?1006l")
    sys.stdout.write("\033[2J\033[H")
    sys.stdout.flush()

    if os.name == "nt":
        try:
            import ctypes
            kernel32 = ctypes.windll.kernel32
            handle = kernel32.GetStdHandle(-11)
            mode = ctypes.c_ulong()
            if kernel32.GetConsoleMode(handle, ctypes.byref(mode)):
                kernel32.SetConsoleMode(handle, mode.value | 0x0004)
        except Exception as exc:
            logger.debug("VT enable on Windows failed: %s", exc)


def _run_legacy_job(func, args, kwargs) -> None:
    logger.debug("Running legacy job: %s", getattr(func, "__name__", func))
    _restore_terminal_for_legacy()

    try:
        func(*args, **kwargs)
    finally:
        try:
            input("\nPress Enter to return to the UI...")
        except EOFError:
            pass


# ---------------------------------------------------------------------------
# 4) PUBLIC ENTRYPOINT
# ---------------------------------------------------------------------------
def run_AirlockTools(api: AirlockAPIWrapper) -> None:
    global _PENDING_JOB

    while True:
        base_dir = get_base_directory()
        env_path = base_dir / ".env"
        dotenv.load_dotenv(dotenv_path=env_path, override=True)

        _PENDING_JOB = None
        app = AirlockTools(api)

        try:
            app.run()
        except SystemExit as exc:
            logger.debug("Caught SystemExit from Textual: %s", exc)

        job = _PENDING_JOB
        logger.debug("After app.run(), _PENDING_JOB = %r", job)

        if not job:
            break

        if job[0] == "legacy":
            _, func, args, kwargs = job
            _run_legacy_job(func, args, kwargs)
            continue

        if job[0] == "restart":
            # just loop again; fresh .env was already loaded at the top
            continue

        break


# ---------------------------------------------------------------------------
# 5) DEV
# ---------------------------------------------------------------------------
if __name__ == "__main__":
    api = AirlockAPIWrapper()
    run_AirlockTools(api)