from textual.app import App, ComposeResult
from textual.containers import Container, Vertical, Horizontal
from textual.widgets import Header, Footer, Static, Button, Input, Select, Label, Checkbox, ProgressBar
from textual.screen import Screen
from textual.binding import Binding
from pathlib import Path
import threading
import time

# archinstall imports
from archinstall.lib.args import ArchConfig, Arguments, arch_config_handler
from archinstall.lib.authentication.authentication_handler import AuthenticationHandler
from archinstall.lib.applications.application_handler import ApplicationHandler
from archinstall.lib.mirrors import MirrorListHandler
from archinstall.lib.disk.device_handler import device_handler
from archinstall.lib.disk.filesystem import FilesystemHandler
from archinstall.scripts.guided import perform_installation
from archinstall.lib.output import logger

# Acreetion-specific imports
from archinstall.acreetion.config_mapping import gui_config_to_arch_config

class WelcomeScreen(Static):
    """
    Initial welcome screen for the AcreetionOS TUI installer.
    Displays the project title and a start button.
    """
    def compose(self) -> ComposeResult:
        yield Label("▲ AcreetionOS Installer", id="title")
        yield Label("Welcome to the modern Arch installation experience.", id="subtitle")
        yield Button("Start Installation", variant="primary", id="start-btn")

class DiskSetupScreen(Static):
    """
    Disk selection and filesystem configuration screen.
    Allows the user to choose a target block device and preferred filesystem.
    """
    def compose(self) -> ComposeResult:
        yield Label("Select Target Disk", classes="section-title")
        disks = []
        for disk in device_handler.devices:
            disks.append((f"{disk.device_info.path} ({disk.device_info.total_size.format()})", str(disk.device_info.path)))
        
        yield Select(disks, id="disk-select", prompt="Choose a disk...")
        yield Label("Filesystem", classes="section-title")
        yield Select([("Btrfs (Recommended)", "btrfs"), ("ext4", "ext4")], id="fs-select", value="btrfs")
        
        yield Horizontal(
            Button("Back", id="back-btn"),
            Button("Next", variant="primary", id="next-btn"),
            classes="nav-buttons"
        )

class UserConfigScreen(Static):
    """
    User account and system identity configuration screen.
    Collects hostname, username, password, and sudo preferences.
    """
    def compose(self) -> ComposeResult:
        yield Label("User Configuration", classes="section-title")
        yield Label("Hostname")
        yield Input(placeholder="acreetion-pc", id="hostname", value="acreetion-pc")
        yield Label("Username")
        yield Input(placeholder="username", id="username")
        yield Label("Password")
        yield Input(placeholder="password", id="password", password=True)
        yield Checkbox("Make administrator (sudo)", value=True, id="sudo-checkbox")
        
        yield Horizontal(
            Button("Back", id="back-btn"),
            Button("Next", variant="primary", id="next-btn"),
            classes="nav-buttons"
        )

class ProgressScreen(Static):
    """
    Installation progress visualization screen.
    Displays a progress bar, status text, and live logs from the installer backend.
    """
    def compose(self) -> ComposeResult:
        yield Label("Installing AcreetionOS...", id="install-title")
        yield ProgressBar(total=100, show_eta=True, id="install-progress")
        yield Static("Initializing...", id="status-text")
        yield Static("", id="log-output")

