"""Gmail data agent — A2A server using pydantic-ai.

Email access is provided by the `gws` CLI (https://github.com/googleworkspace/cli).
Make sure `gws auth login` has been completed before running this agent.
"""
import asyncio
import json
import os
import shlex
from pathlib import Path

from loguru import logger as _log
from pydantic_ai import RunContext

from backend.chanakya.a2a_context import RequestContext, make_metadata_aware_app
from backend.chanakya.base import EnterpriseAgent
from backend.chanakya.config import AGENT_REGISTRY
from pylogue.shell import app_factory as create_core_app

_cfg = AGENT_REGISTRY["gmail"]

gmail_agent = EnterpriseAgent(
    api_base=os.environ["LITELLM_PROVIDER_BASE_URL"],
    instructions=(
        "You are a Gmail data agent for Tata Steel's Enterprise Brain platform. "
        "You have access to the user's real Gmail inbox via the gws CLI. "
        "Answer questions about emails concisely. Always cite the email subject and date. "
        "When calling search_emails, pass short keyword terms only (e.g. 'coking coal'), never full sentences. "
        "Use read_email to fetch the full body of a specific message by ID. "
        "If a tool returns no results, respond immediately with a clear 'nothing found' message. Never retry with alternative queries."
    ),
    logfire_env="chanakya-alpha",
    service_name="gmail",
    deps_type=RequestContext,
)


async def _run_gws(*args: str) -> dict | list:
    """Run a gws CLI command asynchronously and return parsed JSON output."""
    cmd = ["gws", *args, "--format", "json"]
    _log.debug("gws cmd: {}", shlex.join(cmd))
    proc = await asyncio.create_subprocess_exec(
        *cmd,
        stdout=asyncio.subprocess.PIPE,
        stderr=asyncio.subprocess.PIPE,
    )
    stdout, stderr = await proc.communicate()
    if proc.returncode != 0:
        raise RuntimeError(f"gws error (rc={proc.returncode}): {stderr.decode().strip()}")
    text = stdout.decode().strip()
    if not text:
        return []
    try:
        return json.loads(text)
    except json.JSONDecodeError:
        # gws emits plain-text messages (e.g. "No messages found…") on stdout
        # with exit 0 when there are no results — treat as empty.
        _log.debug("gws non-JSON stdout (treating as empty): {!r}", text)
        return []


@gmail_agent.tool
async def search_emails(ctx: RunContext[RequestContext], query: str, max_results: int = 20) -> str:
    """Search the Gmail inbox using a keyword or Gmail search query.

    Returns a summary table of matching messages (id, from, subject, date).
    Use read_email to fetch the full body of any message.

    Args:
        query:       Gmail search query, e.g. 'coking coal', 'from:boss is:unread'.
        max_results: Maximum number of results to return (default 20).
    """
    _log.debug("search_emails caller={} query={!r}", ctx.deps.user_email, query)
    try:
        data = await _run_gws("gmail", "+triage", "--query", query, "--max", str(max_results))
    except RuntimeError as exc:
        _log.warning("search_emails gws error: {}", exc)
        return str(exc)

    if not data:
        return f"No emails found matching '{query}'."

    rows = data if isinstance(data, list) else data.get("messages", [])
    if not rows:
        return f"No emails found matching '{query}'."

    lines = ["id | from | subject | date"]
    for msg in rows:
        lines.append(
            f"{msg.get('id', '?')} | {msg.get('from', '?')} | {msg.get('subject', '?')} | {msg.get('date', '?')}"
        )
    return "\n".join(lines)


@gmail_agent.tool
async def read_email(ctx: RunContext[RequestContext], message_id: str) -> str:
    """Read the full body (and key headers) of a single Gmail message by ID.

    Args:
        message_id: The Gmail message ID returned by search_emails.
    """
    _log.debug("read_email caller={} id={!r}", ctx.deps.user_email, message_id)
    try:
        data = await _run_gws("gmail", "+read", "--id", message_id, "--headers")
    except RuntimeError as exc:
        _log.warning("read_email gws error: {}", exc)
        return str(exc)

    if isinstance(data, dict):
        headers = data.get("headers", {})
        body = data.get("body", "").strip()
        summary = (
            f"From: {headers.get('From', '?')}\n"
            f"To: {headers.get('To', '?')}\n"
            f"Subject: {headers.get('Subject', '?')}\n"
            f"Date: {headers.get('Date', '?')}\n\n"
            f"{body}"
        )
        return summary
    return str(data)


app = make_metadata_aware_app(
    gmail_agent,
    name=_cfg.name,
    description=_cfg.description,
)


@app.on_event("startup")
async def _startup_banner() -> None:
    port = os.environ.get("GMAIL_AGENT_PORT", "8001")
    _log.info("\n  Gmail Agent ready:")
    _log.info("    A2A   → http://0.0.0.0:{}/", port)
    _log.info("    Chat  → http://0.0.0.0:{}/chat", port)
    _log.info("    Card  → http://0.0.0.0:{}/.well-known/agent-card.json", port)
# FastAPI path-strips /chat so pylogue sees / internally.
# Access at: http://localhost:8001/chat
# In Docker, PYLOGUE_DB_PATH=/data/chat_history.db (bind-mounted from ./data/conversation-histories/gmail/).
_PROJECT_ROOT = Path(__file__).parents[3]
_gmail_db_path = os.environ.get(
    "PYLOGUE_DB_PATH",
    str(_PROJECT_ROOT / "data" / "conversation-histories" / "gmail" / "chat_history.db"),
)
Path(_gmail_db_path).parent.mkdir(parents=True, exist_ok=True)

app.mount(
    "/chat",
    create_core_app(
        responder_factory=lambda: gmail_agent,
        hero_title="Gmail Agent",
        hero_subtitle="Search and retrieve Tata Steel emails.",
        db_path=_gmail_db_path,
    ),
)


if __name__ == "__main__":
    import uvicorn
    uvicorn.run(_cfg.module_path, host=_cfg.host, port=_cfg.port, reload=True)
