import json
from typing import TypedDict, cast

import nio
from nio.crypto import InboundGroupSession, OlmDevice, TrustState
from slidge.command import Command, CommandAccess, FormField, TableResult
from slidge.command.base import (
    FormRecipient,
    FormSession,
    FormValues,
    MUCCommand,
    Option,
)
from slidge.command.categories import CONTACTS, GROUPS
from slidge.util.types import LegacyAttachment
from slixmpp import JID
from slixmpp.exceptions import XMPPError

from .group import MUC
from .session import Session


class ListSpaces(Command[Session]):
    NAME = "🌌 Matrix spaces"
    CHAT_COMMAND = NODE = "spaces"
    CATEGORY = GROUPS
    HELP = "List the matrix spaces you're part of"
    ACCESS = CommandAccess.USER_LOGGED

    async def run(
        self,
        session: Session | None,
        _ifrom: JID | None,
        *args: str,
    ) -> FormSession[Session]:
        assert session is not None

        spaces = list[nio.MatrixRoom]()
        for room in session.matrix.rooms.values():
            if room.children:
                spaces.append(room)
        spaces = sorted(spaces, key=lambda r: r.name)
        return FormSession(
            title=self.NAME,
            instructions="Choose a space to list its children rooms. "
            "NB: as of now, you can also see rooms that you are a member of.",
            handler=self.finish,
            handler_args=(spaces,),
            fields=[
                FormField(
                    "space",
                    label="Matrix space",
                    type="list-single",
                    options=[
                        {"label": room.display_name, "value": str(i)}
                        for i, room in enumerate(spaces)
                    ],
                )
            ],
        )

    @staticmethod
    async def finish(
        form_values: FormValues,
        session: Session,
        _ifrom: JID | None,
        rooms: list[nio.MatrixRoom],
    ) -> TableResult:
        space = rooms[int(form_values["space"])]  # type:ignore[arg-type]
        mucs = list[MUC]()
        for room_id in space.children:
            try:
                mucs.append(await session.bookmarks.by_legacy_id(room_id))
            except XMPPError:
                continue

        mucs = sorted(mucs, key=lambda muc: muc.name or "Unnamed room")
        return TableResult(
            fields=[FormField("name"), FormField("jid", type="jid-single")],
            description=f"Rooms of '{space.display_name}'",
            jids_are_mucs=True,
            items=[
                {"name": muc.name or "Unnamed room", "jid": str(muc.jid)}
                for muc in mucs
            ],
        )


class ManageTrust(Command[Session]):
    NAME = "🤝 Manage trust in devices"
    CATEGORY = CONTACTS
    CHAT_COMMAND = NODE = "verify"
    HELP = "Manage which OLM keys you trust or not."
    ACCESS = CommandAccess.USER_LOGGED

    HUMAN_STATES = {0: "unset", 1: "verified", 2: "blacklisted", 3: "ignored"}

    def __human_device(self, d: OlmDevice, state: bool = True) -> str:
        r = f"{d.ed25519} of {d.user_id}"
        if state:
            return r + f" ({self.HUMAN_STATES[d.trust_state.value]})"
        return r

    async def run(
        self,
        session: Session | None,
        _ifrom: JID | None,
        *args: str,
    ) -> FormSession[Session] | str:
        assert session is not None

        devices = list[OlmDevice](session.matrix.olm.device_store)
        device_dict = {d.id: d for d in devices}

        # this part if for chat commands only
        if args:
            if args[0] == "all":
                return await self.step2(
                    {
                        "device": list(device_dict.keys()),
                        "new_state": "verified",
                    },
                    session,
                    None,
                    device_dict,
                )

            else:
                return await self.step2(
                    {"device": args[0].upper(), "new_state": "verified"},
                    session,
                    None,
                    device_dict,
                )

        return FormSession(
            title=self.NAME,
            instructions="Choose the session(s) which trust state you want to change",
            handler=self.step2,
            handler_args=(device_dict,),
            fields=[
                FormField(
                    "device",
                    label="Device(s)",
                    type="list-multi",
                    options=[
                        {"label": self.__human_device(d), "value": d.id}
                        for d in devices
                    ],
                ),
                FormField(
                    "new_state",
                    label="What new status do you want to give the selected devices?",
                    type="list-single",
                    options=[
                        {"label": v, "value": v} for v in self.HUMAN_STATES.values()
                    ],
                ),
            ],
        )

    async def step2(
        self,
        form_values: FormValues,
        session: Session,
        _ifrom: JID | None,
        devices: dict[str, OlmDevice],
    ) -> str:
        new_state = form_values["new_state"]
        matrix = session.matrix
        result = ""
        for device_name in form_values["device"]:  # type:ignore
            assert isinstance(device_name, str)
            device = devices[device_name]
            change = False
            if new_state == "unset":
                if device.trust_state == TrustState.verified:
                    change = matrix.unverify_device(device)
                elif device.trust_state == TrustState.ignored:
                    change = matrix.unignore_device(device)
                elif device.trust_state == TrustState.blacklisted:
                    change = matrix.unblacklist_device(device)
            elif new_state == "verified":
                change = session.matrix.verify_device(device)
            elif new_state == "blacklisted":
                change = session.matrix.blacklist_device(device)
            elif new_state == "ignored":
                change = session.matrix.ignore_device(device)

            if change:
                result += (
                    f"\nThe status of {self.__human_device(device, False)} "
                    f"is now {new_state}."
                )
            else:
                result += (
                    f"\nThe status of {self.__human_device(device, False)} "
                    f"has not changed."
                )

        return result or "Nothing was changed."


