"""
Phase 55: Service Discovery and Inter-Module Communication.

A lightweight service mesh for registering platform services, discovering
them by name or capability, invoking endpoints, and tracking call
statistics. Supports middleware injection for cross-cutting concerns
like logging and auth.

Storage:
  state/broker/services/
    {name}.json      — service registration data
    calls.jsonl      — call log (invocations, responses, errors)
"""

import json
import logging
import time
import uuid
from datetime import datetime, timezone
from pathlib import Path
from typing import Any, Callable, Dict, List, Optional

from brynq.server import BROKER_DIR, _write_json_sync, _read_jsonl_sync, _append_jsonl_sync

logger = logging.getLogger(__name__)

# ── Paths ──────────────────────────────────────────────────────────────────

SERVICES_DIR = BROKER_DIR / "services"
CALLS_LOG = SERVICES_DIR / "calls.jsonl"

SERVICES_DIR.mkdir(parents=True, exist_ok=True)

# ── In-memory registries ──────────────────────────────────────────────────

_endpoint_handlers: Dict[str, Dict[str, Callable]] = {}
_health_checks: Dict[str, Callable] = {}
_middleware: List[Dict[str, Any]] = []


def reset_mesh() -> None:
    """Reset all in-memory state. Used by tests."""
    global _endpoint_handlers, _health_checks, _middleware
    _endpoint_handlers = {}
    _health_checks = {}
    _middleware = []


# ── Internal helpers ──────────────────────────────────────────────────────

def _service_file(name: str) -> Path:
    return SERVICES_DIR / f"{name}.json"


def _read_service(name: str) -> Optional[dict]:
    path = _service_file(name)
    if not path.exists():
        return None
    try:
        return json.loads(path.read_text("utf-8"))
    except (json.JSONDecodeError, OSError):
        return None


def _save_service(name: str, data: dict) -> None:
    _write_json_sync(_service_file(name), data)


# ── Public API ────────────────────────────────────────────────────────────

def register_service(
    name: str,
    version: str,
    endpoints: Dict[str, Callable],
    health_check_fn: Optional[Callable] = None,
    capabilities: Optional[List[str]] = None,
) -> dict:
    """Register a platform service.

    Args:
        name: Unique service name.
        version: Semantic version string.
        endpoints: Dict mapping method names to handler callables.
        health_check_fn: Optional callable returning True/False for health.
        capabilities: Optional list of capability tags.

    Returns:
        The service registration record.
    """
    now = datetime.now(timezone.utc).isoformat()
    record = {
        "name": name,
        "version": version,
        "endpoint_names": list(endpoints.keys()),
        "capabilities": capabilities or [],
        "registered_at": now,
        "status": "active",
        "call_count": 0,
        "error_count": 0,
        "total_latency_ms": 0.0,
    }

    _save_service(name, record)
    _endpoint_handlers[name] = endpoints
    if health_check_fn:
        _health_checks[name] = health_check_fn

    logger.info("Registered service '%s' v%s with endpoints %s", name, version, list(endpoints.keys()))
    return record


def discover(
    name: Optional[str] = None,
    capability: Optional[str] = None,
) -> List[dict]:
    """Find services by name or capability.

    Args:
        name: Exact service name to look up.
        capability: Capability tag to search for.

    Returns:
        List of matching service records.
    """
    if name is not None:
        svc = _read_service(name)
        if svc and svc.get("status") == "active":
            return [svc]
        return []

    results = []
    for path in SERVICES_DIR.glob("*.json"):
        try:
            data = json.loads(path.read_text("utf-8"))
        except (json.JSONDecodeError, OSError):
            continue
        if data.get("status") != "active":
            continue
        if capability is not None:
            if capability not in data.get("capabilities", []):
                continue
        results.append(data)

    results.sort(key=lambda s: s.get("name", ""))
    return results


