"""
Agent Zero Zulip Adapter — Direct Zulip integration for Agent Zero agents.

Architecture:
  Zulip event queue → poll loop → A2A send to Agent Zero → 
  wait for A2A response → post back to Zulip

Runs alongside Agent Zero on the same host. Communicates with Agent Zero
via local A2A protocol (port 8001). No intermediate bridge needed.
"""

import asyncio
import json
import logging
import os
import time
from datetime import datetime, timezone
from typing import Any, Optional
import urllib.parse

try:
    import httpx
except ImportError:
    httpx = None

logger = logging.getLogger("agent-zero-zulip")

# Constants
DEFAULT_A2A_URL = "http://localhost:8001/a2a"
DEFAULT_A2A_TOKEN = "8zNgdOEXzYxjQvTl"
DEFAULT_POLL_INTERVAL = 3.0
MAX_ZULIP_MSG = 10000
RECONNECT_BACKOFF = [2, 5, 10, 30, 60]
HEARTBEAT_INTERVAL = 300  # 5 min


class AgentZeroZulipAdapter:
    """Zulip adapter for Agent Zero agents.
    
    Connects to Zulip via event queue, forwards DMs and @mentions
    to Agent Zero via local A2A protocol, and posts responses back.
    """

    def __init__(self):
        # Zulip config from env vars
        self._site = (os.getenv("ZULIP_SITE") or "").rstrip("/")
        self._email = os.getenv("ZULIP_EMAIL") or ""
        self._api_key = os.getenv("ZULIP_API_KEY") or ""
        self._agent_name = os.getenv("ZULIP_AGENT_NAME") or "kagentz"
        self._stream = os.getenv("ZULIP_STREAM") or "agent-hub"
        
        # A2A config
        self._a2a_url = os.getenv("A2A_URL") or DEFAULT_A2A_URL
        self._a2a_token = os.getenv("A2A_TOKEN") or DEFAULT_A2A_TOKEN
        
        # State
        self._auth_header = self._build_auth()
        self._http_client: Optional[httpx.AsyncClient] = None
        self._queue_id: Optional[str] = None
        self._last_event_id: int = -1
        self._bot_user_id: Optional[int] = None
        self._bot_email: str = self._email
        self._running = False
        self._poll_task: Optional[asyncio.Task] = None
        
        # Stats
        self._messages_processed = 0
        self._poll_count = 0
        self._reconnects = 0
        self._last_heartbeat = 0.0
        self._last_event_time = time.time()
        self._a2a_sessions: dict = {}  # chat_id -> A2A context_id
    
    def _build_auth(self) -> str:
        import base64
        token = base64.b64encode(f"{self._email}:{self._api_key}".encode()).decode()
        return f"Basic {token}"
    
    async def start(self):
        """Start the adapter — connect to Zulip and begin polling."""
        if not all([self._site, self._email, self._api_key]):
            logger.error("Missing Zulip credentials. Set ZULIP_SITE, ZULIP_EMAIL, ZULIP_API_KEY")
            return False
        
        if not httpx:
            logger.error("httpx not installed")
            return False
        
        self._running = True
        self._http_client = httpx.AsyncClient(timeout=30.0)
        
        # Connect to Zulip
        if not await self._connect():
            return False
        
        logger.info(f"Adapter started for {self._agent_name} ({self._email})")
        return True
    
    async def stop(self):
        """Stop the adapter."""
        self._running = False
        if self._poll_task:
            self._poll_task.cancel()
            try:
                await self._poll_task
            except asyncio.CancelledError:
                pass
        if self._http_client:
            await self._http_client.aclose()
    
    async def _connect(self) -> bool:
        """Register Zulip event queue."""
        try:
            resp = await self._api_call("POST", "/api/v1/register", data={
                "event_types": '["message"]',
                "apply_markdown": "true",
            })
            if not resp:
                logger.error("Queue registration failed")
                return False
            
            self._queue_id = resp.get("queue_id")
            self._last_event_id = resp.get("last_event_id", -1)
            self._bot_user_id = resp.get("user_id")
            
            # Resolve bot user_id if not provided by register
            if self._bot_user_id is None:
                me = await self._api_call("GET", "/api/v1/users/me")
                if me:
                    self._bot_user_id = me.get("user_id")
            
            logger.info(
                f"Connected to Zulip as {self._email} "
                f"(queue={self._queue_id}, bot_id={self._bot_user_id})"
            )
            
            # Start poll loop
            self._poll_task = asyncio.create_task(self._poll_forever())
            return True
            
        except Exception as e:
            logger.error(f"Connection failed: {e}")
            return False
    
    async def _poll_forever(self):
        """Main poll loop."""
        backoff = 0
        while self._running:
            if not self._queue_id:
                await asyncio.sleep(1)
                continue
            
            try:
                events = await self._fetch_events()
                self._poll_count += 1
                
                for event in events:
                    try:
                        await self._process_event(event)
                    except Exception as e:
                        logger.warning(f"Event processing error: {e}")
                
                backoff = 0
                
                # Heartbeat
                now = time.time()
                if now - self._last_heartbeat > HEARTBEAT_INTERVAL:
                    self._last_heartbeat = now
                    silence = now - self._last_event_time
                    logger.info(
                        f"Heartbeat — polls={self._poll_count} "
                        f"processed={self._messages_processed} "
                        f"silence={silence:.0f}s reconnects={self._reconnects}"
                    )
                    
            except asyncio.CancelledError:
                break
            except Exception as e:
                err = str(e)
                if "BAD_EVENT_QUEUE_ID" in err:
                    logger.info("Queue expired, reconnecting...")
                    await self._reconnect()
                    continue
                delay = RECONNECT_BACKOFF[min(backoff, len(RECONNECT_BACKOFF) - 1)]
                backoff += 1
                logger.warning(f"Poll error: {e}")
                await asyncio.sleep(delay)
            
            await asyncio.sleep(DEFAULT_POLL_INTERVAL)
    
    async def _fetch_events(self) -> list:
        """Fetch events from Zulip queue. Raises on BAD_EVENT_QUEUE_ID."""
        resp, status = await self._api_call("GET", "/api/v1/events", params={
            "queue_id": self._queue_id,
            "last_event_id": str(self._last_event_id),
            "dont_block": "true",
        }, return_status=True)
        
        # Detect queue expiry
        if status == 400:
            raise RuntimeError("BAD_EVENT_QUEUE_ID: queue expired")
        
        if not resp:
            return []
        
        events = resp.get("events", [])
        for event in events:
            eid = event.get("id", 0)
            if eid > self._last_event_id:
                self._last_event_id = eid
        
        return [e for e in events if e.get("type") == "message"]
    
    async def _reconnect(self):
        """Re-register event queue."""
        self._queue_id = None
        try:
            resp = await self._api_call("POST", "/api/v1/register", data={
                "event_types": '["message"]',
                "apply_markdown": "true",
            })
            if resp:
                self._queue_id = resp.get("queue_id")
                self._last_event_id = resp.get("last_event_id", -1)
                self._reconnects += 1
                logger.info(f"Reconnected, queue={self._queue_id}")
        except Exception as e:
            logger.error(f"Reconnect failed: {e}")
    
    async def _process_event(self, event: dict):
        """Process a Zulip message event."""
        msg = event.get("message", {})
        msg_type = msg.get("type", "")
        sender_email = msg.get("sender_email", "")
        sender_name = msg.get("sender_full_name", "Unknown")
        sender_id = msg.get("sender_id")
        content = msg.get("content", "")
        
        # Ignore own messages
        if sender_email == self._bot_email:
            return
        
        # Determine targeting
        mentioned_users = msg.get("mentioned_users", []) or []
        mentioned_ids = [u.get("user_id") for u in mentioned_users if isinstance(u, dict)]
        is_dm = msg_type == "private"
        is_mention = self._bot_user_id and self._bot_user_id in mentioned_ids
        
        if is_dm:
            self._messages_processed += 1
            self._last_event_time = time.time()
            logger.info(f"DM from {sender_name}: {content[:80]}...")
            
            # Send typing indicator
            await self._typing(sender_id, "start")
            
            # Send placeholder
            placeholder_id = await self._send_msg("private", str(sender_id), 
                ":robot: _Processing your message..._")
            
            # Forward to Agent Zero via A2A
            response = await self._a2a_chat(sender_name, content)
            
            # Stop typing
            await self._typing(sender_id, "stop")
            
            # Post response
            if response and placeholder_id:
                await self._edit_msg(placeholder_id, response[:MAX_ZULIP_MSG])
                logger.info(f"Responded to {sender_name} ({len(response)} chars)")
            elif response:
                await self._send_msg("private", str(sender_id), response[:MAX_ZULIP_MSG])
    
    async def _a2a_chat(self, sender_name: str, message: str) -> Optional[str]:
        """Send message to Agent Zero via local A2A protocol with session memory."""
        try:
            # Use sender_name as stable session key so conversation persists
            session_id = f"zulip-{sender_name.lower().replace(' ', '-')}"
            
            a2a_payload = {
                "jsonrpc": "2.0",
                "method": "tasks/send",
                "params": {
                    "id": f"zulip-{int(time.time())}",
                    "message": {
                        "role": "user",
                        "parts": [{"text": f"[Zulip DM from {sender_name}]: {message}"}]
                    },
                    "metadata": {
                        "session_id": session_id
                    }
                },
                "id": 1
            }
            
            headers = {
                "Content-Type": "application/json",
                "Authorization": f"Bearer {self._a2a_token}",
            }
            
            async with httpx.AsyncClient(timeout=120.0) as client:
                # Send the message to Agent Zero
                resp = await client.post(
                    self._a2a_url,
                    json=a2a_payload,
                    headers=headers,
                )
                
                if resp.status_code != 200:
                    logger.error(f"A2A send failed: HTTP {resp.status_code}")
                    return f":warning: Failed to reach agent. Error: HTTP {resp.status_code}"
                
                result = resp.json()
                task_id = result.get("result", {}).get("id")
                
                if not task_id:
                    logger.error("A2A: no task_id returned")
                    return ":warning: Agent did not create a task."
                
                # Poll for completion
                for _ in range(60):  # max 60s wait
                    await asyncio.sleep(1)
                    status_resp = await client.post(
                        self._a2a_url,
                        json={
                            "jsonrpc": "2.0",
                            "method": "tasks/get",
                            "params": {"id": task_id},
                            "id": 2
                        },
                        headers=headers,
                    )
                    
                    if status_resp.status_code != 200:
                        continue
                    
                    status_data = status_resp.json()
                    result_data = status_data.get("result", {})
                    state = result_data.get("status", {}).get("state", "")
                    
                    if state == "completed":
                        # Extract response text
                        return self._extract_a2a_response(result_data)
                    elif state in ("failed", "error"):
                        error_msg = result_data.get("status", {}).get("message", {}).get("text", "Unknown error")
                        return f":warning: Agent processing failed: {error_msg}"
                
                return ":hourglass: Agent did not complete in time."
                
        except httpx.TimeoutException:
            return ":hourglass: Agent request timed out."
        except Exception as e:
            logger.error(f"A2A error: {e}")
            return f":warning: Communication error: {type(e).__name__}"
    
    def _extract_a2a_response(self, result: dict) -> str:
        """Extract assistant text from A2A response."""
        history = result.get("history", [])
        for msg in reversed(history):
            if isinstance(msg, dict) and msg.get("role") == "assistant":
                parts = msg.get("parts", [])
                texts = [p.get("text", "") for p in parts if isinstance(p, dict)]
                text = "\n".join(t for t in texts if t)
                if text:
                    return text
        
        artifacts = result.get("artifacts", [])
        for artifact in reversed(artifacts):
            if isinstance(artifact, dict):
                text = artifact.get("text") or artifact.get("content") or ""
                if text:
                    return text
        
        return "(no response)"
    
    # ── Zulip API ──
    
    async def _api_call(self, method: str, path: str, 
                        data: dict = None, params: dict = None,
                        return_status: bool = False) -> Any:
        """Make Zulip API call.
        
        Returns: dict response body, or if return_status=True returns (dict, int)
        """
        if not self._http_client:
            return (None, 0) if return_status else None
        
        url = f"{self._site}{path}"
        headers = {"Authorization": self._auth_header}
        
        try:
            if method == "GET":
                resp = await self._http_client.get(url, params=params, headers=headers)
            elif method == "POST":
                resp = await self._http_client.post(url, data=data, headers=headers)
            elif method == "PATCH":
                resp = await self._http_client.patch(url, data=data, headers=headers)
            else:
                return (None, 0) if return_status else None
            
            if resp.status_code >= 400:
                logger.debug(f"API {method} {path}: {resp.status_code}")
                return (None, resp.status_code) if return_status else None
            
            body = resp.json()
            return (body, resp.status_code) if return_status else body
        except Exception as e:
            logger.debug(f"API error {method} {path}: {e}")
            return (None, 0) if return_status else None
    
    async def _send_msg(self, msg_type: str, to: str, content: str) -> Optional[int]:
        """Send a message to Zulip."""
        payload = {"type": msg_type, "content": content}
        if msg_type == "private":
            payload["to"] = json.dumps([int(to)])
        else:
            payload["to"] = to
            payload["subject"] = "general"
        
        resp = await self._api_call("POST", "/api/v1/messages", data=payload)
        if resp and resp.get("id"):
            return resp["id"]
        return None
    
    async def _edit_msg(self, message_id: int, content: str):
        """Edit a Zulip message."""
        await self._api_call("PATCH", f"/api/v1/messages/{message_id}",
                             data={"content": content})
    
    async def _typing(self, user_id: int, op: str):
        """Send typing indicator."""
        if not user_id:
            return
        to_data = json.dumps([user_id])
        await self._api_call("POST", "/api/v1/typing",
                             data={"to": to_data, "op": op})
    
    def run(self):
        """Synchronous entry point."""
        asyncio.run(self._run_async())
    
    async def _run_async(self):
        try:
            if await self.start():
                # Keep running until stopped
                while self._running:
                    await asyncio.sleep(1)
        except asyncio.CancelledError:
            pass
        finally:
            await self.stop()


# ── Entry point ──

def main():
    """CLI entry point."""
    logging.basicConfig(
        level=logging.INFO,
        format="%(asctime)s [%(name)s] %(levelname)s: %(message)s",
        datefmt="%Y-%m-%d %H:%M:%S",
    )
    
    adapter = AgentZeroZulipAdapter()
    
    try:
        adapter.run()
    except KeyboardInterrupt:
        logger.info("Shutdown requested")
    
    # Handle signals
    import signal
    shutdown = asyncio.Event()
    
    def _handle_signal():
        shutdown.set()
    
    loop = asyncio.new_event_loop()
    asyncio.set_event_loop(loop)
    
    try:
        loop.add_signal_handler(signal.SIGTERM, _handle_signal)
        loop.add_signal_handler(signal.SIGINT, _handle_signal)
    except NotImplementedError:
        pass
    
    async def _run():
        if await adapter.start():
            await shutdown.wait()
        await adapter.stop()
    
    try:
        loop.run_until_complete(_run())
    except KeyboardInterrupt:
        pass
    finally:
        loop.close()


if __name__ == "__main__":
    main()