from __future__ import annotations

"""
BloomOS client for talking to the Helix64 + BloomOS AI OS Kernel.

Drop this file somewhere on your PYTHONPATH, e.g. alongside your other
helix_cdc Python modules.

Requires:
    pip install requests
"""

from dataclasses import dataclass, field
from typing import Any, Dict, List, Optional, Union, TYPE_CHECKING
import json
import time

import requests

if TYPE_CHECKING:
    from helix_cdc.bloomos.sym import SymCapsule


@dataclass
class TaskRequest:
    task_type: str
    description: str
    anchors: Optional[List[str]] = None
    gate_policy: Optional[str] = None
    extra: Dict[str, Any] = field(default_factory=dict)

    def to_payload(self) -> Dict[str, Any]:
        payload: Dict[str, Any] = {
            "task_type": self.task_type,
            "description": self.description,
        }
        if self.anchors is not None:
            payload["anchors"] = self.anchors
        if self.gate_policy is not None:
            payload["gate_policy"] = self.gate_policy
        # Merge any extra fields (e.g., model overrides, max_tokens, etc.)
        payload.update(self.extra)
        return payload


@dataclass
class TaskReceipt:
    receipt_hash: str
    experiment_ref: Optional[str] = None


@dataclass
class TaskResult:
    task_id: str
    status: str
    output: Any
    gate_used: Optional[str]
    anchors_used: List[str]
    logits_mean: Optional[float]
    logits_std: Optional[float]
    receipt: Optional[TaskReceipt]
    raw: Dict[str, Any]

    @classmethod
    def from_response(cls, data: Dict[str, Any]) -> "TaskResult":
        # Be defensive: BloomOS may evolve its schema
        receipt_data = data.get("receipt")
        receipt: Optional[TaskReceipt] = None
        if isinstance(receipt_data, dict):
            receipt = TaskReceipt(
                receipt_hash=receipt_data.get("receipt_hash", ""),
                experiment_ref=receipt_data.get("experiment_ref"),
            )

        return cls(
            task_id=str(data.get("task_id", "")),
            status=str(data.get("status", "")),
            output=data.get("output"),
            gate_used=data.get("gate_used"),
            anchors_used=list(data.get("anchors_used") or []),
            logits_mean=data.get("logits_mean"),
            logits_std=data.get("logits_std"),
            receipt=receipt,
            raw=data,
        )


