#!/usr/bin/env python3
"""Relay backend abstractions for lobstrd.

Backends:
- memory: in-process event store (default, test-friendly)
- websocket: Nostr relay over websocket (requires `websockets` package)
"""

from __future__ import annotations

import asyncio
import json
import urllib.request
import uuid
from dataclasses import dataclass
from typing import Any, Dict, List, Optional


_OK_PREFIXES = frozenset([
    "duplicate", "pow", "blocked", "rate-limited",
    "invalid", "restricted", "error",
])


def parse_ok_prefix(message: str) -> Optional[str]:
    """Extract known NIP-01 OK prefix from a relay message string.

    Known prefixes: duplicate, pow, blocked, rate-limited, invalid,
    restricted, error.  Returns the prefix without trailing colon,
    or None if no known prefix is found.
    """
    if not message:
        return None
    for prefix in _OK_PREFIXES:
        token = prefix + ":"
        if message.startswith(token):
            return prefix
    return None


@dataclass
class RelayFailureDetail:
    """Detailed failure info for a single relay in a publish result."""
    url: str
    reason: str
    ok_prefix: Optional[str] = None


@dataclass
class RelayFailureRecord:
    """Tracks the most recent failure for a single relay."""
    url: str
    reason: str
    ok_prefix: Optional[str] = None
    timestamp: int = 0


@dataclass
class RelayPublishResult:
    ok: bool
    published_to: List[str]
    failed_relays: List[RelayFailureDetail]
    error: Optional[str] = None


class RelayClient:
    def publish(self, event: Dict[str, Any]) -> RelayPublishResult:
        raise NotImplementedError

    def fetch(self, nostr_filter: Dict[str, Any]) -> List[Dict[str, Any]]:
        raise NotImplementedError

    def set_relay_url(self, relay_url: str) -> None:
        raise NotImplementedError

    @property
    def relay_url(self) -> str:
        raise NotImplementedError


class InMemoryRelayClient(RelayClient):
    def __init__(self, relay_url: str) -> None:
        self._relay_url = relay_url
        self._events: List[Dict[str, Any]] = []

    @property
    def relay_url(self) -> str:
        return self._relay_url

    def set_relay_url(self, relay_url: str) -> None:
        self._relay_url = relay_url

    def publish(self, event: Dict[str, Any]) -> RelayPublishResult:
        self._events.append(event)
        return RelayPublishResult(ok=True, published_to=[self._relay_url],
                                  failed_relays=[])

    def fetch(self, nostr_filter: Dict[str, Any]) -> List[Dict[str, Any]]:
        out: List[Dict[str, Any]] = []
        for ev in self._events:
            if _matches_nostr_filter(ev, nostr_filter):
                out.append(ev)
        limit = nostr_filter.get("limit")
        if isinstance(limit, int) and limit > 0:
            return out[-limit:]
        return out


