"""
Auditor MCP Server (FastMCP)
- Tools:
  - log_instances(image_path, instances)  [DISABLED - no longer writes to artifacts/logs]
  - log_output(image_path, output_path)   [DISABLED - JSON saved next to output]
  - audit_output(image_path, output_path)  # verifies residual PII on the final image via Analyzer
"""

from __future__ import annotations

import json
from pathlib import Path
from datetime import datetime, timezone
import sys
from typing import Any, Dict
import platform
import os


# LOG_DIR disabled to avoid cluttering artifacts/logs - all logging done in main JSON
# LOG_DIR = Path("artifacts/logs")
# LOG_DIR.mkdir(parents=True, exist_ok=True)


def _run_without_mcp() -> None:
  print(
    "Auditor MCP server requires the 'mcp' package.\n"
    "Install with: pip install mcp\n"
    "Then run: python -m mcp.servers.auditor_server"
  )


try:
  from mcp.server.fastmcp import FastMCP
except Exception:
  FastMCP = None  # type: ignore[misc,assignment]

# Optional MCP client SDK for calling Analyzer from within this server
try:
  import asyncio
  from mcp.client.session import ClientSession  # type: ignore
  from mcp.client.stdio import stdio_client, StdioServerParameters  # type: ignore
except Exception:  # pragma: no cover
  ClientSession = None  # type: ignore
  stdio_client = None  # type: ignore
  StdioServerParameters = None  # type: ignore

def _unwrap_tool_result(res: Any) -> Dict[str, Any]:
  if isinstance(res, dict):
    return res
  try:
    content = getattr(res, "content", None)
    if not content:
      return {}
    for c in content:
      if isinstance(c, dict):
        ctype = c.get("type")
        if ctype in ("json", "application/json"):
          data = c.get("json") or c.get("data") or c.get("value")
          if isinstance(data, dict):
            return data
        if ctype == "text" and isinstance(c.get("text"), str):
          import json as _json
          txt = c.get("text")
          try:
            parsed = _json.loads(txt)  # type: ignore[arg-type]
            if isinstance(parsed, dict):
              return parsed
          except Exception:
            continue
      else:
        ctype = getattr(c, "type", None)
        if ctype in ("json", "application/json"):
          data = getattr(c, "json", None) or getattr(c, "data", None) or getattr(c, "value", None)
          if isinstance(data, dict):
            return data
        if ctype == "text":
          import json as _json
          txt = getattr(c, "text", None)
          if isinstance(txt, str):
            try:
              parsed = _json.loads(txt)
              if isinstance(parsed, dict):
                return parsed
            except Exception:
              continue
  except Exception:
    return {}
  return {}

async def _call_analyzer(image_path: str) -> Dict[str, Any]:
  if ClientSession is None or stdio_client is None or StdioServerParameters is None:
    return {"instances": [], "error": "MCP client SDK not available in auditor_server"}
  params = StdioServerParameters(
    command=sys.executable,
    args=["-m", "mcp.servers.analyzer_server"],
  )
  async with stdio_client(params) as (read, write):
    async with ClientSession(read, write) as session:
      await session.initialize()
      res = await session.call_tool("classify_pii", {"image_path": image_path})
      return _unwrap_tool_result(res)


if FastMCP is None:
  def main() -> None:  # noqa: D401
    _run_without_mcp()