class ChangePassword(Command[Session]):
    NAME = "🔑 Change Password"
    CHAT_COMMAND = NODE = "change-password"
    HELP = "Change your Matrix account password."
    ACCESS = CommandAccess.USER_LOGGED

    async def run(
        self,
        session: Session | None,
        _ifrom: JID | None,
        *args: str,
    ) -> FormSession[Session] | str:
        assert session is not None

        return FormSession(
            title=self.NAME,
            instructions="Enter your current and new Matrix credentials.",
            handler=self.change_password,
            fields=[
                FormField(
                    "current_password",
                    label="Current Password",
                    type="text-private",
                    required=True,
                ),
                FormField(
                    "new_password",
                    label="New Password",
                    type="text-private",
                    required=True,
                ),
                FormField(
                    "confirm_new_password",
                    label="Confirm New Password",
                    type="text-private",
                    required=True,
                ),
            ],
        )

    async def change_password(
        self,
        form_values: FormValues,
        session: Session,
        _ifrom: JID | None,
    ) -> str:
        current_password = form_values["current_password"]
        new_password = form_values["new_password"]
        confirm_new_password = form_values["confirm_new_password"]

        if new_password != confirm_new_password:
            raise XMPPError("not-acceptable", "New passwords do not match.")

        user_id = session.matrix.user_id
        homeserver = session.matrix.homeserver
        username = user_id.split(":")[0][1:]

        url = f"{homeserver.rstrip('/')}/_matrix/client/r0/account/password"
        payload = {
            "auth": {
                "type": "m.login.password",
                "identifier": {"type": "m.id.user", "user": username},
                "password": current_password,
            },
            "new_password": new_password,
        }
        headers = {
            "Authorization": f"Bearer {session.matrix.access_token}",
            "Content-Type": "application/json",
        }

        async with session.http.post(
            url,
            json=payload,
            headers=headers,
        ) as response:
            if response.status == 200:
                return "Password changed successfully!"

            error_data = await response.json()
            error_msg = error_data.get("error", "Unknown error")
            errcode = error_data.get("errcode", "M_UNKNOWN")
            raise XMPPError(
                "internal-server-error",
                f"Server error: [{errcode}] {error_msg}",
            )