class WebSocketRelayClient(RelayClient):
    def __init__(self, relay_url: str, timeout_sec: float = 5.0) -> None:
        self._relay_url = relay_url
        self._timeout_sec = timeout_sec

    @property
    def relay_url(self) -> str:
        return self._relay_url

    def set_relay_url(self, relay_url: str) -> None:
        self._relay_url = relay_url

    def publish(self, event: Dict[str, Any]) -> RelayPublishResult:
        try:
            ok, message = asyncio.run(self.publish_async(event))
        except Exception as exc:  # pragma: no cover
            return RelayPublishResult(
                ok=False,
                published_to=[],
                failed_relays=[RelayFailureDetail(
                    url=self._relay_url, reason=str(exc),
                    ok_prefix=None,
                )],
                error=str(exc),
            )
        if ok:
            return RelayPublishResult(ok=True, published_to=[self._relay_url],
                                      failed_relays=[])
        reason = message or "relay rejected EVENT"
        return RelayPublishResult(
            ok=False,
            published_to=[],
            failed_relays=[RelayFailureDetail(
                url=self._relay_url, reason=reason,
                ok_prefix=parse_ok_prefix(reason),
            )],
            error=reason,
        )

    def fetch(self, nostr_filter: Dict[str, Any]) -> List[Dict[str, Any]]:
        try:
            return asyncio.run(self.fetch_async(nostr_filter))
        except Exception:  # pragma: no cover
            return []

    async def publish_async(self, event: Dict[str, Any]) -> tuple[bool, str]:
        ws_mod = _import_websockets()
        async with ws_mod.connect(self._relay_url, open_timeout=self._timeout_sec) as ws:
            await ws.send(json.dumps(["EVENT", event]))
            try:
                reply = await asyncio.wait_for(ws.recv(), timeout=self._timeout_sec)
            except asyncio.TimeoutError:
                return False, "relay did not respond in time"
            parsed = json.loads(reply)
            # NIP-01: ["OK", <event_id>, <true|false>, <message>]
            if isinstance(parsed, list) and len(parsed) >= 3 and parsed[0] == "OK":
                ok = bool(parsed[2])
                message = parsed[3] if len(parsed) >= 4 and isinstance(parsed[3], str) else ""
                return ok, message
            return False, "unexpected relay response"

    async def fetch_async(self, nostr_filter: Dict[str, Any]) -> List[Dict[str, Any]]:
        ws_mod = _import_websockets()
        sub_id = f"lobstrd-{uuid.uuid4().hex[:12]}"
        events: List[Dict[str, Any]] = []

        async with ws_mod.connect(self._relay_url, open_timeout=self._timeout_sec) as ws:
            await ws.send(json.dumps(["REQ", sub_id, nostr_filter]))
            while True:
                try:
                    msg = await asyncio.wait_for(ws.recv(), timeout=self._timeout_sec)
                except asyncio.TimeoutError:
                    break
                parsed = json.loads(msg)
                if not isinstance(parsed, list) or len(parsed) < 2:
                    continue
                kind = parsed[0]
                if kind == "EVENT" and len(parsed) >= 3 and parsed[1] == sub_id and isinstance(parsed[2], dict):
                    events.append(parsed[2])
                elif kind == "EOSE" and parsed[1] == sub_id:
                    break
            await ws.send(json.dumps(["CLOSE", sub_id]))

        return events


@dataclass
class FetchResult:
    """Result of a multi-relay fetch with deduplication metadata."""
    events: List[Dict[str, Any]]
    queried: List[str]
    timed_out: List[str]