def call_service(
    name: str,
    method: str,
    payload: Optional[dict] = None,
) -> dict:
    """Invoke a service endpoint.

    Args:
        name: Service name.
        method: Method/endpoint name.
        payload: Optional dict payload.

    Returns:
        Dict with 'success', 'result' or 'error', and metadata.
    """
    svc = _read_service(name)
    if svc is None:
        raise ValueError(f"Service '{name}' not found")

    if name not in _endpoint_handlers or method not in _endpoint_handlers.get(name, {}):
        raise ValueError(f"Method '{method}' not found on service '{name}'")

    # Run middleware (before)
    call_ctx = {
        "service": name,
        "method": method,
        "payload": payload,
        "timestamp": datetime.now(timezone.utc).isoformat(),
    }
    for mw in _middleware:
        fn = mw.get("fn")
        if fn:
            try:
                fn("before", call_ctx)
            except Exception as e:
                logger.warning("Middleware '%s' error: %s", mw.get("name"), e)

    start_time = time.monotonic()
    call_id = uuid.uuid4().hex[:12]
    result_record: dict = {
        "call_id": call_id,
        "service": name,
        "method": method,
    }

    try:
        handler = _endpoint_handlers[name][method]
        result = handler(payload) if payload is not None else handler()
        elapsed_ms = (time.monotonic() - start_time) * 1000

        result_record["success"] = True
        result_record["result"] = result
        result_record["latency_ms"] = round(elapsed_ms, 2)

        # Update stats
        svc["call_count"] = svc.get("call_count", 0) + 1
        svc["total_latency_ms"] = svc.get("total_latency_ms", 0.0) + elapsed_ms
        _save_service(name, svc)

    except Exception as e:
        elapsed_ms = (time.monotonic() - start_time) * 1000
        result_record["success"] = False
        result_record["error"] = str(e)
        result_record["latency_ms"] = round(elapsed_ms, 2)

        svc["call_count"] = svc.get("call_count", 0) + 1
        svc["error_count"] = svc.get("error_count", 0) + 1
        svc["total_latency_ms"] = svc.get("total_latency_ms", 0.0) + elapsed_ms
        _save_service(name, svc)

    # Log the call
    _append_jsonl_sync(CALLS_LOG, {
        **result_record,
        "timestamp": datetime.now(timezone.utc).isoformat(),
    })

    # Run middleware (after)
    call_ctx["result"] = result_record
    for mw in _middleware:
        fn = mw.get("fn")
        if fn:
            try:
                fn("after", call_ctx)
            except Exception as e:
                logger.warning("Middleware '%s' after-error: %s", mw.get("name"), e)

    return result_record


def get_service_health(name: Optional[str] = None) -> Dict[str, Any]:
    """Check health of one or all services.

    Args:
        name: Optional service name. If None, checks all.

    Returns:
        Dict mapping service names to health status.
    """
    if name is not None:
        svc = _read_service(name)
        if svc is None:
            return {name: {"status": "not_found"}}
        healthy = True
        if name in _health_checks:
            try:
                healthy = bool(_health_checks[name]())
            except Exception:
                healthy = False
        return {name: {"status": "healthy" if healthy else "unhealthy", "active": svc.get("status") == "active"}}

    results = {}
    for path in SERVICES_DIR.glob("*.json"):
        try:
            data = json.loads(path.read_text("utf-8"))
        except (json.JSONDecodeError, OSError):
            continue
        svc_name = data.get("name", path.stem)
        healthy = True
        if svc_name in _health_checks:
            try:
                healthy = bool(_health_checks[svc_name]())
            except Exception:
                healthy = False
        results[svc_name] = {
            "status": "healthy" if healthy else "unhealthy",
            "active": data.get("status") == "active",
        }

    return results


def register_middleware(name: str, fn: Callable) -> None:
    """Add middleware that runs on every service call.

    The middleware fn receives (phase, context) where phase is 'before' or
    'after' and context is a dict with call details.
    """
    _middleware.append({"name": name, "fn": fn})
    logger.info("Registered middleware '%s'", name)


def get_service_stats(name: Optional[str] = None) -> Dict[str, Any]:
    """Get call stats for one or all services.

    Returns:
        Dict mapping service names to stats (calls, errors, avg_latency_ms).
    """
    if name is not None:
        svc = _read_service(name)
        if svc is None:
            return {}
        calls = svc.get("call_count", 0)
        errors = svc.get("error_count", 0)
        total_lat = svc.get("total_latency_ms", 0.0)
        avg_lat = round(total_lat / calls, 2) if calls > 0 else 0.0
        return {
            name: {
                "calls": calls,
                "errors": errors,
                "avg_latency_ms": avg_lat,
                "error_rate": round(errors / calls * 100, 2) if calls > 0 else 0.0,
            }
        }

    results = {}
    for path in SERVICES_DIR.glob("*.json"):
        try:
            data = json.loads(path.read_text("utf-8"))
        except (json.JSONDecodeError, OSError):
            continue
        svc_name = data.get("name", path.stem)
        calls = data.get("call_count", 0)
        errors = data.get("error_count", 0)
        total_lat = data.get("total_latency_ms", 0.0)
        avg_lat = round(total_lat / calls, 2) if calls > 0 else 0.0
        results[svc_name] = {
            "calls": calls,
            "errors": errors,
            "avg_latency_ms": avg_lat,
            "error_rate": round(errors / calls * 100, 2) if calls > 0 else 0.0,
        }

    return results


def deregister(name: str) -> bool:
    """Remove a service.

    Returns:
        True if the service was found and removed, False otherwise.
    """
    svc = _read_service(name)
    if svc is None:
        return False

    svc["status"] = "deregistered"
    _save_service(name, svc)

    # Clean up in-memory state
    _endpoint_handlers.pop(name, None)
    _health_checks.pop(name, None)

    logger.info("Deregistered service '%s'", name)
    return True
