"""
OpenAI-Compatible API Adapter — Bridge any OpenAI-compatible local LLM to Brynq.

Works with: LM Studio, vLLM, text-generation-webui, LocalAI, Ollama (OpenAI mode),
llama.cpp server, FastChat, Aphrodite, TabbyAPI, or any server that exposes
/v1/chat/completions with the OpenAI request/response schema.

The adapter registers with the Brynq REST API, polls a channel for new messages,
forwards conversation context to the LLM with Brynq tool definitions in OpenAI
function-calling format, executes any tool calls the LLM returns, and posts
text responses back to the channel.

Usage:
  python -m brynq.openai_adapter --api-url http://localhost:1234/v1 --model local-model --channel general
  python -m brynq.openai_adapter --api-url http://localhost:8080/v1 --model mistral --channel tasks --name task-bot

Only stdlib: urllib.request, json, argparse, time.
"""

import argparse
import json
import logging
import sys
import time
import urllib.error
import urllib.request
from typing import Any, Optional

logging.basicConfig(
    level=logging.INFO,
    format="%(asctime)s [%(levelname)s] %(message)s",
    datefmt="%H:%M:%S",
)
log = logging.getLogger("brynq.openai_adapter")

# ── System Prompt ────────────────────────────────────────────────────────────

SYSTEM_PROMPT = (
    "You are an AI agent connected to Brynq, a multi-agent coordination platform. "
    "You share channels with other AI agents (Claude, Gemini, Cursor, Codex, etc.) "
    "and human operators. You can read and post messages, manage tasks, list agents, "
    "and use other Brynq tools.\n\n"
    "Guidelines:\n"
    "- Be concise and helpful. Other agents and humans will read your messages.\n"
    "- Use tools when appropriate — post responses to channels, check inbox, list tasks.\n"
    "- When asked to do something, use the available Brynq tools to accomplish it.\n"
    "- If you receive a question or request from another agent, respond thoughtfully.\n"
    "- Coordinate with other agents when tasks require collaboration.\n"
)

# ── OpenAI-Compatible Client ────────────────────────────────────────────────


class OpenAICompatClient:
    """Minimal client for any OpenAI-compatible /v1/chat/completions endpoint."""

    def __init__(
        self,
        base_url: str = "http://localhost:1234/v1",
        model: str = "local-model",
        api_key: str = "not-needed",
    ) -> None:
        # Strip trailing slash so URL joins are clean
        self.base_url = base_url.rstrip("/")
        self.model = model
        self.api_key = api_key

    def chat_completion(
        self,
        messages: list[dict],
        tools: Optional[list[dict]] = None,
        tool_choice: str = "auto",
    ) -> dict:
        """Send a chat completion request and return the parsed response dict.

        Raises RuntimeError on HTTP or JSON errors.
        """
        url = f"{self.base_url}/chat/completions"

        payload: dict[str, Any] = {
            "model": self.model,
            "messages": messages,
        }
        if tools:
            payload["tools"] = tools
            payload["tool_choice"] = tool_choice

        body = json.dumps(payload).encode("utf-8")
        req = urllib.request.Request(
            url,
            data=body,
            headers={
                "Content-Type": "application/json",
                "Authorization": f"Bearer {self.api_key}",
            },
            method="POST",
        )

        try:
            with urllib.request.urlopen(req, timeout=120) as resp:
                return json.loads(resp.read().decode("utf-8"))
        except urllib.error.HTTPError as exc:
            detail = ""
            try:
                detail = exc.read().decode("utf-8", errors="replace")
            except Exception:
                pass
            raise RuntimeError(
                f"LLM API returned HTTP {exc.code}: {detail}"
            ) from exc
        except urllib.error.URLError as exc:
            raise RuntimeError(
                f"Cannot reach LLM API at {url}: {exc.reason}"
            ) from exc


# ── Brynq REST Client ───────────────────────────────────────────────────────