class RelayPool(RelayClient):
    """Wraps multiple relay clients with fanout publish and deduplicated fetch."""

    def __init__(self, relay_urls: List[str],
                 fanout_min_success: int = 1,
                 clients: Optional[List[RelayClient]] = None) -> None:
        self._relay_urls = list(relay_urls)
        self.fanout_min_success = fanout_min_success
        self.failure_history: Dict[str, RelayFailureRecord] = {}
        if clients is not None:
            self._clients = {url: c for url, c in zip(relay_urls, clients)}
        else:
            self._clients = {url: WebSocketRelayClient(url) for url in relay_urls}

    @property
    def relay_url(self) -> str:
        return self._relay_urls[0] if self._relay_urls else ""

    @property
    def relay_urls(self) -> List[str]:
        return list(self._relay_urls)

    def set_relay_url(self, relay_url: str) -> None:
        self._relay_urls = [relay_url]
        self._clients = {relay_url: WebSocketRelayClient(relay_url)}
        self.failure_history.clear()

    def set_relay_urls(self, relay_urls: List[str],
                       fanout_min_success: Optional[int] = None) -> None:
        self._relay_urls = list(relay_urls)
        self._clients = {url: WebSocketRelayClient(url) for url in relay_urls}
        if fanout_min_success is not None:
            self.fanout_min_success = fanout_min_success
        self.failure_history.clear()

    def publish(self, event: Dict[str, Any]) -> RelayPublishResult:
        results = asyncio.run(self._publish_fanout(event))
        published_to: List[str] = []
        failed_relays: List[RelayFailureDetail] = []
        for url, result in results:
            if result.ok:
                published_to.append(url)
                self.failure_history.pop(url, None)
            else:
                detail = result.failed_relays[0] if result.failed_relays else \
                    RelayFailureDetail(url=url, reason=result.error or "unknown")
                failed_relays.append(detail)
                self._record_failure(url, detail.reason, detail.ok_prefix)
        ok = len(published_to) >= self.fanout_min_success
        error = None
        if not ok:
            error = (f"publish failed: {len(published_to)}/{len(self._relay_urls)} "
                     f"relays accepted, need {self.fanout_min_success}")
        return RelayPublishResult(ok=ok, published_to=published_to,
                                  failed_relays=failed_relays, error=error)

    def fetch(self, nostr_filter: Dict[str, Any]) -> List[Dict[str, Any]]:
        result = self.fetch_with_status(nostr_filter)
        return result.events

    def fetch_with_status(self, nostr_filter: Dict[str, Any]) -> FetchResult:
        per_relay = asyncio.run(self._fetch_all(nostr_filter))
        seen: Dict[str, Dict[str, Any]] = {}
        timed_out: List[str] = []
        queried: List[str] = list(self._relay_urls)
        for url, events, ok in per_relay:
            if not ok:
                timed_out.append(url)
                self._record_failure(url, "fetch timeout", None)
                continue
            self.failure_history.pop(url, None)
            for ev in events:
                eid = ev.get("id")
                if eid and eid not in seen:
                    seen[eid] = ev
        return FetchResult(events=list(seen.values()), queried=queried,
                           timed_out=timed_out)

    async def _publish_fanout(self, event: Dict[str, Any]) -> List[tuple[str, RelayPublishResult]]:
        tasks = []
        for url in self._relay_urls:
            client = self._clients[url]
            tasks.append(self._publish_to_relay(url, client, event))
        return await asyncio.gather(*tasks)

    async def _publish_to_relay(self, url: str, client: RelayClient,
                                event: Dict[str, Any]) -> tuple[str, RelayPublishResult]:
        max_retries = 3
        backoff = 0.5
        last_result: Optional[RelayPublishResult] = None
        for attempt in range(max_retries + 1):
            try:
                result = await self._do_publish(client, event)
                last_result = result
                if result.ok:
                    return url, result
                # Explicit relay rejection — don't retry
                if result.failed_relays and result.failed_relays[0].ok_prefix is not None:
                    return url, result
                # Non-prefixed rejection could be transient
                if attempt < max_retries:
                    await asyncio.sleep(backoff * (2 ** attempt))
                    continue
                return url, result
            except Exception as exc:
                reason = str(exc)
                last_result = RelayPublishResult(
                    ok=False, published_to=[],
                    failed_relays=[RelayFailureDetail(url=url, reason=reason)],
                    error=reason,
                )
                if attempt < max_retries:
                    await asyncio.sleep(backoff * (2 ** attempt))
                    continue
        return url, last_result  # type: ignore[return-value]

    @staticmethod
    async def _do_publish(client: RelayClient, event: Dict[str, Any]) -> RelayPublishResult:
        """Publish using async path for WebSocketRelayClient, sync for others."""
        if isinstance(client, WebSocketRelayClient):
            try:
                ok, message = await client.publish_async(event)
            except Exception as exc:
                return RelayPublishResult(
                    ok=False, published_to=[],
                    failed_relays=[RelayFailureDetail(
                        url=client.relay_url, reason=str(exc),
                    )],
                    error=str(exc),
                )
            if ok:
                return RelayPublishResult(ok=True, published_to=[client.relay_url],
                                          failed_relays=[])
            reason = message or "relay rejected EVENT"
            return RelayPublishResult(
                ok=False, published_to=[],
                failed_relays=[RelayFailureDetail(
                    url=client.relay_url, reason=reason,
                    ok_prefix=parse_ok_prefix(reason),
                )],
                error=reason,
            )
        return client.publish(event)

    async def _fetch_all(self, nostr_filter: Dict[str, Any]) -> List[tuple[str, List[Dict[str, Any]], bool]]:
        tasks = []
        for url in self._relay_urls:
            client = self._clients[url]
            tasks.append(self._fetch_from_relay(url, client, nostr_filter))
        return await asyncio.gather(*tasks)

    async def _fetch_from_relay(self, url: str, client: RelayClient,
                                nostr_filter: Dict[str, Any]) -> tuple[str, List[Dict[str, Any]], bool]:
        try:
            if isinstance(client, WebSocketRelayClient):
                events = await client.fetch_async(nostr_filter)
            else:
                events = client.fetch(nostr_filter)
            return url, events, True
        except Exception:
            return url, [], False

    def _record_failure(self, url: str, reason: str,
                        ok_prefix: Optional[str]) -> None:
        import time
        self.failure_history[url] = RelayFailureRecord(
            url=url, reason=reason, ok_prefix=ok_prefix,
            timestamp=int(time.time()),
        )