class DeleteDevices(Command[Session]):
    NAME = "🔑 Delete devices"
    CHAT_COMMAND = NODE = "delete-devices"
    HELP = "Remove devices from your Matrix account."
    ACCESS = CommandAccess.USER_LOGGED

    def _device_label(self, device: nio.Device) -> str:
        return f"{device.id}: {device.display_name or '(no name)'}"

    async def run(
        self,
        session: Session | None,
        _ifrom: JID | None,
        *args: str,
    ) -> FormSession[Session] | str:
        assert session is not None

        response = await session.matrix.devices()

        if isinstance(response, nio.DevicesError):
            raise XMPPError(
                "internal-server-error",
                f"Failed to list devices: {response.message}",
            )

        if not (devices := response.devices):
            return "No devices found."

        device_opts: list[Option] = [
            {"label": self._device_label(d), "value": d.id} for d in devices
        ]
        return FormSession(
            title=self.NAME,
            instructions="Select device to delete from your account.",
            handler=self.delete_step,
            fields=[
                FormField(
                    "devices",
                    label="Devices to delete",
                    type="list-multi",
                    options=device_opts,
                    required=True,
                ),
            ],
        )

    async def delete_step(
        self,
        form_values: FormValues,
        session: Session,
        _ifrom: JID | None,
    ) -> str | FormSession[Session]:
        devices_to_delete: list[str] = form_values["devices"]  # type:ignore

        response = await session.matrix.delete_devices(devices_to_delete)

        if isinstance(response, nio.DeleteDevicesError):
            raise XMPPError(
                "internal-server-error",
                f"Failed to delete devices: {response.message}",
            )

        if isinstance(response, nio.DeleteDevicesAuthResponse):
            return FormSession(
                title=self.NAME,
                instructions="Authentication required to delete devices.",
                handler=self.delete_with_auth,
                handler_args=(devices_to_delete, response),
                fields=[
                    FormField(
                        "password",
                        label="Password",
                        type="text-private",
                        required=True,
                    ),
                ],
            )

        # XXX: this branch untested; my server requires auth.
        deleted_names = ", ".join(devices_to_delete)
        return f"Successfully deleted devices: {deleted_names}"

    async def delete_with_auth(
        self,
        form_values: FormValues,
        session: Session,
        _ifrom: JID | None,
        devices_to_delete: list[str],
        auth_response: nio.DeleteDevicesAuthResponse,
    ) -> str:
        """Prompts for a password and completes the deletion.

        nio doesn't store the password, so we need to prompt for it.
        """
        password = form_values["password"]
        user_id = session.matrix.user_id
        username = user_id.split(":")[0][1:]

        auth = {
            "type": "m.login.password",
            "identifier": {"type": "m.id.user", "user": username},
            "password": password,
            "session": auth_response.session,
        }

        response = await session.matrix.delete_devices(devices_to_delete, auth)

        if isinstance(response, nio.DeleteDevicesError):
            raise XMPPError(
                "internal-server-error",
                f"Failed to delete devices: {response.message}",
            )

        names = ", ".join(devices_to_delete)
        return f"Successfully deleted devices: {names}"


class JoinRoom(Command[Session]):
    NAME = "🚪 Join Matrix room"
    CHAT_COMMAND = NODE = "join-room"
    CATEGORY = GROUPS
    HELP = "Join a Matrix room by name or alias."
    ACCESS = CommandAccess.USER_LOGGED

    async def run(
        self,
        session: Session | None,
        _ifrom: JID | None,
        *args: str,
    ) -> FormSession[Session] | str:
        assert session is not None

        if args:  # For chat command.
            room = args[0]
            return await self.join_room({"room": room}, session, None)

        return FormSession(
            title=self.NAME,
            instructions="Enter the Matrix room name (format: #public:example.com)",
            handler=self.join_room,
            fields=[
                FormField(
                    "room",
                    label="Room name.",
                    type="text-single",
                    required=True,
                ),
            ],
        )

    async def join_room(
        self,
        form_values: FormValues,
        session: Session,
        _ifrom: JID | None,
    ) -> str:
        room = form_values["room"]

        response = await session.matrix.join(room)

        if isinstance(response, nio.JoinError):
            raise XMPPError(
                "internal-server-error",
                f"Failed to join room: {response.message}",
            )

        room_id = response.room_id

        muc = await session.bookmarks.by_legacy_id(room_id)

        try:
            await muc.add_to_bookmarks()
        except XMPPError as e:
            # Reconnecting to the gateway should make the room show up.
            return f"Joined room {room}, but couldn't add bookmark: {e}"

        return f"Successfully joined room {room} (ID: {room_id})"


