"""b2a tui – Terminal UI for B2Alpha messaging."""

from __future__ import annotations

import base64
import json
import threading
from datetime import datetime

from textual import work
from textual.app import App, ComposeResult
from textual.binding import Binding
from textual.containers import Horizontal, Vertical
from textual.reactive import reactive
from textual.screen import ModalScreen
from textual.widgets import Footer, Header, Input, Label, ListItem, ListView, Static


# ── Helpers ──────────────────────────────────────────────────────────────────

def _decode_payload(raw: str) -> str:
    try:
        padding = 4 - len(raw) % 4
        if padding != 4:
            raw += "=" * padding
        decoded = base64.urlsafe_b64decode(raw)
        body = json.loads(decoded)
        return body.get("text", decoded.decode(errors="replace"))
    except Exception:
        return "(encrypted)"


def _short_time(ts: str) -> str:
    try:
        dt = datetime.fromisoformat(ts.replace("Z", "+00:00"))
        return dt.strftime("%b %d %H:%M")
    except Exception:
        return ""


def _did_snippet(did: str) -> str:
    """Truncate a DID for display: did:b2a:zXYZ... -> zXYZ..."""
    if did.startswith("did:b2a:"):
        tag = did[8:]
        return tag[:12] + "..." if len(tag) > 12 else tag
    return did[:20]


def _display_name(did: str, names: dict[str, str]) -> str:
    """Format as 'Name (zSnippet...)' or just 'zSnippet...' if no name."""
    name = names.get(did)
    snippet = _did_snippet(did)
    if name:
        return f"{name} ({snippet})"
    return snippet


# ── Styles ───────────────────────────────────────────────────────────────────

CSS = """
Screen {
    layout: horizontal;
}

#sidebar {
    width: 42;
    border-right: solid $surface-lighten-2;
}

#sidebar-title {
    dock: top;
    height: 1;
    padding: 0 1;
    background: $primary;
    color: $text;
    text-style: bold;
}

#convo-list {
    height: 1fr;
}

#convo-list > ListItem {
    height: 3;
    padding: 0 1;
}

#convo-list > ListItem.--highlight {
    background: $primary 30%;
}

#main-panel {
    width: 1fr;
}

#messages-title {
    dock: top;
    height: 1;
    padding: 0 1;
    background: $surface-lighten-1;
    text-style: bold;
}

#messages-scroll {
    height: 1fr;
    padding: 0 1;
    overflow-y: auto;
}

#input-bar {
    dock: bottom;
    height: 3;
    padding: 0 1;
}

#msg-input {
    width: 1fr;
}

.msg-row {
    height: auto;
    margin-bottom: 1;
}

.msg-sender-you {
    color: $success;
    text-style: bold;
}

.msg-sender-them {
    color: $accent;
    text-style: bold;
}

.msg-time {
    color: $text-muted;
}

.convo-partner {
    text-style: bold;
}

.convo-preview {
    color: $text-muted;
}

.convo-unread .convo-partner {
    color: $warning;
}

.empty-state {
    width: 1fr;
    height: 1fr;
    content-align: center middle;
    color: $text-muted;
}

#status-bar {
    dock: bottom;
    height: 1;
    padding: 0 1;
    background: $surface-lighten-1;
    color: $text-muted;
}
"""

SEARCH_CSS = """
PhonebookSearchScreen {
    align: center middle;
}

#search-dialog {
    width: 60;
    max-height: 80%;
    background: $surface;
    border: thick $primary;
    padding: 1 2;
}

#search-title {
    text-style: bold;
    margin-bottom: 1;
}

#search-input {
    margin-bottom: 1;
}

#search-results {
    height: 1fr;
    min-height: 6;
}

#search-results > ListItem {
    height: 3;
    padding: 0 1;
}

.search-name {
    text-style: bold;
}

.search-detail {
    color: $text-muted;
}

#search-hint {
    color: $text-muted;
    height: 1;
    margin-top: 1;
}
"""


# ── Widgets ──────────────────────────────────────────────────────────────────