class AcreetionTUI(App):
    """
    Main Terminal User Interface application for AcreetionOS.
    Manages the multi-step installation wizard and handles the background installation process.
    """
    CSS = """
    Screen {
        align: center middle;
        background: #0f172a;
    }

    Container {
        width: 60;
        height: auto;
        background: #1e293b;
        border: solid #38bdf8;
        padding: 2 4;
    }

    #title {
        text-align: center;
        width: 100%;
        color: #38bdf8;
        text-style: bold;
        margin-bottom: 1;
        font-size: 150%;
    }

    #subtitle {
        text-align: center;
        width: 100%;
        color: #94a3b8;
        margin-bottom: 2;
    }

    .section-title {
        text-style: bold;
        color: #f1f5f9;
        margin-top: 1;
        margin-bottom: 1;
    }

    Label {
        color: #cbd5e1;
        margin-top: 1;
    }

    Input {
        margin-bottom: 1;
        border: tall #334155;
    }

    Select {
        margin-bottom: 1;
    }

    .nav-buttons {
        margin-top: 2;
        height: 3;
        align: center middle;
    }

    #start-btn {
        width: 100%;
        margin-top: 2;
    }

    #back-btn {
        margin-right: 2;
    }

    #install-progress {
        margin: 2 0;
    }

    #status-text {
        text-align: center;
        width: 100%;
        color: #38bdf8;
    }

    #log-output {
        height: 5;
        border: sunken #334155;
        margin-top: 1;
        padding: 0 1;
        color: #94a3b8;
        font-size: 80%;
    }
    """

    BINDINGS = [
        Binding("q", "quit", "Quit"),
    ]

    def __init__(self):
        super().__init__()
        self.install_config = {
            "disk": None,
            "filesystem": "btrfs",
            "hostname": "acreetion-pc",
            "username": "",
            "password": "",
            "sudo": True
        }
        self.current_step = 0
        self.steps = [WelcomeScreen(), DiskSetupScreen(), UserConfigScreen(), ProgressScreen()]

    def compose(self) -> ComposeResult:
        yield Header()
        with Container():
            yield self.steps[0]
        yield Footer()

    def on_button_pressed(self, event: Button.Pressed) -> None:
        if event.button.id == "start-btn" or event.button.id == "next-btn":
            if self.current_step == 1:
                disk = self.query_one("#disk-select", Select).value
                if not disk:
                    return
                self.install_config["disk"] = disk
                self.install_config["filesystem"] = self.query_one("#fs-select", Select).value
            
            elif self.current_step == 2:
                self.install_config["hostname"] = self.query_one("#hostname", Input).value
                self.install_config["username"] = self.query_one("#username", Input).value
                self.install_config["password"] = self.query_one("#password", Input).value
                self.install_config["sudo"] = self.query_one("#sudo-checkbox", Checkbox).value
                
                if not self.install_config["username"] or not self.install_config["password"]:
                    return

            self.next_step()
            
            if self.current_step == 3:
                self.start_installation()

        elif event.button.id == "back-btn":
            self.prev_step()

    def next_step(self):
        self.current_step += 1
        self.update_view()

    def prev_step(self):
        self.current_step -= 1
        self.update_view()

    def update_view(self):
        container = self.query_one(Container)
        container.remove_children()
        container.mount(self.steps[self.current_step])

    def start_installation(self):
        arch_config = gui_config_to_arch_config(self.install_config)
        thread = threading.Thread(target=self.run_installation, args=(arch_config,))
        thread.start()

    def run_installation(self, arch_config):
        def log_to_tui(message, *args, **kwargs):
            self.call_from_thread(self.update_log, str(message))
            
        logger.info = log_to_tui
        logger.debug = log_to_tui
        logger.error = log_to_tui
        logger.warn = log_to_tui

        try:
            mirror_list_handler = MirrorListHandler(offline=False)
            fs_handler = FilesystemHandler(arch_config.disk_config)
            fs_handler.perform_filesystem_operations()
            
            arch_config_handler.config = arch_config
            arch_config_handler.args = Arguments(mountpoint=Path('/mnt'), silent=True)
            
            perform_installation(
                mountpoint=Path('/mnt'),
                mirror_list_handler=mirror_list_handler,
                auth_handler=AuthenticationHandler(),
                application_handler=ApplicationHandler()
            )
            self.call_from_thread(self.update_status, "Installation Complete!")
        except Exception as e:
            self.call_from_thread(self.update_status, f"Error: {str(e)}")

    def update_log(self, message):
        log_widget = self.query_one("#log-output", Static)
        log_widget.update(message)
        
        progress = self.query_one("#install-progress", ProgressBar)
        if "formatting" in message.lower(): progress.progress = 20
        elif "syncing" in message.lower(): progress.progress = 40
        elif "base installation" in message.lower(): progress.progress = 60
        elif "bootloader" in message.lower(): progress.progress = 80
        elif "complete" in message.lower(): progress.progress = 100

    def update_status(self, message):
        self.query_one("#status-text", Static).update(message)

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

if __name__ == "__main__":
    main()