class SessionKeyData(TypedDict):
    algorithm: str
    sender_key: str
    sender_claimed_keys: dict[str, str]
    forwarding_curve25519_key_chain: list[str]
    room_id: str
    session_id: str
    session_key: str


class ExportRoomKeys(Command[Session]):
    NAME = "🔐 Export room keys"
    CHAT_COMMAND = NODE = "export-keys"
    HELP = "Export all Megolm session keys."
    ACCESS = CommandAccess.USER_LOGGED

    async def run(
        self,
        session: Session | None,
        _ifrom: JID | None,
        *args: str,
    ) -> str:
        assert session is not None

        matrix = session.matrix
        sessions_list: list[SessionKeyData] = []

        inbound_store = matrix.store.load_inbound_group_sessions()
        for sess in inbound_store:
            sessions_list.append(
                {
                    "algorithm": "m.megolm.v1.aes-sha2",
                    "sender_key": sess.sender_key,
                    "sender_claimed_keys": {"ed25519": sess.ed25519},
                    "forwarding_curve25519_key_chain": sess.forwarding_chain,
                    "room_id": sess.room_id,
                    "session_id": sess.id,
                    "session_key": sess.export_session(sess.first_known_index),
                }
            )

        device_curve25519 = matrix.olm.account.identity_keys["curve25519"]
        device_ed25519 = matrix.olm.account.identity_keys["ed25519"]
        for room_id, sess in matrix.olm.outbound_group_sessions.items():
            sessions_list.append(
                {
                    "algorithm": "m.megolm.v1.aes-sha2",
                    "sender_key": device_curve25519,
                    "sender_claimed_keys": {"ed25519": device_ed25519},
                    "forwarding_curve25519_key_chain": [],
                    "room_id": room_id,
                    "session_id": sess.session_id,
                    "session_key": sess.export_session(sess.first_known_index),
                }
            )

        data = json.dumps(sessions_list, indent=2).encode("utf-8")

        attachment = LegacyAttachment(
            data=data,
            name="matrix_room_keys.json",
            content_type="application/json",
        )
        await session.xmpp.send_file(attachment, mto=session.user_jid)

        return f"Room keys exported ({len(sessions_list)} sessions)"


