from __future__ import annotations

import json
import time
from typing import Any, Optional

from mubit import Client


def _compact(payload: dict[str, Any]) -> dict[str, Any]:
    return {key: value for key, value in payload.items() if value is not None}


def _json_string(value: Any) -> Optional[str]:
    if value is None:
        return None
    return json.dumps(value)


def _default_signal(outcome: str) -> float:
    if outcome == "success":
        return 1.0
    if outcome == "failure":
        return -1.0
    return 0.0


OPS = {
    "ingest": {
        "key": "control.ingest",
        "summary": "Ingest memory items",
        "runIdField": "run_id",
        "http": {"method": "POST", "path": "/v2/control/ingest"},
    },
    "query": {
        "key": "control.query",
        "summary": "Query memories",
        "runIdField": "run_id",
        "http": {"method": "POST", "path": "/v2/control/query"},
    },
    "checkpoint": {
        "key": "control.checkpoint",
        "summary": "Create checkpoint",
        "runIdField": "run_id",
        "http": {"method": "POST", "path": "/v2/control/checkpoint"},
    },
    "record_outcome": {
        "key": "control.record_outcome",
        "summary": "Record outcome",
        "runIdField": "run_id",
        "http": {"method": "POST", "path": "/v2/control/outcome"},
    },
    "surface_strategies": {
        "key": "control.surface_strategies",
        "summary": "Surface strategies",
        "runIdField": "run_id",
        "http": {"method": "POST", "path": "/v2/control/strategies"},
    },
    "context": {
        "key": "control.context",
        "summary": "Get context",
        "runIdField": "run_id",
        "http": {"method": "POST", "path": "/v2/control/context"},
    },
    "archive_block": {
        "key": "control.archive_block",
        "summary": "Archive exact artifact block",
        "runIdField": "run_id",
        "http": {"method": "POST", "path": "/v2/control/archive"},
    },
    "dereference": {
        "key": "control.dereference",
        "summary": "Dereference exact artifact",
        "runIdField": "run_id",
        "http": {"method": "POST", "path": "/v2/control/dereference"},
    },
    "memory_health": {
        "key": "control.memory_health",
        "summary": "Memory health",
        "runIdField": "run_id",
        "http": {"method": "POST", "path": "/v2/control/memory_health"},
    },
    "diagnose": {
        "key": "control.diagnose",
        "summary": "Diagnose failures",
        "runIdField": "run_id",
        "http": {"method": "POST", "path": "/v2/control/diagnose"},
    },
    "register_agent": {
        "key": "control.register_agent",
        "summary": "Register agent",
        "runIdField": "run_id",
        "http": {"method": "POST", "path": "/v2/control/agents/register"},
    },
    "list_agents": {
        "key": "control.list_agents",
        "summary": "List agents",
        "runIdField": "run_id",
        "http": {"method": "POST", "path": "/v2/control/agents"},
    },
    "handoff": {
        "key": "control.create_handoff",
        "summary": "Create handoff",
        "runIdField": "run_id",
        "http": {"method": "POST", "path": "/v2/control/handoff"},
    },
    "feedback": {
        "key": "control.submit_feedback",
        "summary": "Submit feedback",
        "runIdField": "run_id",
        "http": {"method": "POST", "path": "/v2/control/feedback"},
    },
}


def normalize_metadata(evidence: dict[str, Any]) -> dict[str, Any]:
    metadata = evidence.get("metadata")
    if isinstance(metadata, dict):
        return metadata
    metadata_json = evidence.get("metadata_json")
    if isinstance(metadata_json, str) and metadata_json:
        try:
            parsed = json.loads(metadata_json)
            if isinstance(parsed, dict):
                return parsed
        except json.JSONDecodeError:
            return {}
    return {}


