Medios-Macina/Provider/matrix.py

from __future__ import annotations

import mimetypes
import time
import uuid
from pathlib import Path
from typing import Any, Dict, List, Optional, Tuple
from urllib.parse import quote

import requests

from ProviderCore.base import Provider


_MATRIX_INIT_CHECK_CACHE: Dict[str, Tuple[bool, Optional[str]]] = {}


def _normalize_homeserver(value: str) -> str:
	text = str(value or "").strip()
	if not text:
		return ""
	if not text.startswith("http"):
		text = f"https://{text}"
	return text.rstrip("/")


def _matrix_health_check(*, homeserver: str, access_token: Optional[str]) -> Tuple[bool, Optional[str]]:
	"""Lightweight Matrix reachability/auth validation.

	- Always checks `/versions` (no auth).
	- If `access_token` is present, also checks `/whoami`.
	"""
	try:
		base = _normalize_homeserver(homeserver)
		if not base:
			return False, "Matrix homeserver missing"

		resp = requests.get(f"{base}/_matrix/client/versions", timeout=5)
		if resp.status_code != 200:
			return False, f"Homeserver returned {resp.status_code}"

		if access_token:
			headers = {"Authorization": f"Bearer {access_token}"}
			resp = requests.get(f"{base}/_matrix/client/v3/account/whoami", headers=headers, timeout=5)
			if resp.status_code != 200:
				return False, f"Authentication failed: {resp.status_code}"

		return True, None
	except Exception as exc:
		return False, str(exc)