class ConvoItem(ListItem):
    """A single conversation in the sidebar list."""

    def __init__(self, convo: dict, my_did: str, names: dict[str, str]) -> None:
        super().__init__()
        self.convo = convo
        self.my_did = my_did
        self.names = names
        partner = convo["participant_b_did"] if convo["participant_a_did"] == my_did else convo["participant_a_did"]
        self.partner_did = partner

        last_msg = convo.get("last_message")
        is_unread = last_msg and last_msg.get("sender_did") != my_did
        if is_unread:
            self.add_class("convo-unread")

    def compose(self) -> ComposeResult:
        last_msg = self.convo.get("last_message")
        count = self.convo.get("message_count", 0)
        updated = _short_time(self.convo.get("updated_at", ""))

        preview = ""
        if last_msg:
            text = _decode_payload(last_msg.get("payload", ""))
            preview = text[:30] + "..." if len(text) > 30 else text

        label = _display_name(self.partner_did, self.names)
        yield Label(f"{label}  ({count})  {updated}", classes="convo-partner")
        yield Label(preview or "(no messages)", classes="convo-preview")


class MessageDisplay(Static):
    """A single message bubble."""

    def __init__(self, msg: dict, my_did: str, names: dict[str, str]) -> None:
        sender = msg.get("sender_did", "?")
        is_me = sender == my_did
        text = _decode_payload(msg.get("payload", ""))
        ts = _short_time(msg.get("created_at", ""))

        sender_cls = "msg-sender-you" if is_me else "msg-sender-them"
        sender_label = "you" if is_me else _display_name(sender, names)

        markup = (
            f"[@click=noop]{ts}[/] "
            f"[{sender_cls}]{sender_label}[/{sender_cls}]: "
            f"{text}"
        )
        super().__init__(markup, classes="msg-row")


# ── Phonebook Search ─────────────────────────────────────────────────────────

class SearchResultItem(ListItem):
    """A single phonebook search result."""

    def __init__(self, agent: dict) -> None:
        super().__init__()
        self.agent_data = agent

    def compose(self) -> ComposeResult:
        a = self.agent_data
        name = a.get("display_name", "(unnamed)")
        did = a.get("did", "")
        caps = ", ".join(a.get("capabilities", []))
        type_labels = {0: "unknown", 1: "user", 2: "business"}
        type_label = type_labels.get(a.get("agent_type", 0), "unknown")
        yield Label(f"{name}  [{type_label}]", classes="search-name")
        detail = _did_snippet(did)
        if caps:
            detail += f"  ({caps})"
        yield Label(detail, classes="search-detail")


class PhonebookSearchScreen(ModalScreen[str | None]):
    """Modal overlay for searching the phonebook."""

    CSS = SEARCH_CSS
    BINDINGS = [Binding("escape", "cancel", "Close", priority=True)]

    def __init__(self, token: str) -> None:
        super().__init__()
        self._token = token
        self._results: list[dict] = []

    def compose(self) -> ComposeResult:
        with Vertical(id="search-dialog"):
            yield Label("Phonebook Search", id="search-title")
            yield Input(placeholder="Search by name or capability...", id="search-input")
            yield ListView(id="search-results")
            yield Label("Enter to search · ↑↓ to select · Enter to chat · Esc to close", id="search-hint")

    def on_mount(self) -> None:
        self.query_one("#search-input", Input).focus()

    def on_input_submitted(self, event: Input.Submitted) -> None:
        query = event.value.strip()
        if query:
            self._do_search(query)

    @work(thread=True)
    def _do_search(self, query: str) -> None:
        from .api import search_phonebook
        try:
            result = search_phonebook(query=query, limit=20, token=self._token)
            self._results = result.get("agents", [])
            self.call_from_thread(self._render_results)
        except RuntimeError as e:
            self.call_from_thread(self.notify, f"Search failed: {e}", severity="error")

    def _render_results(self) -> None:
        lv = self.query_one("#search-results", ListView)
        lv.clear()
        if not self._results:
            self.notify("No agents found.", severity="warning")
            return
        for agent in self._results:
            lv.append(SearchResultItem(agent))
        lv.focus()

    def on_list_view_selected(self, event: ListView.Selected) -> None:
        item = event.item
        if isinstance(item, SearchResultItem):
            did = item.agent_data.get("did", "")
            if did:
                self.dismiss(did)

    def action_cancel(self) -> None:
        self.dismiss(None)


# ── Main App ─────────────────────────────────────────────────────────────────