class BrynqClient:
    """Client for the Brynq REST API (api_server.py on port 9998)."""

    def __init__(
        self,
        base_url: str = "http://localhost:9998",
        api_key: Optional[str] = None,
    ) -> None:
        self.base_url = base_url.rstrip("/")
        self.api_key = api_key

    # ── Internal ──

    def _request(
        self,
        method: str,
        path: str,
        body: Optional[dict] = None,
        query: Optional[dict[str, str]] = None,
    ) -> dict:
        """Make an authenticated HTTP request to Brynq and return JSON."""
        url = f"{self.base_url}{path}"
        if query:
            qs = "&".join(f"{k}={urllib.request.quote(str(v))}" for k, v in query.items())
            url = f"{url}?{qs}"

        data = json.dumps(body).encode("utf-8") if body else None
        headers: dict[str, str] = {"Content-Type": "application/json"}
        if self.api_key:
            headers["X-API-Key"] = self.api_key

        req = urllib.request.Request(url, data=data, headers=headers, method=method)

        try:
            with urllib.request.urlopen(req, timeout=30) as resp:
                return json.loads(resp.read().decode("utf-8"))
        except urllib.error.HTTPError as exc:
            detail = ""
            try:
                detail = exc.read().decode("utf-8", errors="replace")
            except Exception:
                pass
            raise RuntimeError(
                f"Brynq API {method} {path} returned HTTP {exc.code}: {detail}"
            ) from exc
        except urllib.error.URLError as exc:
            raise RuntimeError(
                f"Cannot reach Brynq API at {url}: {exc.reason}"
            ) from exc

    # ── Registration ──

    def register(self, name: str, platform: str = "openai-compat") -> str:
        """Register this adapter as an agent and return the API key."""
        resp = self._request("POST", "/auth/register", body={
            "name": name,
            "platform": platform,
            "capabilities": ["chat", "tool-use"],
        })
        key = resp.get("api_key", "")
        if not key:
            raise RuntimeError(f"Registration failed: {resp}")
        self.api_key = key
        log.info("Registered with Brynq as '%s' (session: %s)", name, resp.get("session_id"))
        return key

    # ── Channel Operations ──

    def post_message(
        self,
        channel: str,
        message: str,
        msg_type: str = "info",
    ) -> dict:
        """Post a message to a Brynq channel."""
        return self._request("POST", f"/channels/{channel}/messages", body={
            "channel": channel,
            "message": message,
            "msg_type": msg_type,
        })

    def read_messages(
        self,
        channel: str,
        limit: int = 50,
        offset: int = 0,
    ) -> list[dict]:
        """Read message history from a channel."""
        resp = self._request("GET", f"/channels/{channel}/messages", query={
            "limit": str(limit),
            "offset": str(offset),
        })
        return resp.get("messages", [])

    def inbox(self) -> dict:
        """Check all channels for unread messages."""
        return self._request("GET", "/inbox")

    def mark_read(self, channel: str) -> dict:
        """Mark all messages in a channel as read."""
        return self._request("POST", f"/inbox/{channel}/mark_read")

    # ── Agent / Task Operations ──

    def list_agents(self) -> list[dict]:
        """List all registered agents."""
        resp = self._request("GET", "/agents")
        return resp.get("agents", [])

    def list_channels(self) -> list[dict]:
        """List all broker channels."""
        resp = self._request("GET", "/channels")
        return resp.get("channels", [])

    def list_tasks(self, channel: Optional[str] = None, status: Optional[str] = None) -> list[dict]:
        """List tasks with optional filters."""
        query: dict[str, str] = {}
        if channel:
            query["channel"] = channel
        if status:
            query["status"] = status
        resp = self._request("GET", "/tasks", query=query or None)
        return resp.get("tasks", [])

    def create_task(self, title: str, channel: str = "tasks", priority: str = "normal") -> dict:
        """Create a new task."""
        return self._request("POST", "/tasks", body={
            "title": title,
            "channel": channel,
            "priority": priority,
        })

    def get_stats(self) -> dict:
        """Get broker statistics."""
        return self._request("GET", "/stats")

    def call_tool(self, tool_name: str, arguments: dict) -> Any:
        """Dispatch a tool call to the appropriate Brynq REST endpoint.

        Maps OpenAI function names to Brynq REST API calls.
        Returns the result to be sent back to the LLM.
        """
        try:
            if tool_name == "brynq_post":
                return self.post_message(
                    channel=arguments.get("channel", "general"),
                    message=arguments["message"],
                    msg_type=arguments.get("msg_type", "info"),
                )
            elif tool_name == "brynq_read":
                return self.read_messages(
                    channel=arguments.get("channel", "general"),
                    limit=arguments.get("limit", 20),
                )
            elif tool_name == "brynq_inbox":
                return self.inbox()
            elif tool_name == "brynq_list_channels":
                return self.list_channels()
            elif tool_name == "brynq_list_agents":
                return self.list_agents()
            elif tool_name == "brynq_list_tasks":
                return self.list_tasks(
                    channel=arguments.get("channel"),
                    status=arguments.get("status"),
                )
            elif tool_name == "brynq_create_task":
                return self.create_task(
                    title=arguments["title"],
                    channel=arguments.get("channel", "tasks"),
                    priority=arguments.get("priority", "normal"),
                )
            elif tool_name == "brynq_stats":
                return self.get_stats()
            elif tool_name == "brynq_mark_read":
                return self.mark_read(channel=arguments.get("channel", "general"))
            else:
                return {"error": f"Unknown tool: {tool_name}"}
        except Exception as exc:
            return {"error": str(exc)}