class MubitControlClient:
    def __init__(
        self,
        *,
        endpoint: str = "http://127.0.0.1:3000",
        api_key: str = "",
        sdk_client: Optional[Any] = None,
    ) -> None:
        self._client = sdk_client or Client(
            endpoint=endpoint.rstrip("/"), api_key=api_key, transport="http"
        )
        self._transport = self._client._transport

    def _wait_for_ingest(self, session_id: str, job_id: Optional[str]) -> None:
        if not job_id or not hasattr(self._client, "control"):
            return
        if not hasattr(self._client.control, "get_ingest_job"):
            return
        for _ in range(30):
            status = self._client.control.get_ingest_job(
                {"run_id": session_id, "job_id": job_id}
            )
            if status.get("done"):
                return
            time.sleep(0.05)
        raise TimeoutError(f"Timed out waiting for ingest job {job_id}")

    def ingest_items(
        self,
        *,
        session_id: str,
        agent_id: str,
        items: list[dict[str, Any]],
        idempotency_key: str = "",
    ) -> Any:
        serialized_items = []
        for item in items:
            serialized_items.append(
                _compact(
                    {
                        "item_id": item["item_id"],
                        "content_type": item.get("content_type", "text/plain"),
                        "text": item["text"],
                        "metadata_json": _json_string(item.get("metadata", {})),
                        "user_id": item.get("user_id"),
                        "intent": item.get("intent"),
                    }
                )
            )
        accepted = self._transport.invoke(
            OPS["ingest"],
            {
                "run_id": session_id,
                "agent_id": agent_id,
                "idempotency_key": idempotency_key,
                "parallel": False,
                "items": serialized_items,
            },
        )
        self._wait_for_ingest(session_id, accepted.get("job_id"))
        return accepted

    def recall(self, **payload: Any) -> Any:
        return self._transport.invoke(
            OPS["query"],
            _compact(
                {
                    "run_id": payload["session_id"],
                    "query": payload["query"],
                    "user_id": payload.get("user_id"),
                    "limit": payload.get("limit", 10),
                    "include_working_memory": payload.get(
                        "include_working_memory", True
                    ),
                }
            ),
        )

    def checkpoint(self, **payload: Any) -> Any:
        return self._transport.invoke(
            OPS["checkpoint"],
            _compact(
                {
                    "run_id": payload["session_id"],
                    "context_snapshot": payload["snapshot"],
                    "label": payload.get("label"),
                    "metadata_json": _json_string(payload.get("metadata")),
                    "agent_id": payload.get("agent_id", "adk-agent"),
                    "user_id": payload.get("user_id"),
                }
            ),
        )

    def record_outcome(self, **payload: Any) -> Any:
        return self._transport.invoke(
            OPS["record_outcome"],
            _compact(
                {
                    "run_id": payload["session_id"],
                    "reference_id": payload["reference_id"],
                    "outcome": payload["outcome"],
                    "signal": payload.get("signal", _default_signal(payload["outcome"])),
                    "rationale": payload.get("rationale"),
                    "agent_id": payload.get("agent_id", "adk-agent"),
                    "user_id": payload.get("user_id"),
                }
            ),
        )

    def surface_strategies(self, **payload: Any) -> Any:
        return self._transport.invoke(
            OPS["surface_strategies"],
            _compact(
                {
                    "run_id": payload["session_id"],
                    "lesson_types": payload.get("lesson_types"),
                    "max_strategies": payload.get("max_strategies", 5),
                    "user_id": payload.get("user_id"),
                }
            ),
        )

    def get_context(self, **payload: Any) -> Any:
        return self._transport.invoke(
            OPS["context"],
            _compact(
                {
                    "run_id": payload["session_id"],
                    "query": payload["query"],
                    "user_id": payload.get("user_id"),
                    "agent_id": payload.get("agent_id"),
                    "mode": payload.get("mode"),
                    "sections": payload.get("sections"),
                    "entry_types": payload.get("entry_types"),
                    "limit": payload.get("limit"),
                    "max_token_budget": payload.get("max_token_budget"),
                    "include_working_memory": True,
                }
            ),
        )

    def archive(self, **payload: Any) -> Any:
        return self._transport.invoke(
            OPS["archive_block"],
            _compact(
                {
                    "run_id": payload["session_id"],
                    "content": payload["content"],
                    "artifact_kind": payload["artifact_kind"],
                    "metadata_json": _json_string(payload.get("metadata")),
                    "user_id": payload.get("user_id"),
                    "agent_id": payload.get("agent_id"),
                    "origin_agent_id": payload.get("origin_agent_id")
                    or payload.get("agent_id"),
                    "source_attempt_id": payload.get("source_attempt_id"),
                    "source_tool": payload.get("source_tool"),
                    "labels": payload.get("labels"),
                    "family": payload.get("family"),
                    "importance": payload.get("importance"),
                }
            ),
        )

    def dereference(self, **payload: Any) -> Any:
        return self._transport.invoke(
            OPS["dereference"],
            _compact(
                {
                    "run_id": payload["session_id"],
                    "reference_id": payload["reference_id"],
                    "user_id": payload.get("user_id"),
                    "agent_id": payload.get("agent_id"),
                }
            ),
        )

    def memory_health(self, **payload: Any) -> Any:
        return self._transport.invoke(
            OPS["memory_health"],
            _compact(
                {
                    "run_id": payload["session_id"],
                    "user_id": payload.get("user_id"),
                    "limit": payload.get("limit", 100),
                }
            ),
        )

    def diagnose(self, **payload: Any) -> Any:
        return self._transport.invoke(
            OPS["diagnose"],
            _compact(
                {
                    "run_id": payload["session_id"],
                    "error_text": payload["error_text"],
                    "error_type": payload.get("error_type"),
                    "user_id": payload.get("user_id"),
                    "limit": payload.get("limit", 10),
                }
            ),
        )

    def register_agent(self, **payload: Any) -> Any:
        return self._transport.invoke(
            OPS["register_agent"],
            _compact(
                {
                    "run_id": payload["session_id"],
                    "agent_id": payload["agent_id"],
                    "role": payload.get("role"),
                    "status": payload.get("status"),
                    "read_scopes": payload.get("read_scopes"),
                    "write_scopes": payload.get("write_scopes"),
                    "shared_memory_lanes": payload.get("shared_memory_lanes"),
                    "capabilities_json": _json_string(payload.get("capabilities")),
                    "user_id": payload.get("user_id"),
                }
            ),
        )

    def list_agents(self, **payload: Any) -> Any:
        return self._transport.invoke(
            OPS["list_agents"],
            _compact(
                {
                    "run_id": payload["session_id"],
                    "user_id": payload.get("user_id"),
                }
            ),
        )

    def handoff(self, **payload: Any) -> Any:
        return self._transport.invoke(
            OPS["handoff"],
            _compact(
                {
                    "run_id": payload["session_id"],
                    "from_agent_id": payload["from_agent_id"],
                    "to_agent_id": payload["to_agent_id"],
                    "content": payload["content"],
                    "requested_action": payload.get("requested_action"),
                    "metadata_json": _json_string(payload.get("metadata")),
                    "user_id": payload.get("user_id"),
                }
            ),
        )

    def feedback(self, **payload: Any) -> Any:
        return self._transport.invoke(
            OPS["feedback"],
            _compact(
                {
                    "run_id": payload["session_id"],
                    "handoff_id": payload["handoff_id"],
                    "verdict": payload["verdict"],
                    "comments": payload.get("comments"),
                    "from_agent_id": payload.get("from_agent_id"),
                    "metadata_json": _json_string(payload.get("metadata")),
                    "user_id": payload.get("user_id"),
                }
            ),
        )