class B2AApp(App):
    TITLE = "B2Alpha"
    CSS = CSS
    BINDINGS = [
        Binding("q", "quit", "Quit", priority=False),
        Binding("r", "refresh", "Refresh", priority=False),
        Binding("n", "new_chat", "New chat", priority=False),
        Binding("s", "search_phonebook", "Search", priority=False),
        Binding("escape", "blur_input", "Back", show=False),
    ]

    selected_convo: reactive[dict | None] = reactive(None)
    my_did: str = ""
    token: str = ""
    conversations: list[dict] = []
    _names: dict[str, str] = {}  # did -> display_name cache
    _stop_realtime = None

    def compose(self) -> ComposeResult:
        yield Header()
        with Horizontal():
            with Vertical(id="sidebar"):
                yield Label(" Conversations", id="sidebar-title")
                yield ListView(id="convo-list")
            with Vertical(id="main-panel"):
                yield Label(" Select a conversation", id="messages-title")
                yield Vertical(id="messages-scroll")
                with Horizontal(id="input-bar"):
                    yield Input(placeholder="Type a message...", id="msg-input")
        yield Label("", id="status-bar")
        yield Footer()

    def on_mount(self) -> None:
        from .auth import get_valid_token
        from .config import get_active_did

        try:
            self.token = get_valid_token()
        except RuntimeError:
            self.notify("Not logged in. Run 'b2a login' first.", severity="error")
            self.exit()
            return

        self.my_did = get_active_did() or ""
        if not self.my_did:
            self.notify("No agent. Run 'b2a register' first.", severity="error")
            self.exit()
            return

        self._set_status(f"  {self.my_did}")
        self._load_inbox()
        self._start_realtime()

    def _set_status(self, text: str) -> None:
        self.query_one("#status-bar", Label).update(text)

    @work(thread=True)
    def _load_inbox(self) -> None:
        from .api import list_conversations, lookup_agent
        try:
            result = list_conversations(self.token, limit=50)
            self.conversations = result.get("conversations", [])

            # Resolve display names for all partner DIDs we haven't seen
            partner_dids = set()
            for c in self.conversations:
                partner = c["participant_b_did"] if c["participant_a_did"] == self.my_did else c["participant_a_did"]
                if partner not in self._names:
                    partner_dids.add(partner)
            for did in partner_dids:
                try:
                    info = lookup_agent(did, token=self.token)
                    name = info.get("agent", {}).get("display_name", "")
                    if name:
                        self._names[did] = name
                except RuntimeError:
                    pass

            self.call_from_thread(self._render_convo_list)
        except RuntimeError as e:
            self.call_from_thread(self.notify, f"Failed to load inbox: {e}", severity="error")

    def _render_convo_list(self) -> None:
        lv = self.query_one("#convo-list", ListView)
        lv.clear()
        for c in self.conversations:
            lv.append(ConvoItem(c, self.my_did, self._names))

    def on_list_view_selected(self, event: ListView.Selected) -> None:
        item = event.item
        if isinstance(item, ConvoItem):
            self.selected_convo = item.convo
            self._load_messages(item.convo["interaction_id"])

    @work(thread=True)
    def _load_messages(self, interaction_id: str) -> None:
        from .api import get_conversation
        try:
            result = get_conversation(self.token, interaction_id)
            conv = result.get("conversation", {})
            if conv:
                self.call_from_thread(self._render_messages, conv)
        except RuntimeError as e:
            self.call_from_thread(self.notify, f"Failed to load: {e}", severity="error")

    def _render_messages(self, conv: dict) -> None:
        partner = conv["participant_b_did"] if conv["participant_a_did"] == self.my_did else conv["participant_a_did"]
        title = self.query_one("#messages-title", Label)
        title.update(f" Chat with {_display_name(partner, self._names)}")

        scroll = self.query_one("#messages-scroll", Vertical)
        scroll.remove_children()

        msgs = conv.get("messages", [])
        for msg in msgs:
            scroll.mount(MessageDisplay(msg, self.my_did, self._names))

        # Scroll to bottom
        scroll.scroll_end(animate=False)

        # Focus the input
        self.query_one("#msg-input", Input).focus()

    def _start_realtime(self) -> None:
        from .realtime import subscribe_to_conversations

        def on_incoming(record: dict) -> None:
            # If we're viewing this conversation, reload it
            convo = self.selected_convo
            if convo and record.get("interaction_id") == convo["interaction_id"]:
                self._load_messages(convo["interaction_id"])
            # Also refresh the sidebar
            self._load_inbox()

        def on_error(err: str) -> None:
            pass  # Silently handle realtime errors

        self._stop_realtime = subscribe_to_conversations(
            self.my_did, self.token, on_incoming, on_error,
        )

    def action_refresh(self) -> None:
        self._load_inbox()
        if self.selected_convo:
            self._load_messages(self.selected_convo["interaction_id"])
        self.notify("Refreshed")

    def action_new_chat(self) -> None:
        """Prompt for a DID and start a new conversation."""
        inp = self.query_one("#msg-input", Input)
        inp.placeholder = "Enter recipient DID, then press Enter..."
        inp.value = ""
        inp.focus()
        self._awaiting_new_chat = True

    def action_search_phonebook(self) -> None:
        """Open phonebook search overlay."""
        def on_dismiss(did: str | None) -> None:
            if did:
                # Cache the name if we got it from the search
                self._start_chat_with(did)

        self.push_screen(PhonebookSearchScreen(self.token), callback=on_dismiss)

    def _start_chat_with(self, did: str) -> None:
        """Start a new chat with the given DID (from search or manual entry)."""
        # Resolve display name
        if did not in self._names:
            try:
                from .api import lookup_agent
                info = lookup_agent(did, token=self.token)
                name = info.get("agent", {}).get("display_name", "")
                if name:
                    self._names[did] = name
            except RuntimeError:
                pass

        self.selected_convo = {
            "interaction_id": None,
            "participant_a_did": self.my_did,
            "participant_b_did": did,
            "message_count": 0,
        }
        title = self.query_one("#messages-title", Label)
        title.update(f" New chat with {_display_name(did, self._names)}")
        scroll = self.query_one("#messages-scroll", Vertical)
        scroll.remove_children()
        self.query_one("#msg-input", Input).focus()
        self.notify(f"Chatting with {_display_name(did, self._names)}. Type your message.")

    _awaiting_new_chat: bool = False

    def on_input_submitted(self, event: Input.Submitted) -> None:
        text = event.value.strip()
        if not text:
            return

        if self._awaiting_new_chat:
            self._awaiting_new_chat = False
            event.input.value = ""
            event.input.placeholder = "Type a message..."
            if text.startswith("did:b2a:"):
                # Create a synthetic convo so we can send to this DID
                self.selected_convo = {
                    "interaction_id": None,
                    "participant_a_did": self.my_did,
                    "participant_b_did": text,
                    "message_count": 0,
                }
                title = self.query_one("#messages-title", Label)
                title.update(f" New chat with {_display_name(text, self._names)}")
                scroll = self.query_one("#messages-scroll", Vertical)
                scroll.remove_children()
                self.notify(f"Chatting with {_display_name(text, self._names)}. Type your message.")
            else:
                self.notify("Invalid DID. Must start with did:b2a:", severity="warning")
            return

        if not self.selected_convo:
            return
        event.input.value = ""
        self._send_message(text)

    @work(thread=True)
    def _send_message(self, text: str) -> None:
        from .crypto import sign_message
        from .api import send_message

        convo = self.selected_convo
        if not convo:
            return

        partner = convo["participant_b_did"] if convo["participant_a_did"] == self.my_did else convo["participant_a_did"]
        interaction_id = convo.get("interaction_id")

        payload_bytes = json.dumps({"text": text}).encode()
        envelope = sign_message(self.my_did, partner, payload_bytes)

        try:
            result = send_message(self.token, self.my_did, partner, envelope, interaction_id=interaction_id)
            new_id = result.get("interaction_id")
            if new_id and not interaction_id:
                # Update the convo with the real interaction ID
                convo["interaction_id"] = new_id
                self.selected_convo = convo
            self._load_messages(new_id or interaction_id)
            self._load_inbox()
        except RuntimeError as e:
            self.call_from_thread(self.notify, f"Send failed: {e}", severity="error")

    def action_blur_input(self) -> None:
        self.query_one("#convo-list", ListView).focus()

    def on_unmount(self) -> None:
        if self._stop_realtime:
            self._stop_realtime()


def run_tui() -> None:
    app = B2AApp()
    app.run()