# ── Tool Definitions (OpenAI Function Calling Format) ────────────────────────

BRYNQ_TOOLS: list[dict[str, Any]] = [
    {
        "type": "function",
        "function": {
            "name": "brynq_post",
            "description": (
                "Post a message to a Brynq channel. Other agents and terminals "
                "connected to the same broker will see it."
            ),
            "parameters": {
                "type": "object",
                "properties": {
                    "channel": {
                        "type": "string",
                        "description": "Channel name (e.g. 'general', 'tasks', 'alerts'). Auto-created if new.",
                        "default": "general",
                    },
                    "message": {
                        "type": "string",
                        "description": "The message content to post.",
                    },
                    "msg_type": {
                        "type": "string",
                        "description": "Message type: info, request, alert, status, question, file_ready.",
                        "default": "info",
                    },
                },
                "required": ["message"],
            },
        },
    },
    {
        "type": "function",
        "function": {
            "name": "brynq_read",
            "description": (
                "Read recent messages from a Brynq channel. Returns the latest messages "
                "in the channel."
            ),
            "parameters": {
                "type": "object",
                "properties": {
                    "channel": {
                        "type": "string",
                        "description": "Channel to read from.",
                        "default": "general",
                    },
                    "limit": {
                        "type": "integer",
                        "description": "Max number of messages to return.",
                        "default": 20,
                    },
                },
            },
        },
    },
    {
        "type": "function",
        "function": {
            "name": "brynq_inbox",
            "description": "Check all channels for unread messages. Quick overview of what you missed.",
            "parameters": {
                "type": "object",
                "properties": {},
            },
        },
    },
    {
        "type": "function",
        "function": {
            "name": "brynq_list_channels",
            "description": "List all broker channels with message counts and last activity.",
            "parameters": {
                "type": "object",
                "properties": {},
            },
        },
    },
    {
        "type": "function",
        "function": {
            "name": "brynq_list_agents",
            "description": "List all registered agents with their capabilities, platform, and status.",
            "parameters": {
                "type": "object",
                "properties": {},
            },
        },
    },
    {
        "type": "function",
        "function": {
            "name": "brynq_list_tasks",
            "description": "List tasks with optional filters by channel or status.",
            "parameters": {
                "type": "object",
                "properties": {
                    "channel": {
                        "type": "string",
                        "description": "Filter by channel (optional).",
                    },
                    "status": {
                        "type": "string",
                        "description": "Filter by status: open, accepted, completed (optional).",
                    },
                },
            },
        },
    },
    {
        "type": "function",
        "function": {
            "name": "brynq_create_task",
            "description": "Create a new task in the Brynq task system.",
            "parameters": {
                "type": "object",
                "properties": {
                    "title": {
                        "type": "string",
                        "description": "Task title.",
                    },
                    "channel": {
                        "type": "string",
                        "description": "Channel to post the task in.",
                        "default": "tasks",
                    },
                    "priority": {
                        "type": "string",
                        "description": "Priority: low, normal, high, critical.",
                        "default": "normal",
                    },
                },
                "required": ["title"],
            },
        },
    },
    {
        "type": "function",
        "function": {
            "name": "brynq_stats",
            "description": "View global broker statistics: total messages, storage, active sessions.",
            "parameters": {
                "type": "object",
                "properties": {},
            },
        },
    },
    {
        "type": "function",
        "function": {
            "name": "brynq_mark_read",
            "description": "Mark all messages in a channel as read.",
            "parameters": {
                "type": "object",
                "properties": {
                    "channel": {
                        "type": "string",
                        "description": "Channel to mark as read.",
                        "default": "general",
                    },
                },
            },
        },
    },
]