class Matrix(Provider):
	"""File provider for Matrix (Element) chat rooms."""

	def __init__(self, config: Optional[Dict[str, Any]] = None):
		super().__init__(config)
		self._init_ok: Optional[bool] = None
		self._init_reason: Optional[str] = None

		matrix_conf = self.config.get("provider", {}).get("matrix", {}) if isinstance(self.config, dict) else {}
		homeserver = matrix_conf.get("homeserver")
		access_token = matrix_conf.get("access_token")
		password = matrix_conf.get("password")

		# Not configured: keep instance but mark invalid via validate().
		# Note: `room_id` is intentionally NOT required, since the CLI can prompt
		# the user to select a room dynamically.
		if not (homeserver and (access_token or password)):
			self._init_ok = None
			self._init_reason = None
			return

		cache_key = f"{_normalize_homeserver(str(homeserver))}|has_token:{bool(access_token)}"
		cached = _MATRIX_INIT_CHECK_CACHE.get(cache_key)
		if cached is None:
			ok, reason = _matrix_health_check(homeserver=str(homeserver), access_token=str(access_token) if access_token else None)
			_MATRIX_INIT_CHECK_CACHE[cache_key] = (ok, reason)
		else:
			ok, reason = cached

		self._init_ok = ok
		self._init_reason = reason
		if not ok:
			raise Exception(reason or "Matrix unavailable")

	def validate(self) -> bool:
		if not self.config:
			return False
		if self._init_ok is False:
			return False
		matrix_conf = self.config.get("provider", {}).get("matrix", {})
		return bool(
			matrix_conf.get("homeserver")
			and (matrix_conf.get("access_token") or matrix_conf.get("password"))
		)

	def _get_homeserver_and_token(self) -> Tuple[str, str]:
		matrix_conf = self.config.get("provider", {}).get("matrix", {})
		homeserver = matrix_conf.get("homeserver")
		access_token = matrix_conf.get("access_token")
		if not homeserver:
			raise Exception("Matrix homeserver missing")
		if not access_token:
			raise Exception("Matrix access_token missing")
		base = _normalize_homeserver(str(homeserver))
		if not base:
			raise Exception("Matrix homeserver missing")
		return base, str(access_token)

	def list_joined_room_ids(self) -> List[str]:
		"""Return joined room IDs for the current user.

		Uses `GET /_matrix/client/v3/joined_rooms`.
		"""
		base, token = self._get_homeserver_and_token()
		headers = {"Authorization": f"Bearer {token}"}
		resp = requests.get(f"{base}/_matrix/client/v3/joined_rooms", headers=headers, timeout=10)
		if resp.status_code != 200:
			raise Exception(f"Matrix joined_rooms failed: {resp.text}")
		data = resp.json() or {}
		rooms = data.get("joined_rooms") or []
		out: List[str] = []
		for rid in rooms:
			if not isinstance(rid, str) or not rid.strip():
				continue
			out.append(rid.strip())
		return out

	def list_rooms(self, *, room_ids: Optional[List[str]] = None) -> List[Dict[str, Any]]:
		"""Return joined rooms, optionally limited to a subset.

		Performance note: room names require additional per-room HTTP requests.
		If `room_ids` is provided, only those rooms will have name lookups.
		"""
		base, token = self._get_homeserver_and_token()
		headers = {"Authorization": f"Bearer {token}"}

		joined = self.list_joined_room_ids()
		if room_ids:
			allowed = {str(v).strip().casefold() for v in room_ids if str(v).strip()}
			if allowed:
				# Accept either full IDs (!id:hs) or short IDs (!id).
				def _is_allowed(rid: str) -> bool:
					r = str(rid or "").strip()
					if not r:
						return False
					rc = r.casefold()
					if rc in allowed:
						return True
					short = r.split(":", 1)[0].strip().casefold()
					return bool(short) and short in allowed

				joined = [rid for rid in joined if _is_allowed(rid)]

		out: List[Dict[str, Any]] = []
		for room_id in joined:
			name = ""
			# Best-effort room name lookup (safe to fail).
			try:
				encoded = quote(room_id, safe="")
				name_resp = requests.get(
					f"{base}/_matrix/client/v3/rooms/{encoded}/state/m.room.name",
					headers=headers,
					timeout=5,
				)
				if name_resp.status_code == 200:
					payload = name_resp.json() or {}
					maybe = payload.get("name")
					if isinstance(maybe, str):
						name = maybe
			except Exception:
				pass
			out.append({"room_id": room_id, "name": name})
		return out

	def upload_to_room(self, file_path: str, room_id: str, **kwargs: Any) -> str:
		"""Upload a file and send it to a specific room."""
		from models import ProgressFileReader
		path = Path(file_path)
		if not path.exists():
			raise FileNotFoundError(f"File not found: {file_path}")
		if not room_id:
			raise Exception("Matrix room_id missing")

		base, token = self._get_homeserver_and_token()
		headers = {
			"Authorization": f"Bearer {token}",
			"Content-Type": "application/octet-stream",
		}

		mime_type, _ = mimetypes.guess_type(path)
		if mime_type:
			headers["Content-Type"] = mime_type

		filename = path.name

		# Upload media
		upload_url = f"{base}/_matrix/media/v3/upload"
		with open(path, "rb") as handle:
			wrapped = ProgressFileReader(handle, total_bytes=int(path.stat().st_size), label="upload")
			resp = requests.post(upload_url, headers=headers, data=wrapped, params={"filename": filename})
		if resp.status_code != 200:
			raise Exception(f"Matrix upload failed: {resp.text}")
		content_uri = (resp.json() or {}).get("content_uri")
		if not content_uri:
			raise Exception("No content_uri returned")

		# Build a fragment-free URL suitable for storage backends.
		# `matrix.to` links use fragments (`#/...`) which some backends normalize away.
		download_url_for_store = ""
		try:
			curi = str(content_uri or "").strip()
			if curi.startswith("mxc://"):
				rest = curi[len("mxc://"):]
				if "/" in rest:
					server_name, media_id = rest.split("/", 1)
					server_name = str(server_name).strip()
					media_id = str(media_id).strip()
					if server_name and media_id:
						download_url_for_store = f"{base}/_matrix/media/v3/download/{quote(server_name, safe='')}/{quote(media_id, safe='')}"
		except Exception:
			download_url_for_store = ""

		# Determine message type
		msgtype = "m.file"
		ext = path.suffix.lower()
		audio_exts = {".mp3", ".flac", ".wav", ".m4a", ".aac", ".ogg", ".opus", ".wma", ".mka", ".alac"}
		video_exts = {".mp4", ".mkv", ".webm", ".mov", ".avi", ".flv", ".mpg", ".mpeg", ".ts", ".m4v", ".wmv"}
		image_exts = {".jpg", ".jpeg", ".png", ".gif", ".webp", ".bmp", ".tiff"}
		if ext in audio_exts:
			msgtype = "m.audio"
		elif ext in video_exts:
			msgtype = "m.video"
		elif ext in image_exts:
			msgtype = "m.image"

		info = {"mimetype": mime_type, "size": path.stat().st_size}
		payload = {"msgtype": msgtype, "body": filename, "url": content_uri, "info": info}

		# Correct Matrix client API send endpoint requires a transaction ID.
		txn_id = f"mm_{int(time.time())}_{uuid.uuid4().hex[:8]}"
		encoded_room = quote(str(room_id), safe="")
		send_url = f"{base}/_matrix/client/v3/rooms/{encoded_room}/send/m.room.message/{txn_id}"
		send_headers = {"Authorization": f"Bearer {token}"}
		send_resp = requests.put(send_url, headers=send_headers, json=payload)
		if send_resp.status_code != 200:
			raise Exception(f"Matrix send message failed: {send_resp.text}")

		event_id = (send_resp.json() or {}).get("event_id")
		link = f"https://matrix.to/#/{room_id}/{event_id}" if event_id else f"https://matrix.to/#/{room_id}"

		# Optional: if a PipeObject is provided and it already has store+hash,
		# attach the uploaded URL back to the stored file.
		try:
			pipe_obj = kwargs.get("pipe_obj")
			if pipe_obj is not None:
				from Store import Store

				# Prefer the direct media download URL for storage backends.
				Store(self.config, suppress_debug=True).try_add_url_for_pipe_object(
					pipe_obj,
					download_url_for_store or link,
				)
		except Exception:
			pass

		return link

	def send_text_to_room(self, text: str, room_id: str) -> str:
		"""Send a plain text message to a specific room."""
		message = str(text or "").strip()
		if not message:
			return ""
		if not room_id:
			raise Exception("Matrix room_id missing")

		base, token = self._get_homeserver_and_token()
		encoded_room = quote(str(room_id), safe="")
		txn_id = f"mm_{int(time.time())}_{uuid.uuid4().hex[:8]}"
		send_url = f"{base}/_matrix/client/v3/rooms/{encoded_room}/send/m.room.message/{txn_id}"
		send_headers = {"Authorization": f"Bearer {token}"}
		payload = {"msgtype": "m.text", "body": message}
		send_resp = requests.put(send_url, headers=send_headers, json=payload)
		if send_resp.status_code != 200:
			raise Exception(f"Matrix send text failed: {send_resp.text}")

		event_id = (send_resp.json() or {}).get("event_id")
		return f"https://matrix.to/#/{room_id}/{event_id}" if event_id else f"https://matrix.to/#/{room_id}"

	def upload(self, file_path: str, **kwargs: Any) -> str:
		matrix_conf = self.config.get("provider", {}).get("matrix", {})
		room_id = matrix_conf.get("room_id")
		if not room_id:
			raise Exception("Matrix room_id missing")
		return self.upload_to_room(file_path, str(room_id))

	def selector(self, selected_items: List[Any], *, ctx: Any, stage_is_last: bool = True, **_kwargs: Any) -> bool:
		"""Handle Matrix room selection via `@N`.

		If the CLI has a pending upload stash, selecting a room triggers an upload.
		"""
		if not stage_is_last:
			return False

		pending = None
		try:
			pending = ctx.load_value('matrix_pending_uploads', default=None)
		except Exception:
			pending = None

		pending_list = list(pending) if isinstance(pending, list) else []
		if not pending_list:
			return False

		room_ids: List[str] = []
		for item in selected_items or []:
			rid = None
			if isinstance(item, dict):
				rid = item.get('room_id') or item.get('id')
			else:
				rid = getattr(item, 'room_id', None) or getattr(item, 'id', None)
			if rid and str(rid).strip():
				room_ids.append(str(rid).strip())

		if not room_ids:
			print("No Matrix room selected\n")
			return True

		any_failed = False
		for room_id in room_ids:
			for payload in pending_list:
				try:
					file_path = ''
					delete_after = False
					pipe_obj = None
					if isinstance(payload, dict):
						file_path = str(payload.get('path') or '')
						delete_after = bool(payload.get('delete_after', False))
						pipe_obj = payload.get('pipe_obj')
					else:
						file_path = str(getattr(payload, 'path', '') or '')
					if not file_path:
						any_failed = True
						continue

					media_path = Path(file_path)
					if not media_path.exists():
						any_failed = True
						print(f"Matrix upload file missing: {file_path}")
						continue

					link = self.upload_to_room(str(media_path), str(room_id), pipe_obj=pipe_obj)
					if link:
						print(link)

					if delete_after:
						try:
							media_path.unlink(missing_ok=True)  # type: ignore[arg-type]
						except TypeError:
							try:
								if media_path.exists():
									media_path.unlink()
							except Exception:
								pass
				except Exception as exc:
					any_failed = True
					print(f"Matrix upload failed: {exc}")

		try:
			ctx.store_value('matrix_pending_uploads', [])
		except Exception:
			pass

		if any_failed:
			print("\nOne or more Matrix uploads failed\n")
		return True