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 FileProvider


_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(FileProvider):
	"""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_rooms(self) -> List[Dict[str, Any]]:
		"""Return the rooms the current user has joined.

		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[Dict[str, Any]] = []
		for rid in rooms:
			if not isinstance(rid, str) or not rid.strip():
				continue
			room_id = rid.strip()
			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) -> str:
		"""Upload a file and send it to a specific room."""
		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:
			resp = requests.post(upload_url, headers=headers, data=handle, 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")

		# 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")
		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))