import os
import time
import json
from typing import Any, Callable, Dict, Optional

import requests


def _env(name: str, default: str = "") -> str:
    v = os.getenv(name)
    return v if v is not None else default


def _safe_json(value: Any) -> Any:
    if value is None:
        return None
    if isinstance(value, (str, int, float, bool)):
        return value
    if isinstance(value, (list, tuple)):
        return [_safe_json(v) for v in value]
    if isinstance(value, dict):
        return {str(k): _safe_json(v) for k, v in value.items()}
    return str(value)


class KybernisClient:
    def __init__(
        self,
        base_url: Optional[str] = None,
        api_token: Optional[str] = None,
        tenant_id: Optional[str] = None,
        timeout_s: int = 10,
    ) -> None:
        self.base_url = (base_url or _env("KYB_API_BASE", "https://api.kybernis.io")).rstrip("/")
        self.api_token = api_token or _env("KYB_API_KEY", "") or _env("KYBERNIS_ADMIN_TOKEN", "")
        self.tenant_id = tenant_id or _env("KYB_TENANT_ID", "") or _env("DEMO_TENANT_ID", "")
        self.timeout_s = timeout_s
        # Debug: surface client configuration (don't print secrets)
        try:
            print(f"[kybernis] client init base_url={self.base_url} timeout={self.timeout_s} api_token_present={bool(self.api_token)} tenant_id={self.tenant_id}", flush=True)
        except Exception:
            pass

    def _headers(self) -> Dict[str, str]:
        headers = {"Content-Type": "application/json"}
        if self.api_token:
            headers["Authorization"] = f"Bearer {self.api_token}"
        if self.tenant_id:
            headers["X-Kyb-Tenant-Id"] = self.tenant_id
        return headers

    def _post(self, path: str, payload: Dict[str, Any]) -> Dict[str, Any]:
        url = f"{self.base_url}{path}"
        resp = requests.post(url, headers=self._headers(), data=json.dumps(payload), timeout=self.timeout_s)
        try:
            resp.raise_for_status()
        except Exception:
            # Print server body to debug 500s
            try:
                print(f"[kybernis] HTTP {resp.status_code} {url} body={resp.text}", flush=True)
            except Exception:
                pass
            raise
        return resp.json()

    def _get(self, path: str, params: Dict[str, Any]) -> Dict[str, Any]:
        url = f"{self.base_url}{path}"
        resp = requests.get(url, headers=self._headers(), params=params, timeout=self.timeout_s)
        resp.raise_for_status()
        return resp.json()

    def tool_register(self, name: str, kind: str = "tool", endpoint: str = "", labels: Optional[Dict[str, str]] = None) -> str:
        payload = {
            "tenant_id": self.tenant_id,
            "name": name,
            "kind": kind,
            "endpoint": endpoint,
            "labels": labels or {},
        }
        data = self._post("/infra/tools/register", payload)
        return data.get("data", {}).get("tool_id", "")

    def tool_heartbeat(self, tool_id: str) -> None:
        self._post("/infra/tools/heartbeat", {"tool_id": tool_id})

    def tool_report(self, tool_id: str, status: str = "UNKNOWN", circuit_state: str = "CLOSED", failure_inc: int = 0) -> None:
        payload = {
            "tool_id": tool_id,
            "status": status,
            "circuit_state": circuit_state,
            "failure_inc": failure_inc,
            "checked_at": time.strftime("%Y-%m-%dT%H:%M:%SZ", time.gmtime()),
        }
        self._post("/infra/tools/report", payload)

    def emit_event(self, run_id: str, event_type: str, payload: Optional[Dict[str, Any]] = None, parent_id: str = "") -> None:
        payload = payload or {}
        body = {
            "run_id": run_id,
            "parent_id": parent_id,
            "type": event_type,
            "payload": payload,
        }
        # Debug: ensure we at least print the emission attempt
        try:
            print("[kybernis] EMIT", run_id, event_type, payload, flush=True)
        except Exception:
            pass
        self._post("/infra/graph/event", body)

    def record_cost(self, run_id: str, cost_micros: int, currency: str = "USD", source: str = "sdk") -> None:
        body = {
            "tenant_id": self.tenant_id,
            "run_id": run_id,
            "cost_micros": int(cost_micros),
            "currency": currency,
            "source": source,
        }
        self._post("/infra/cost/record", body)

    def gate_request(
        self,
        run_id: str,
        tool_name: str,
        action: str = "call",
        tool_id: str = "",
        payload: Optional[Dict[str, Any]] = None,
        policy: Optional[Dict[str, Any]] = None,
        reason: str = "",
    ) -> Dict[str, Any]:
        body = {
            "tenant_id": self.tenant_id,
            "run_id": run_id,
            "tool_id": tool_id,
            "tool_name": tool_name,
            "action": action,
            "reason": reason,
            "payload": payload or {},
            "policy": policy or {},
        }
        return self._post("/infra/gate/request", body)

    def gate_request_span(
        self,
        run_id: str,
        span_id: str,
        tool_name: str,
        action: str = "call",
        tool_id: str = "",
        payload: Optional[Dict[str, Any]] = None,
        policy: Optional[Dict[str, Any]] = None,
        reason: str = "",
    ) -> Dict[str, Any]:
        payload = payload or {}
        payload = {**payload, "span_id": span_id}
        res = self.gate_request(
            run_id=run_id,
            tool_name=tool_name,
            action=action,
            tool_id=tool_id,
            payload=payload,
            policy=policy,
            reason=reason,
        )
        gate_id = (res.get("data", {}) or {}).get("gate_id", "")
        self.emit_event(run_id, "gate.request", {"gate_id": gate_id, "span_id": span_id, "tool": tool_name})
        return res

    def gate_status(self, gate_id: str) -> Dict[str, Any]:
        return self._get("/infra/gate/status", {"gate_id": gate_id})

    def list_gates(
        self,
        *,
        run_id: str = "",
        status: str = "",
        limit: int = 100,
    ) -> Dict[str, Any]:
        params: Dict[str, Any] = {
            "tenant_id": self.tenant_id,
            "limit": int(limit),
        }
        if run_id:
            params["run_id"] = run_id
        if status:
            params["status"] = status
        return self._get("/infra/gates", params)

    def gate_wait(
        self,
        gate_id: str,
        timeout_s: int = 30,
        poll_interval_s: float = 1.0,
        span_id: str = "",
        run_id: str = "",
    ) -> Dict[str, Any]:
        deadline = time.time() + timeout_s
        last = {}
        while time.time() < deadline:
            last = self.gate_status(gate_id)
            status = (last.get("data", {}) or {}).get("status", "")
            if status and status != "PENDING":
                if span_id and run_id:
                    self.emit_event(run_id, "gate.decision", {"gate_id": gate_id, "status": status, "span_id": span_id})
                return last
            time.sleep(poll_interval_s)
        return last
    
    # kybernis/client.py
    def get_or_create_run(self, input_context_ref: str, workflow_version_id: str = "") -> str:
        body = {
            "tenant_id": self.tenant_id,
            "input_context_ref": input_context_ref,
        }
        if workflow_version_id:
            body["workflow_version_id"] = workflow_version_id
        data = self._post("/infra/runs/get_or_create", body)
        run_id = (data.get("data", {}) or {}).get("run_id", "")
        if not run_id:
            raise RuntimeError(f"missing run_id in response: {data}")
        return run_id

    def task_run_start(
        self,
        *,
        run_id: str,
        task_run_id: str,
        kind: str,
        name: str,
        idempotency_key: str,
        attributes: Optional[Dict[str, Any]] = None,
        attempt: int = 1,
    ) -> Dict[str, Any]:
        body = {
            "tenant_id": self.tenant_id,
            "run_id": run_id,
            "task_run_id": task_run_id,
            "kind": kind,
            "name": name,
            "idempotency_key": idempotency_key,
            "attributes": attributes or {},
            "attempt": int(attempt),
        }
        data = self._post("/infra/task_runs/start", body)
        return data.get("data", {}) or {}

    def task_run_finish(
        self,
        *,
        run_id: str,
        task_run_id: str,
        status: str,
        result: Optional[Dict[str, Any]] = None,
        error: str = "",
        error_code: str = "",
    ) -> Dict[str, Any]:
        body = {
            "tenant_id": self.tenant_id,
            "run_id": run_id,
            "task_run_id": task_run_id,
            "status": status,
            "result": result or {},
            "error": error,
            "error_code": error_code,
            "ended_at": time.time(),
        }
        data = self._post("/infra/task_runs/finish", body)
        return data.get("data", {}) or {}

    def task_run_artifact(
        self,
        *,
        run_id: str,
        task_run_id: str,
        kind: str,
        payload: Optional[Dict[str, Any]] = None,
    ) -> Dict[str, Any]:
        body = {
            "tenant_id": self.tenant_id,
            "run_id": run_id,
            "task_run_id": task_run_id,
            "kind": kind,
            "payload": payload or {},
        }
        data = self._post("/infra/task_runs/artifacts", body)
        return data.get("data", {}) or {}

    def intent_guard(
        self,
        *,
        run_id: str,
        idempotency_key: str,
        intent: Dict[str, Any],
        original_payload: Optional[Dict[str, Any]] = None,
        requested_mode: str = "",
    ) -> Dict[str, Any]:
        body = {
            "tenant_id": self.tenant_id,
            "run_id": run_id,
            "idempotency_key": idempotency_key,
            "intent": intent or {},
            "original_payload": original_payload if original_payload is not None else {"intent": intent or {}},
            "requested_mode": requested_mode,
        }
        data = self._post("/infra/intent/guard", body)
        return data.get("data", {}) or {}

    def intent_commit(
        self,
        *,
        mutation_id: str,
        runtime_intent_hash: str = "",
        intent: Optional[Dict[str, Any]] = None,
        receipt: Optional[Dict[str, Any]] = None,
        provider: str = "sdk",
    ) -> None:
        body = {
            "tenant_id": self.tenant_id,
            "mutation_id": mutation_id,
            "runtime_intent_hash": runtime_intent_hash,
            "intent": intent or {},
            "receipt": receipt or {},
            "provider": provider,
        }
        self._post("/infra/intent/commit", body)

    def intent_fail(self, *, mutation_id: str, reason: str) -> None:
        body = {
            "tenant_id": self.tenant_id,
            "mutation_id": mutation_id,
            "reason": reason,
        }
        self._post("/infra/intent/fail", body)

    def intent_unknown(self, *, mutation_id: str, reason: str) -> None:
        body = {
            "tenant_id": self.tenant_id,
            "mutation_id": mutation_id,
            "reason": reason,
        }
        self._post("/infra/intent/unknown", body)

    def intent_reconcile(
        self,
        *,
        mutation_id: str,
        outcome: str,
        runtime_intent_hash: str = "",
        receipt: Optional[Dict[str, Any]] = None,
        provider: str = "reconciler",
        reason: str = "",
    ) -> Dict[str, Any]:
        body = {
            "tenant_id": self.tenant_id,
            "mutation_id": mutation_id,
            "outcome": outcome,
            "runtime_intent_hash": runtime_intent_hash,
            "receipt": receipt or {},
            "provider": provider,
            "reason": reason,
        }
        data = self._post("/infra/intent/reconcile", body)
        return data.get("data", {}) or {}