# ── Agent Loop ───────────────────────────────────────────────────────────────

MAX_HISTORY = 20


class AgentLoop:
    """Polls Brynq for messages and drives the LLM conversation loop."""

    def __init__(
        self,
        llm: OpenAICompatClient,
        brynq: BrynqClient,
        channel: str = "general",
        agent_name: str = "openai-agent",
        poll_interval: float = 5.0,
    ) -> None:
        self.llm = llm
        self.brynq = brynq
        self.channel = channel
        self.agent_name = agent_name
        self.poll_interval = poll_interval

        # Conversation history sent to the LLM (capped at MAX_HISTORY)
        self.history: list[dict[str, Any]] = [
            {"role": "system", "content": SYSTEM_PROMPT},
        ]
        # Track which messages we have already processed (by id)
        self._seen_ids: set[str] = set()

    def _trim_history(self) -> None:
        """Keep the system prompt plus the most recent MAX_HISTORY messages."""
        if len(self.history) > MAX_HISTORY + 1:
            self.history = [self.history[0]] + self.history[-(MAX_HISTORY):]

    def _fetch_new_messages(self) -> list[dict]:
        """Fetch channel messages and return only unseen ones not from this agent."""
        try:
            messages = self.brynq.read_messages(self.channel, limit=50)
        except Exception as exc:
            log.warning("Failed to fetch messages: %s", exc)
            return []

        new = []
        for msg in messages:
            msg_id = msg.get("id", "")
            if not msg_id or msg_id in self._seen_ids:
                continue
            # Skip our own messages to avoid echo loops
            if msg.get("session_name") == self.agent_name:
                self._seen_ids.add(msg_id)
                continue
            new.append(msg)
            self._seen_ids.add(msg_id)
        return new

    def _format_channel_message(self, msg: dict) -> str:
        """Format a Brynq channel message for the LLM context."""
        sender = msg.get("session_name", "unknown")
        platform = msg.get("platform", "?")
        content = msg.get("message", "")
        ts = msg.get("timestamp", "")[:19]
        return f"[{ts}] {sender}@{platform}: {content}"

    def _process_tool_calls(self, tool_calls: list[dict]) -> list[dict]:
        """Execute tool calls and return tool-role messages for the LLM."""
        results = []
        for tc in tool_calls:
            fn_name = tc.get("function", {}).get("name", "")
            raw_args = tc.get("function", {}).get("arguments", "{}")
            call_id = tc.get("id", "")

            try:
                arguments = json.loads(raw_args) if isinstance(raw_args, str) else raw_args
            except json.JSONDecodeError:
                arguments = {}

            log.info("Tool call: %s(%s)", fn_name, json.dumps(arguments, default=str))
            result = self.brynq.call_tool(fn_name, arguments)
            log.info("Tool result: %s", json.dumps(result, default=str)[:200])

            results.append({
                "role": "tool",
                "tool_call_id": call_id,
                "content": json.dumps(result, default=str),
            })
        return results

    def _call_llm(self) -> None:
        """Send the current history to the LLM, handle tool calls and text responses."""
        max_tool_rounds = 5
        for _ in range(max_tool_rounds):
            try:
                resp = self.llm.chat_completion(
                    messages=self.history,
                    tools=BRYNQ_TOOLS,
                    tool_choice="auto",
                )
            except RuntimeError as exc:
                log.error("LLM call failed: %s", exc)
                return

            choices = resp.get("choices", [])
            if not choices:
                log.warning("LLM returned empty choices")
                return

            message = choices[0].get("message", {})
            finish_reason = choices[0].get("finish_reason", "")

            # Append the assistant message to history
            self.history.append(message)

            tool_calls = message.get("tool_calls")
            if tool_calls:
                tool_results = self._process_tool_calls(tool_calls)
                self.history.extend(tool_results)
                # If finish_reason is tool_calls, loop to let LLM continue
                if finish_reason == "tool_calls":
                    continue

            # Extract text content
            content = message.get("content")
            if content and content.strip():
                log.info("LLM response: %s", content[:200])
                try:
                    self.brynq.post_message(self.channel, content.strip())
                except Exception as exc:
                    log.error("Failed to post response: %s", exc)

            # Done with this turn (either text response or no more tool calls)
            break

        self._trim_history()

    def run(self) -> None:
        """Main polling loop. Runs until interrupted."""
        log.info(
            "Agent '%s' monitoring #%s (poll every %ds)",
            self.agent_name, self.channel, self.poll_interval,
        )
        log.info("LLM: %s @ %s", self.llm.model, self.llm.base_url)

        # Seed the seen-IDs with existing messages to avoid replaying history on start
        try:
            existing = self.brynq.read_messages(self.channel, limit=50)
            for msg in existing:
                msg_id = msg.get("id", "")
                if msg_id:
                    self._seen_ids.add(msg_id)
            log.info("Skipped %d existing messages in #%s", len(existing), self.channel)
        except Exception as exc:
            log.warning("Could not pre-seed messages: %s", exc)

        # Mark channel as read so inbox starts clean
        try:
            self.brynq.mark_read(self.channel)
        except Exception:
            pass

        # Announce presence
        try:
            self.brynq.post_message(
                self.channel,
                f"{self.agent_name} is online and monitoring this channel.",
                msg_type="status",
            )
        except Exception as exc:
            log.warning("Could not post join announcement: %s", exc)

        while True:
            try:
                new_messages = self._fetch_new_messages()
                if new_messages:
                    # Add new messages to conversation history
                    for msg in new_messages:
                        formatted = self._format_channel_message(msg)
                        self.history.append({"role": "user", "content": formatted})
                        log.info("New message: %s", formatted[:120])

                    # Drive the LLM
                    self._call_llm()

                time.sleep(self.poll_interval)

            except KeyboardInterrupt:
                log.info("Shutting down (Ctrl+C)")
                try:
                    self.brynq.post_message(
                        self.channel,
                        f"{self.agent_name} is going offline.",
                        msg_type="status",
                    )
                except Exception:
                    pass
                break
            except Exception as exc:
                log.error("Loop error: %s", exc)
                time.sleep(self.poll_interval)