class BloomOSClient:
    """
    Thin HTTP client for the BloomOS FastAPI runtime.

    Example:
        client = BloomOSClient()
        result = client.run_math("Integrate sin(x) from 0 to pi")
        print(result.output)
    """

    def __init__(
        self,
        base_url: str = "http://localhost:7070",
        timeout: float = 60.0,
        log_path: Optional[str] = None,
    ) -> None:
        # No trailing slash for base_url
        self.base_url = base_url.rstrip("/")
        self.timeout = timeout
        self.log_path = log_path

    # ------------------------------
    # Low-level HTTP helpers
    # ------------------------------

    def _post_json(self, path: str, payload: Dict[str, Any]) -> Dict[str, Any]:
        url = f"{self.base_url}{path}"
        resp = requests.post(
            url,
            json=payload,
            headers={"Content-Type": "application/json"},
            timeout=self.timeout,
        )
        resp.raise_for_status()
        try:
            return resp.json()
        except json.JSONDecodeError as e:
            raise RuntimeError(f"Non-JSON response from BloomOS at {url}: {resp.text}") from e

    # ------------------------------
    # Public API
    # ------------------------------

    def run_task(
        self,
        task_type: str,
        description: str,
        anchors: Optional[List[str]] = None,
        gate_policy: Optional[str] = None,
        extra: Optional[Dict[str, Any]] = None,
    ) -> TaskResult:
        """
        Generic task runner. All other helpers (run_math, run_code, etc.)
        delegate here.
        """
        request = TaskRequest(
            task_type=task_type,
            description=description,
            anchors=anchors,
            gate_policy=gate_policy,
            extra=extra or {},
        )
        payload = request.to_payload()
        data = self._post_json("/task", payload)
        result = TaskResult.from_response(data)
        self._maybe_log(request, result)
        return result

    # Convenience wrappers

    def run_math(
        self,
        description: str,
        gate_policy: str = "math_bias",
        anchors: Optional[List[str]] = None,
        extra: Optional[Dict[str, Any]] = None,
    ) -> TaskResult:
        return self.run_task(
            task_type="math",
            description=description,
            anchors=anchors or ["math"],
            gate_policy=gate_policy,
            extra=extra,
        )

    def run_code(
        self,
        description: str,
        gate_policy: str = "code_bias",
        anchors: Optional[List[str]] = None,
        extra: Optional[Dict[str, Any]] = None,
    ) -> TaskResult:
        return self.run_task(
            task_type="code",
            description=description,
            anchors=anchors or ["code"],
            gate_policy=gate_policy,
            extra=extra,
        )

    def run_json(
        self,
        description: str,
        gate_policy: str = "json_bias",
        anchors: Optional[List[str]] = None,
        extra: Optional[Dict[str, Any]] = None,
    ) -> TaskResult:
        return self.run_task(
            task_type="json",
            description=description,
            anchors=anchors or ["json"],
            gate_policy=gate_policy,
            extra=extra,
        )

    def run_reasoning(
        self,
        description: str,
        gate_policy: str = "balanced",
        anchors: Optional[List[str]] = None,
        extra: Optional[Dict[str, Any]] = None,
    ) -> TaskResult:
        """
        Example of a "reason" style task type if you choose to add it on
        the BloomOS side.
        """
        return self.run_task(
            task_type="reason",
            description=description,
            anchors=anchors or ["math", "code", "json"],
            gate_policy=gate_policy,
            extra=extra,
        )

    # ------------------------------
    # Logging
    # ------------------------------

    def _maybe_log(self, request: TaskRequest, result: TaskResult) -> None:
        if not self.log_path:
            return
        entry = {
            "timestamp": time.time(),
            "timestamp_iso": time.strftime(
                "%Y-%m-%dT%H:%M:%S", time.localtime()
            ),
            "request": request.to_payload(),
            "result": {
                "task_id": result.task_id,
                "status": result.status,
                "gate_used": result.gate_used,
                "anchors_used": result.anchors_used,
                "logits_mean": result.logits_mean,
                "logits_std": result.logits_std,
                "receipt": {
                    "receipt_hash": result.receipt.receipt_hash
                    if result.receipt
                    else None,
                    "experiment_ref": result.receipt.experiment_ref
                    if result.receipt
                    else None,
                }
                if result.receipt
                else None,
            },
        }
        try:
            with open(self.log_path, "a", encoding="utf-8") as f:
                f.write(json.dumps(entry) + "\n")
        except OSError:
            # Non-fatal: fail silently, runtime shouldn't break
            pass

    # ------------------------------
    # SymCapsule Integration
    # ------------------------------

    def result_to_capsule(
        self,
        result: TaskResult,
        request: Optional[TaskRequest] = None,
    ) -> "SymCapsule":
        """
        Convert a TaskResult to a SymCapsule.

        Args:
            result: TaskResult from BloomOS
            request: Original TaskRequest (for context)

        Returns:
            SymCapsule wrapping the result
        """
        from helix_cdc.bloomos.sym import lower_task_to_sym

        task_type = request.task_type if request else "unknown"
        description = request.description if request else ""
        program = result.raw.get("helix_program", "ABCDEFGH")

        return lower_task_to_sym(
            task_type=task_type,
            description=description,
            program=program,
            result=result.raw,
        )

    def run_task_capsule(
        self,
        task_type: str,
        description: str,
        anchors: Optional[List[str]] = None,
        gate_policy: Optional[str] = None,
        extra: Optional[Dict[str, Any]] = None,
    ) -> "SymCapsule":
        """
        Run a task and return result as SymCapsule.

        Same as run_task() but returns a SymCapsule instead of TaskResult.
        """
        request = TaskRequest(
            task_type=task_type,
            description=description,
            anchors=anchors,
            gate_policy=gate_policy,
            extra=extra or {},
        )
        payload = request.to_payload()
        data = self._post_json("/task", payload)
        result = TaskResult.from_response(data)
        self._maybe_log(request, result)
        return self.result_to_capsule(result, request)