else:
  app = FastMCP("AuditorServer")

  def _ts() -> str:
    return datetime.now(timezone.utc).isoformat()

  def _workstation_info() -> Dict[str, Any]:
    info: Dict[str, Any] = {
      "os": platform.platform(),
      "python": platform.python_version(),
      "machine": platform.machine(),
      "processor": platform.processor(),
    }
    # Try to add GPU/VRAM hints when available
    try:
      import torch  # type: ignore
      if torch.cuda.is_available():
        gname = torch.cuda.get_device_name(0)
        vram = torch.cuda.get_device_properties(0).total_memory  # bytes
        info["gpu"] = {"name": gname, "vram_bytes": int(vram)}
      else:
        info["gpu"] = {"available": False}
    except Exception:
      # Fallback to environment hints
      info["gpu"] = {
        "available": bool(os.environ.get("CUDA_VISIBLE_DEVICES")),
      }
    return info

  @app.tool()
  def log_instances(image_path: str, instances: Dict[str, Any]) -> Dict[str, bool]:  # noqa: D401
    """Log detected PII instances. (Disabled - no longer writes to artifacts/logs)"""
    # Disabled to avoid cluttering artifacts/logs - instances are saved in main JSON
    return {"ok": True}

  @app.tool()
  def log_output(image_path: str, output_path: str) -> Dict[str, bool]:  # noqa: D401
    # No centralized logging - JSON log is saved next to output as {output_stem}.json
    return {"ok": True}

  @app.tool()
  async def audit_output(image_path: str, output_path: str) -> Dict[str, Any]:  # noqa: D401
    """Verify that the final image has no residual PII using the Analyzer server.

    Returns
    -------
    {"ok": bool, "count": int, "residual": list, "analyzer_meta": dict, "baseline_subset": bool,
     "baseline_count": int, "error"?: str}
    Also appends an entry with "audit" to outputs.jsonl for traceability.
    """
    op = Path(output_path)
    if not op.exists():
      result = {"ok": False, "count": 0, "residual": [], "analyzer_meta": {}, "baseline_subset": False, "baseline_count": 0, "error": f"output not found: {output_path}"}
      return result

    # Run analyzer on the output image
    try:
      analysis = await _call_analyzer(str(op))
    except Exception as e:  # pragma: no cover
      analysis = {"instances": [], "error": f"audit analyzer failed: {e}"}

    residual = analysis.get("instances", []) if isinstance(analysis, dict) else []
    meta = analysis.get("meta", {}) if isinstance(analysis, dict) else {}
    ok = isinstance(residual, list) and len(residual) == 0 and not analysis.get("error")

    # Helper for instance signature
    def _sig(v: Any) -> str:
      """Best-effort stable signature for an instance.
      Combines common fields like type/category/label/text. Geometry is ignored.
      """
      if not isinstance(v, dict):
        return str(v)
      # Prefer most specific identifiers first
      keys = [
        v.get("id"),
        v.get("uuid"),
        v.get("signature"),
        v.get("type"), v.get("category"), v.get("class"), v.get("label"), v.get("name"), v.get("kind"),
        v.get("text"), v.get("value"), v.get("content"),
      ]
      parts = [str(k).strip().lower() for k in keys if k is not None]
      if not parts:
        # Fallback: sorted key=value for a few safe fields
        fields = {k: v[k] for k in ("type", "category", "class", "label", "name", "kind", "text") if k in v}
        parts = [f"{k}={str(val).strip().lower()}" for k, val in sorted(fields.items())]
      return "|".join(parts)

    # Baseline comparison disabled (no longer reading from artifacts/logs/instances.jsonl)
    # All data is now in the main JSON file next to the output
    baseline_instances = []
    baseline_set = set()
    residual_set = { _sig(x) for x in residual if x is not None }
    baseline_subset = False  # No baseline to compare against

    report: Dict[str, Any] = {
      "ok": bool(ok),
      "count": len(residual) if isinstance(residual, list) else 0,
      "residual": residual if isinstance(residual, list) else [],
      "analyzer_meta": meta,
      "baseline_subset": bool(baseline_subset),
      "baseline_count": len(baseline_set),
    }
    if analysis.get("error"):
      report["error"] = analysis.get("error")

    # Logging to artifacts/logs disabled - all audit info is in main JSON
    # (No longer writing to artifacts/logs/outputs.jsonl)
    
    return report

  def main() -> None:  # noqa: D401
    print("[AuditorServer] Starting (transport=stdio)…", file=sys.stderr, flush=True)
    try:
      app.run(transport="stdio")
    except BrokenPipeError:
      # Client closed connection during shutdown - this is expected behavior
      pass
    except KeyboardInterrupt:
      # Graceful shutdown on Ctrl+C
      pass
    except BaseException as e:
      # Check if it's an ExceptionGroup containing BrokenPipeError
      if hasattr(e, '__cause__') or hasattr(e, '__context__'):
        cause = getattr(e, '__cause__', None) or getattr(e, '__context__', None)
        if isinstance(cause, BrokenPipeError):
          pass  # Expected shutdown
          return
      # Check for ExceptionGroup with BrokenPipeError
      if e.__class__.__name__ in ('ExceptionGroup', 'BaseExceptionGroup'):
        try:
          exceptions = getattr(e, 'exceptions', [])
          if any(isinstance(ex, BrokenPipeError) or 
                 'BrokenPipeError' in str(type(ex)) or
                 'Broken pipe' in str(ex) for ex in exceptions):
            pass  # Expected shutdown
            return
        except Exception:
          pass
      print(f"[AuditorServer] Unexpected error: {e}", file=sys.stderr, flush=True)
      raise


if __name__ == "__main__":
  main()
