import logging
import os
import sys

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

from flows.otp import otp_activities_by_agent, otp_revoke
from flows.prepPolicy import menu_policy_enforce
from models.agent import Agent
from models.policy import Policy
from screens.moveagentworkflowscreen import MoveAgentWorkflowScreen
from screens.otpworkflowscreen import OTPWorkflowScreen
from services.API import AirlockAPIWrapper
from services.policyhandler import confirmUpdateAfromE
from utils.configmanager import load_env
from utils.setup import get_base_directory, load_user_config
from utils.utils import open_directory
from widgets.agentmoveoperations import AgentMoveOperations
from widgets.amber_terminal_theme import get_amber_terminal_theme
from widgets.multiagentselector import MultiAgentSelector
from widgets.OTP_generate import OTPGenerator
from widgets.policytreewidget import PolicyTreeWidget
from widgets.resultsdisplay import ResultsDisplay
from widgets.retro_terminal_theme import get_retro_terminal_theme
from widgets.themeselector import ThemeSelector

dotenv.load_dotenv()

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

_PENDING_JOB = None

logger = logging.getLogger(__name__)


# ---------------------------------------------------------------------------
# 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 = {
        "agent_actions": [
            (
                "🖥️ - Find, Move, or Generate OTP for Agents",
                "move_agent_workflow_button",
            ),
            ("🔇 - Find Quiet Hosts", "find_quiet_button"),
        ],
        "policy": [
            ("🔒 - Prepare Policy For Enforcement", "policy_prep_button"),
            ("🔄 - Update Audit Policies", "policy_audit_update_button"),
            ("📊 - OTP Activities By Agent", "otp_activities_button"),
            ("❌ - Revoke OTPs", "otp_revoke_button"),
        ],
    }

    def __init__(self, api: AirlockAPIWrapper) -> None:
        super().__init__()
        self.api = api
        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%"
            buttons.append(btn)
        return Vertical(*buttons)

    def compose(self) -> ComposeResult:
        yield Header(show_clock=True, icon="⚙")

        tabs = [
            Tab("Tree View", id="p_tree"),
            Tab("Agents", id="agent_actions"),
            Tab("Directory", id="dir"),
            Tab("Settings", id="settings"),
        ]

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

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

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

    # 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 == "p_tree":
            content.mount(PolicyTreeWidget(self.app.policies, self.app.devices))
        elif tab_id == "settings":
            content.mount(ThemeSelector())
        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_multi_agent_selector_agents_selected(
        self, message: MultiAgentSelector.AgentsSelected
    ) -> None:
        """Handle selected agents from AgentSelector."""
        global _PENDING_JOB
        selected_agents = message.selected_agents
        logger.info("Selected agents: %s", selected_agents)
        # TODO: Implement actual handling of selected agents
        _PENDING_JOB = ("multi_agent_action", selected_agents)
        self.app.exit()

    def on_theme_selector_theme_selected(
        self, message: ThemeSelector.ThemeSelected
    ) -> None:
        """Handle theme selection from ThemeSelector."""
        global _PENDING_JOB
        _persist_user_theme(message.theme_name)
        _PENDING_JOB = ("restart",)
        self.app.exit()

    def on_otp_generator_otp_info(self, message: OTPGenerator.OTPInfo) -> None:
        """Handle OTP generation request from the workflow."""
        global _PENDING_JOB

        # Log what we received
        logger.info(
            "OTP Generation requested: %d devices, requestor=%s, reason=%s, duration=%d",
            len(message.devices),
            message.requestor,
            message.reasoning,
            message.duration,
        )

        # Set up the job to run the OTP generation
        _PENDING_JOB = (
            "otp_workflow",
            message.devices,
            message.requestor,
            message.reasoning,
            message.duration,
        )

        self.app.exit()

    def on_agent_move_operations_operation_complete(
        self, message: AgentMoveOperations.OperationComplete
    ) -> None:
        """Handle completion of agent move operation - show results."""
        logger.info(
            "Agent move operation completed: %s, %d successful, %d unsuccessful",
            message.operation,
            len(message.successful),
            len(message.unsuccessful),
        )

        # Format results for display
        successful_text = "\n".join(
            [f"{agent.hostname}" for agent, _ in message.successful]
        )
        unsuccessful_text = "\n".join(
            [f"{agent.hostname}: {error}" for agent, error in message.unsuccessful]
        )

        # Remove the operations widget
        try:
            ops_widget = self.query_one(AgentMoveOperations)
            ops_widget.remove()
        except Exception:
            pass

        # Show results
        self.query_one("#content", Vertical).mount(
            ResultsDisplay(message.operation, successful_text, unsuccessful_text)
        )

    def on_results_display_go_back(self, message: ResultsDisplay.GoBack) -> None:
        """Handle back button from results display."""
        try:
            results_widget = self.query_one(ResultsDisplay)
            results_widget.remove()
        except Exception:
            pass
        # Return to main menu
        self.app.pop_screen()

    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)

        match button_id:
            case "move_agent_workflow_button":
                # Push Move Agent workflow screen
                self.app.push_screen(MoveAgentWorkflowScreen(self.app.devices))
                event.stop()
                return  # Don't exit the app
            case "otp_generate_button":
                # NEW: Push OTP workflow screen instead of legacy function
                self.app.push_screen(OTPWorkflowScreen(self.app.devices))
                event.stop()
                return  # Don't exit the app
            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 Loxide(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

        # Add error handling for API calls
        try:
            self.policies = [
                Policy(**row.to_dict()) for _, row in api.policy_find_all().iterrows()
            ]
            self.devices = [
                Agent(**row.to_dict()) for _, row in api.agent_find_all().iterrows()
            ]

            # Enrich agents with policy information
            if self.policies and self.devices:
                for agent in self.devices:
                    agent.enrich_with_policies(self.policies)
                logger.debug(
                    f"Enriched {len(self.devices)} agents with policy information"
                )
        except Exception as exc:
            logger.error("Failed to load policies/devices: %s", exc)
            self.policies = None
            self.devices = None

    def on_mount(self, api: AirlockAPIWrapper) -> None:
        self.register_theme(get_retro_terminal_theme())
        self.register_theme(get_amber_terminal_theme())
        self.theme = self._textual_theme
        self.push_screen(MainMenuScreen(api))

    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_Loxide(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 = Loxide(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

        if job[0] == "multi_agent_action":
            # Handle multi-agent selection
            logger.info("Multi-agent action with selected agents: %s", job[1])
            continue

        # NEW: Handle OTP workflow
        if job[0] == "otp_workflow":
            _, devices, requestor, reasoning, duration = job

            # Call your OTP generation with the parameters
            def otp_generate_with_params():

                print(f"\n{'='*60}")
                print("OTP GENERATION")
                print(f"{'='*60}")
                print(f"Requestor: {requestor}")
                print(f"Reasoning: {reasoning}")
                print(f"Duration: {duration} minutes")
                print(f"\nGenerating OTPs for {len(devices)} devices:")
                print(f"{'='*60}\n")

                # Call your actual OTP generation function
                # You'll need to adapt otp_generate to accept these parameters
                # For now, this is a placeholder showing the structure
                for device in devices:
                    print(f"Device: {device}")
                    print(f"  Requestor: {requestor}")
                    print(f"  Reason: {reasoning}")
                    print(f"  Duration: {duration} minutes")
                    # TODO: Actually call your API to generate OTP
                    # result = api.generate_otp(device, requestor, reasoning, duration)
                    print()

                print(f"{'='*60}")
                print("OTP Generation Complete!")
                print(f"{'='*60}")

            _run_legacy_job(otp_generate_with_params, (), {})
            continue

        break


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