import json
from typing import TypedDict

import nio
from nio.crypto import InboundGroupSession, OlmDevice, TrustState
from slidge.command import Command, CommandAccess, Form, FormField, TableResult
from slidge.command.base import FormValues, 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):
    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,  # type: ignore[override]
        _ifrom: JID | None,
        *args: str,
    ) -> Form:
        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 Form(
            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,  # type:ignore
            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
        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):
    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,  # type: ignore[override]
        _ifrom: JID | None,
        *args: str,
    ) -> Form | str:
        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()),  # type:ignore
                        "new_state": "verified",
                    },
                    session,
                    None,
                    device_dict,
                )

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

        return Form(
            title=self.NAME,
            instructions="Choose the session(s) which trust state you want to change",
            handler=self.step2,  # type:ignore
            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
            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):
    NAME = "🔑 Change Password"
    CHAT_COMMAND = NODE = "change-password"
    HELP = "Change your Matrix account password."
    ACCESS = CommandAccess.USER_LOGGED

    async def run(
        self,
        session: Session,  # type: ignore[override]
        _ifrom: JID | None,
        *args: str,
    ) -> Form | str:
        return Form(
            title=self.NAME,
            instructions="Enter your current and new Matrix credentials.",
            handler=self.change_password,  # type:ignore
            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):
    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,  # type: ignore[override]
        _ifrom: JID | None,
        *args: str,
    ) -> Form | str:
        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 Form(
            title=self.NAME,
            instructions="Select device to delete from your account.",
            handler=self.delete_step,  # type:ignore
            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 | Form:
        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 Form(
                title=self.NAME,
                instructions="Authentication required to delete devices.",
                handler=self.delete_with_auth,  # type:ignore
                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):
    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,  # type: ignore[override]
        _ifrom: JID | None,
        *args: str,
    ) -> Form | str:
        if args:  # For chat command.
            room = args[0]
            return await self.join_room({"room": room}, session, None)

        return Form(
            title=self.NAME,
            instructions="Enter the Matrix room name (format: #public:example.com)",
            handler=self.join_room,  # type:ignore
            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):
    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,  # type: ignore[override]
        _ifrom: JID | None,
        *args: str,
    ) -> str:
        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):
    NAME = "🔐 Import room keys"
    CHAT_COMMAND = NODE = "import-keys"
    HELP = "Import Megolm session keys."
    ACCESS = CommandAccess.USER_LOGGED

    async def run(
        self,
        session: Session,  # type: ignore[override]
        _ifrom: JID | None,
        *args: str,
    ) -> Form | str:
        return Form(
            title=self.NAME,
            instructions="Paste the JSON content from the exported room keys file.",
            handler=self.import_keys,  # type:ignore
            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 = form_values["json_content"]
        assert isinstance(raw_lines, list)
        # Text is a list of lines; need to concat them.
        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."
        )
