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 flows.quietAgent import findQuietAgents
from models.agent import Agent
from models.policy import Policy
from screens.moveagentworkflowscreen import MoveAgentWorkflowScreen
from screens.otpworkflowscreen import OTPWorkflowScreen
from services.agenthandler import findAgents, moveAgents, toggleEnforcement
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.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 = {
        "find": [
            ("🔍 - Device Search", "find_device_button"),
            ("🔇 - Find Quiet Hosts", "find_quiet_button"),
        ],
        "move": [
            ("🔄 - Move Agent Workflow", "move_agent_workflow_button"),
            ("✅ - 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"),
        ],
    }

    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("Policy Tree", id="p_tree"),
            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 == "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 MultiAgentSelector."""
        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 "find_device_button":
                _PENDING_JOB = ("legacy", findAgents, (self.app.api, False), {})
            case "find_quiet_button":
                _PENDING_JOB = ("legacy", findQuietAgents, (self.app.api,), {})
            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 "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":
                # 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.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)