def _import_websockets():
    try:
        import websockets.asyncio.client as ws_mod  # type: ignore
    except Exception as exc:  # pragma: no cover
        raise RuntimeError(
            "websocket relay backend requires the 'websockets' package (>=14.0). "
            "Install it with: pip install lobstr-core[websocket] "
            "or: pip install lobstrd[websocket]"
        ) from exc
    return ws_mod


def _matches_nostr_filter(event: Dict[str, Any], filt: Dict[str, Any]) -> bool:
    kinds = filt.get("kinds")
    if kinds and event.get("kind") not in kinds:
        return False

    authors = filt.get("authors")
    if authors and event.get("pubkey") not in authors:
        return False

    since = filt.get("since")
    if since is not None and event.get("created_at", 0) < since:
        return False

    for k, v in filt.items():
        if not k.startswith("#") or not isinstance(v, list) or not v:
            continue
        tag_name = k[1:]
        tag_values = []
        for t in event.get("tags", []):
            if isinstance(t, list) and len(t) > 1 and t[0] == tag_name:
                tag_values.append(t[1])
        if not any(tv in v for tv in tag_values):
            return False

    return True


def fetch_nip11_info(relay_url: str, timeout_sec: float = 5.0) -> Dict[str, Any]:
    """Fetch NIP-11 relay information document over HTTP.

    Converts wss:// to https:// (or ws:// to http://) and sends a GET
    with Accept: application/nostr+json.  Returns a dict with 'info'
    (parsed JSON or None) and 'error' (str or None).
    """
    http_url = relay_url.replace("wss://", "https://", 1).replace("ws://", "http://", 1)
    req = urllib.request.Request(http_url, headers={"Accept": "application/nostr+json"})
    try:
        with urllib.request.urlopen(req, timeout=timeout_sec) as resp:
            body = resp.read()
            info = json.loads(body.decode("utf-8"))
            if not isinstance(info, dict):
                return {"info": None, "error": "NIP-11 response is not a JSON object"}
            return {"info": info, "error": None}
    except Exception as exc:
        return {"info": None, "error": str(exc)}


def make_relay_client(backend: str, relay_url: str,
                      relay_urls: Optional[List[str]] = None,
                      fanout_min_success: int = 1) -> RelayClient:
    if backend == "websocket":
        urls = relay_urls or [relay_url]
        if len(urls) > 1:
            return RelayPool(urls, fanout_min_success=fanout_min_success)
        return WebSocketRelayClient(urls[0])
    return InMemoryRelayClient(relay_url)
