agents/common/as_type.py

import json
import re
from typing import Any

from mcp_types import TextContent


def _as_dict(result: Any) -> dict:
    """Normalize MCP call_tool output to a dict."""
    if result is None:
        raise ValueError("MCP returned None")

    # Direct dict
    if isinstance(result, dict):
        return result

    if isinstance(result, str):
        return json.loads(result)

    if not isinstance(result, list):
        raise ValueError(f"Unexpected MCP response shape: {type(result)}")

    if not result:
        raise ValueError("MCP returned an empty list")

    # Find the first text block (dict or TextContent object)
    text: str | None = None
    for block in result:
        if isinstance(block, TextContent):
            text = block.text
            break
        if isinstance(block, dict) and block.get("type") == "text":
            text = block.get("text")
            break
        # Last-resort: the block itself is structured data
        if isinstance(block, dict) and "id" in block:
            return block

    if text is None:
        raise ValueError(f"No text block found in {repr(result)[:300]}")

    try:
        parsed = json.loads(text)
    except json.JSONDecodeError as e:
        raise ValueError(f"Text block is not JSON: {text[:300]}") from e

    if isinstance(parsed, dict):
        return parsed
    if isinstance(parsed, list) and parsed and isinstance(parsed[0], dict):
        return parsed[0]
    raise ValueError(f"Parsed JSON is neither dict nor list[dict]: {type(parsed)}")


def _as_list(result: Any) -> list[dict]:
    if result is None:
        return []
    if isinstance(result, list) and result and isinstance(result[0], TextContent):
        text = result[0].text
    elif isinstance(result, list) and result and isinstance(result[0], dict) and "type" in result[0]:
        text = result[0].get("text", "")
    elif isinstance(result, list):
        return result  # already a list of dicts
    elif isinstance(result, str):
        return json.loads(result)
    else:
        raise ValueError(f"Unexpected MCP response shape: {type(result)}")

    parsed = json.loads(text)
    if isinstance(parsed, list):
        return parsed
    if isinstance(parsed, dict):
        for key in ("items", "comments", "issues", "results"):
            if isinstance(parsed.get(key), list):
                return parsed[key]
        return [parsed]
    raise ValueError(f"Parsed JSON is not a list: {type(parsed)}")

def _first(d: dict, *keys, default=None):
    for k in keys:
        if k in d and d[k] is not None:
            return d[k]
    return default

def _extract_base64(result: Any) -> str | None:
    """Pull a base64 image payload out of an MCP call_tool result.

    Handles all the shapes a YouTrack MCP server might return:
      - ImageContent block (has .data and .mimeType)
      - TextContent block with raw base64 or a JSON wrapper
      - dict with {"data": "..."} or {"content": "..."}
      - plain string of base64
    """
    if result is None:
        return None

    # Plain string
    if isinstance(result, str):
        return result if _looks_like_base64(result) else None

    # Unwrap CallToolResult-like objects
    content = getattr(result, "content", None)
    if content is not None:
        return _extract_base64(content)

    # List of content blocks
    if isinstance(result, list):
        for block in result:
            got = _extract_base64(block)
            if got:
                return got
        return None

    # Dict wrapper
    if isinstance(result, dict):
        # Image content: {"type": "image", "data": "...", "mimeType": "..."}
        if result.get("type") == "image" and result.get("data"):
            return result["data"]

        # Direct data field
        for key in ("data", "content", "base64", "image"):
            if isinstance(result.get(key), str):
                val = result[key]
                if _looks_like_base64(val):
                    return val

        # Text wrapper whose text is base64 or JSON containing base64
        text = result.get("text")
        if isinstance(text, str):
            try:
                parsed = json.loads(text)
                return _extract_base64(parsed)
            except json.JSONDecodeError:
                return text if _looks_like_base64(text) else None

        return None

    # SDK TextContent / ImageContent objects (attribute access)
    if isinstance(result, TextContent):
        text = result.text
        try:
            return _extract_base64(json.loads(text))
        except json.JSONDecodeError:
            return text if _looks_like_base64(text) else None

    # ImageContent has .data and .mime_type (or .mimeType)
    if hasattr(result, "data") and hasattr(result, "type") and getattr(result, "type") == "image":
        return getattr(result, "data")

    return None


B64_RE = re.compile(r"^[A-Za-z0-9+/=\s]+$")
def _looks_like_base64(s: str) -> bool:
    """Very cheap check — base64 has no whitespace-heavy content and uses a limited alphabet."""
    s = s.strip()
    if len(s) < 32:
        return False
    return bool(B64_RE.match(s))

async def _fetch_image_base64(self, issue_id: str, attachment_id: str) -> str:
    result = await self.mcp.call_tool(
        "issue_attachment_download",
        {"issueId": issue_id, "attachmentId": attachment_id, "downloadToFile": False},
    )
    # Result contains base64 or a URL you can fetch
    return result  # base64 string

def _extract_text(result: Any) -> str:
    """Pull a single string out of an MCP result, whichever shape it is."""
    content = getattr(result, "content", None) or result
    if isinstance(content, list) and content:
        block = content[0]
        if isinstance(block, dict):
            return block.get("text") or json.dumps(block)
    if isinstance(content, dict):
        return content.get("text") or json.dumps(content)
    return str(content)


NOT_FOUND_HINTS = (
    "not found",
    "does not exist",
    "doesn't exist",
    "no such",
    "unknown issue",
    "404",
)


def _looks_like_not_found(text: str) -> bool:
    lowered = text.lower()
    return any(h in lowered for h in NOT_FOUND_HINTS)