class ImportRoomKeys(Command[Session]):
    NAME = "🔐 Import room keys"
    CHAT_COMMAND = NODE = "import-keys"
    HELP = "Import Megolm session keys."
    ACCESS = CommandAccess.USER_LOGGED

    async def run(
        self,
        session: Session | None,
        _ifrom: JID | None,
        *args: str,
    ) -> FormSession[Session] | str:
        assert session is not None

        return FormSession(
            title=self.NAME,
            instructions="Paste the JSON content from the exported room keys file.",
            handler=self.import_keys,
            fields=[
                FormField(
                    "json_content",
                    label="Room keys JSON",
                    type="text-multi",
                    required=True,
                ),
            ],
        )

    async def import_keys(
        self,
        form_values: FormValues,
        session: Session,
        _ifrom: JID | None,
    ) -> str:
        raw_lines = cast(list[str], form_values["json_content"])
        json_content = "".join(raw_lines)

        try:
            sessions_list = json.loads(json_content)
        except json.JSONDecodeError as e:
            raise XMPPError(
                "bad-request",
                f"Invalid JSON: {e}",
            )

        if not isinstance(sessions_list, list):
            raise XMPPError(
                "bad-request",
                "Invalid format: expected a JSON array of sessions",
            )

        matrix = session.matrix

        imported = 0
        skipped = 0
        failed = 0

        store = matrix.olm.inbound_group_store
        for key_data in sessions_list:
            try:
                room_id = key_data["room_id"]
                session_id = key_data["session_id"]
                sender_key = key_data["sender_key"]

                # Nio always returns True, even for existing keys.
                # See: https://github.com/matrix-nio/matrix-nio/issues/561
                existing = store.get(room_id, sender_key, session_id)
                if existing:
                    skipped += 1
                    session.log.debug(
                        f"Existed: {room_id} sender={sender_key} session={session_id}"
                    )
                    continue

                inbound_session = InboundGroupSession.import_session(
                    session_key=key_data["session_key"],
                    signing_key=key_data["sender_claimed_keys"]["ed25519"],
                    sender_key=sender_key,
                    room_id=room_id,
                    forwarding_chain=key_data.get(
                        "forwarding_curve25519_key_chain", []
                    ),
                )

                if store.add(inbound_session):
                    matrix.store.save_inbound_group_session(inbound_session)
                    imported += 1
                    session.log.debug(
                        f"Imported: {room_id} "
                        f"sender={sender_key} session={session_id} "
                        f"first_known_index={inbound_session.first_known_index}"
                    )
                else:
                    skipped += 1
                    session.log.debug(
                        f"Skipped: {room_id} sender={sender_key} session={session_id}"
                    )

            except KeyError as e:
                session.log.warning(f"Missing required field {e} in session entry")
                failed += 1
            except Exception as e:
                session.log.warning(f"Failed to import session: {e}")
                failed += 1

        session.log.debug(f"Imported {imported}, skipped {skipped}, failed {failed}")

        # Previous backfills discarded messages which required these decryption keys.
        # Re-fetch them for processing with imported keys.

        # Slidge's deduplication (by stanza_id/legacy_id) ensures that
        # already-backfilled messages are updated rather than duplicated.
        imported_rooms = {key_data["room_id"] for key_data in sessions_list}
        rebackfilled = 0
        for room_id in imported_rooms:
            try:
                muc = await session.bookmarks.by_legacy_id(room_id)
            except XMPPError:
                continue  # No longer present in room.
            muc.stored.history_filled = False
            try:
                await muc.backfill()
                muc.commit()
                rebackfilled += 1
            except Exception as e:
                session.log.warning("Failed to re-backfill %s: %s", room_id, e)
            finally:
                muc.stored.history_filled = True
            # FIXME: change to info once we've finished testing.
            session.log.error(f"Re-backfilled {room_id}.")

        return (
            f"Imported {imported} new session keys,"
            f" {skipped} already present, {failed} failed."
            f" Re-backfilled {rebackfilled} rooms."
        )


class TrustUnverified(MUCCommand[MUC]):
    NAME = "🔐 Trust unverified devices"
    NODE = CHAT_COMMAND = "trust-unverified"
    HELP = "Toggle whether to trust unverified devices in this room"

    @staticmethod
    async def run(muc: MUC, *args: str) -> FormRecipient[MUC]:
        current = muc.ignore_unverified_devices
        current_str = "ENABLED" if current else "DISABLED"
        return FormRecipient(
            title=TrustUnverified.NAME,
            instructions=(
                f"Room: {muc.name or muc.legacy_id}\n\n"
                f"Current setting: {current_str}\n\n"
                "When enabled, messages will be sent to unverified devices "
                "in this room. This is less secure but may be necessary if "
                "room members have unverified devices.\n\n"
                "When disabled, only verified devices will "
                "receive your messages."
            ),
            handler=TrustUnverified.toggle,
            fields=[
                FormField(
                    "ignore_unverified",
                    label="Trust unverified devices in this room?",
                    type="boolean",
                    value=str(current).lower(),
                    required=True,
                )
            ],
        )

    @staticmethod
    async def toggle(muc: MUC, form_values: FormValues) -> str:
        new_value = form_values["ignore_unverified"] == "true"
        muc.ignore_unverified_devices = new_value
        if new_value:
            return (
                f"Room '{muc.name or muc.legacy_id}': Unverified devices are now TRUSTED.\n"
                "Your messages will be sent to all devices, even if unverified.\n"
                "This reduces end-to-end encryption security."
            )
        else:
            return (
                f"Room '{muc.name or muc.legacy_id}': Unverified devices are now IGNORED.\n"
                "Your messages will only be sent to verified devices.\n"
                "This is the recommended secure setting."
            )