def kyb_tool(
    client: KybernisClient,
    run_id: str,
    tool_name: str,
    action: str = "call",
    policy: Optional[Dict[str, Any]] = None,
    block: bool = True,
    timeout_s: int = 30,
    tool_id: str = "",
    on_deny: str = "raise",
    span_id: str = "",
) -> Callable[[Callable[..., Any]], Callable[..., Any]]:
    def decorator(fn: Callable[..., Any]) -> Callable[..., Any]:
        def wrapper(*args: Any, **kwargs: Any) -> Any:
            payload = {"args": _safe_json(args), "kwargs": _safe_json(kwargs)}
            if span_id:
                gate = client.gate_request_span(
                    run_id=run_id,
                    span_id=span_id,
                    tool_name=tool_name,
                    action=action,
                    tool_id=tool_id,
                    payload=payload,
                    policy=policy,
                )
            else:
                gate = client.gate_request(
                    run_id=run_id,
                    tool_name=tool_name,
                    action=action,
                    tool_id=tool_id,
                    payload=payload,
                    policy=policy,
                )
            status = (gate.get("data", {}) or {}).get("status", "PENDING")
            gate_id = (gate.get("data", {}) or {}).get("gate_id", "")
            if status == "PENDING" and block and gate_id:
                gate = client.gate_wait(gate_id, timeout_s=timeout_s, span_id=span_id, run_id=run_id)
                status = (gate.get("data", {}) or {}).get("status", "PENDING")
            if status not in {"APPROVED", "OVERRIDDEN"}:
                if on_deny == "raise":
                    raise RuntimeError(f"Kybernis gate denied: {status}")
                return {"kyb_blocked": True, "status": status}

            client.emit_event(run_id, "tool.call", {"tool": tool_name, "action": action})
            try:
                result = fn(*args, **kwargs)
                client.emit_event(run_id, "tool.result", {"tool": tool_name, "ok": True})
                if tool_id:
                    client.tool_report(tool_id, status="UP", circuit_state="CLOSED", failure_inc=0)
                if isinstance(result, dict) and "cost_micros" in result:
                    client.record_cost(run_id, int(result.get("cost_micros", 0)))
                return result
            except Exception as exc:
                client.emit_event(run_id, "tool.error", {"tool": tool_name, "error": str(exc)})
                if tool_id:
                    client.tool_report(tool_id, status="DOWN", circuit_state="OPEN", failure_inc=1)
                raise

        return wrapper

    return decorator
