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"},
    },
    "context": {
        "key": "control.context",
        "summary": "Get context",
        "runIdField": "run_id",
        "http": {"method": "POST", "path": "/v2/control/context"},
    },
    "delete_run": {
        "key": "control.delete_run",
        "summary": "Delete run",
        "runIdField": "run_id",
        "http": {"method": "POST", "path": "/v2/control/delete_run"},
    },
    "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"},
    },
}


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 remember(self, **payload: Any) -> Any:
        accepted = self._transport.invoke(
            OPS["ingest"],
            {
                "run_id": payload["session_id"],
                "agent_id": payload.get("agent_id", "langchain"),
                "parallel": False,
                "items": [
                    _compact(
                        {
                            "item_id": f"lc-{payload.get('upsert_key', 'memory')}",
                            "content_type": "text/plain",
                            "text": payload["text"],
                            "user_id": payload.get("user_id"),
                            "intent": payload.get("intent"),
                            "source": payload.get("source"),
                            "metadata_json": _json_string(payload.get("metadata", {})),
                            "upsert_key": payload.get("upsert_key"),
                        }
                    )
                ],
            },
        )
        self._wait_for_ingest(payload["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),
                    "entry_types": payload.get("entry_types"),
                    "include_working_memory": True,
                    "agent_id": payload.get("agent_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 delete_run(self, **payload: Any) -> Any:
        return self._transport.invoke(
            OPS["delete_run"],
            _compact(
                {
                    "run_id": payload["session_id"],
                }
            ),
        )