# ── CLI Entry Point ──────────────────────────────────────────────────────────


def main() -> None:
    parser = argparse.ArgumentParser(
        prog="brynq.openai_adapter",
        description=(
            "Bridge any OpenAI-compatible local LLM to Brynq. "
            "Works with LM Studio, vLLM, text-generation-webui, LocalAI, and more."
        ),
    )
    parser.add_argument(
        "--api-url",
        default="http://localhost:1234/v1",
        help="OpenAI-compatible API base URL (default: http://localhost:1234/v1)",
    )
    parser.add_argument(
        "--model",
        default="local-model",
        help="Model name to request (default: local-model)",
    )
    parser.add_argument(
        "--api-key",
        default="not-needed",
        help="API key for the LLM endpoint, if required (default: not-needed)",
    )
    parser.add_argument(
        "--channel",
        default="general",
        help="Brynq channel to monitor (default: general)",
    )
    parser.add_argument(
        "--name",
        default=None,
        help="Agent session name (default: openai-{model})",
    )
    parser.add_argument(
        "--brynq-url",
        default="http://localhost:9998",
        help="Brynq REST API URL (default: http://localhost:9998)",
    )
    parser.add_argument(
        "--poll-interval",
        type=float,
        default=5.0,
        help="Seconds between polling for new messages (default: 5)",
    )
    args = parser.parse_args()

    agent_name = args.name or f"openai-{args.model}"

    # Initialize clients
    llm = OpenAICompatClient(
        base_url=args.api_url,
        model=args.model,
        api_key=args.api_key,
    )

    brynq = BrynqClient(base_url=args.brynq_url)

    # Register with Brynq to get an API key
    log.info("Registering '%s' with Brynq at %s ...", agent_name, args.brynq_url)
    try:
        brynq.register(name=agent_name, platform=f"openai-compat/{args.model}")
    except RuntimeError as exc:
        log.error("Registration failed: %s", exc)
        sys.exit(1)

    # Verify LLM connectivity
    log.info("Testing LLM at %s (model: %s) ...", args.api_url, args.model)
    try:
        llm.chat_completion(
            messages=[{"role": "user", "content": "Say OK."}],
        )
        log.info("LLM connection verified")
    except RuntimeError as exc:
        log.error("Cannot reach LLM: %s", exc)
        sys.exit(1)

    # Start the agent loop
    loop = AgentLoop(
        llm=llm,
        brynq=brynq,
        channel=args.channel,
        agent_name=agent_name,
        poll_interval=args.poll_interval,
    )
    loop.run()


if __name__ == "__main__":
    main()
