import base64
import json
import mimetypes
import os
from pathlib import Path
from typing import Any, Dict, List, Optional

import requests

_INTERNAL_API_URL = "https://toxo-api-cwsddklqcq-uc.a.run.app"


class ToxoCloudClient:
    """Lightweight client for TOXO Cloud endpoints."""

    def __init__(self, api_key: Optional[str] = None, timeout: int = 120):
        # Service URL is intentionally fixed and internal-only.
        self.api_url = _INTERNAL_API_URL.rstrip("/")
        self.api_key = api_key or os.getenv("GEMINI_API_KEY") or os.getenv("GOOGLE_API_KEY")
        self.timeout = timeout

    def _layer_base64(self, layer: str | Path) -> str:
        p = Path(layer)
        if not p.exists():
            raise FileNotFoundError(f"Layer file not found: {p}")
        return base64.b64encode(p.read_bytes()).decode("ascii")

    def _doc_payload(self, doc_path: str | Path) -> Dict[str, Any]:
        p = Path(doc_path)
        if not p.exists():
            raise FileNotFoundError(f"Document not found: {p}")
        mt, _ = mimetypes.guess_type(str(p))
        return {
            "path": str(p),
            "data_base64": base64.b64encode(p.read_bytes()).decode("ascii"),
            "mime_type": mt,
        }

    def _post(self, path: str, payload: Dict[str, Any], timeout: Optional[int] = None) -> Dict[str, Any]:
        r = requests.post(
            f"{self.api_url}{path}",
            json=payload,
            timeout=timeout or self.timeout,
        )
        try:
            r.raise_for_status()
        except requests.HTTPError as exc:
            detail = ""
            try:
                body = r.json()
                if isinstance(body, dict) and body.get("detail"):
                    detail = str(body["detail"])
            except Exception:
                detail = r.text.strip()[:400]
            if detail:
                raise requests.HTTPError(
                    f"{exc} | detail: {detail}",
                    response=r,
                    request=r.request,
                ) from exc
            raise
        return r.json()

    def root(self) -> Dict[str, Any]:
        r = requests.get(f"{self.api_url}/", timeout=20)
        r.raise_for_status()
        return r.json()

    def health(self) -> Dict[str, Any]:
        r = requests.get(f"{self.api_url}/health", timeout=20)
        r.raise_for_status()
        return r.json()

    # /v1/query
    def query(
        self,
        layer: str | Path,
        question: str,
        *,
        api_key: Optional[str] = None,
        model: str = "gemini-2.5-flash-lite",
        provider: str = "gemini",
        context: Optional[Dict[str, Any]] = None,
        response_depth: str = "balanced",
    ) -> str:
        data = self._post(
            "/v1/query",
            {
                "layer_base64": self._layer_base64(layer),
                "question": question,
                "api_key": api_key or self.api_key,
                "model": model,
                "provider": provider,
                "context": context,
                "response_depth": response_depth,
            },
        )
        return str(data.get("response", ""))

    # /v1/query_multimodal
    def query_multimodal(
        self,
        layer: str | Path,
        question: str,
        *,
        image_path: Optional[str | Path] = None,
        document_path: Optional[str | Path] = None,
        image_base64: Optional[str] = None,
        document_base64: Optional[str] = None,
        mime_type: Optional[str] = None,
        api_key: Optional[str] = None,
        model: str = "gemini-2.5-flash-lite",
        provider: str = "gemini",
        context: Optional[Dict[str, Any]] = None,
        response_depth: str = "balanced",
    ) -> str:
        sources = [image_path, document_path, image_base64, document_base64]
        if sum(1 for x in sources if x) != 1:
            raise ValueError("Provide exactly one of image_path, document_path, image_base64, document_base64")

        img_b64 = image_base64
        doc_b64 = document_base64
        effective_mime_type = mime_type
        if image_path:
            image_p = Path(image_path)
            img_b64 = base64.b64encode(image_p.read_bytes()).decode("ascii")
            if not effective_mime_type:
                guessed, _ = mimetypes.guess_type(str(image_p))
                effective_mime_type = guessed
            if img_b64 and not str(img_b64).startswith("data:"):
                mt = effective_mime_type or "image/png"
                img_b64 = f"data:{mt};base64,{img_b64}"
        if document_path:
            document_p = Path(document_path)
            doc_b64 = base64.b64encode(document_p.read_bytes()).decode("ascii")
            if not effective_mime_type:
                guessed, _ = mimetypes.guess_type(str(document_p))
                effective_mime_type = guessed or "application/octet-stream"
            if doc_b64 and not str(doc_b64).startswith("data:"):
                mt = effective_mime_type or "application/pdf"
                doc_b64 = f"data:{mt};base64,{doc_b64}"
        if document_base64 and not str(document_base64).startswith("data:"):
            mt = effective_mime_type or "application/pdf"
            doc_b64 = f"data:{mt};base64,{document_base64}"
        if image_base64 and not str(image_base64).startswith("data:"):
            mt = effective_mime_type or "image/png"
            img_b64 = f"data:{mt};base64,{image_base64}"

        data = self._post(
            "/v1/query_multimodal",
            {
                "layer_base64": self._layer_base64(layer),
                "question": question,
                "image_base64": img_b64,
                "document_base64": doc_b64,
                "mime_type": effective_mime_type,
                "api_key": api_key or self.api_key,
                "model": model,
                "provider": provider,
                "context": context,
                "response_depth": response_depth,
            },
            timeout=max(self.timeout, 180),
        )
        return str(data.get("response", ""))

    # /v1/feedback
    def feedback(
        self,
        layer: str | Path,
        question: str,
        response: str,
        *,
        rating: Optional[float] = None,
        suggestions: Optional[List[str]] = None,
    ) -> Dict[str, Any]:
        return self._post(
            "/v1/feedback",
            {
                "layer_base64": self._layer_base64(layer),
                "question": question,
                "response": response,
                "rating": rating,
                "suggestions": suggestions,
            },
            timeout=40,
        )

    # /v1/train_from_data
    def train_from_data(
        self,
        *,
        description: str,
        out_path: str | Path,
        examples: Optional[List[Dict[str, str]]] = None,
        contexts: Optional[List[str]] = None,
        epochs: int = 10,
        llm_provider: str = "gemini",
        llm_model: Optional[str] = None,
        domain: Optional[str] = None,
        api_key: Optional[str] = None,
        mode: Optional[str] = None,
        num_examples: int = 5,
    ) -> Dict[str, Any]:
        payload: Dict[str, Any] = {
            "description": description,
            "epochs": epochs,
            "llm_provider": llm_provider,
            "llm_model": llm_model,
            "api_key": api_key or self.api_key,
            "domain": domain,
            "examples": [{"input": e["input"], "output": e["output"]} for e in (examples or [])],
            "contexts": [{"text": c} for c in (contexts or [])],
            "mode": mode,
            "num_examples": num_examples,
        }
        data = self._post("/v1/train_from_data", payload, timeout=max(self.timeout, 600))
        layer_b64 = data.get("layer_base64")
        if not isinstance(layer_b64, str) or not layer_b64:
            raise RuntimeError("Missing layer_base64 in train_from_data response")
        Path(out_path).parent.mkdir(parents=True, exist_ok=True)
        Path(out_path).write_bytes(base64.b64decode(layer_b64))
        return data

    # /v1/train/extract_contexts
    def train_extract_contexts(
        self,
        *,
        docs: List[str | Path],
        api_key: Optional[str] = None,
        index_id: str = "train_extract_contexts",
        domain: str = "general",
    ) -> Dict[str, Any]:
        payload = {
            "index_id": index_id,
            "domain": domain,
            "api_key": api_key or self.api_key,
            "docs": [self._doc_payload(d) for d in docs],
        }
        return self._post("/v1/train/extract_contexts", payload, timeout=max(self.timeout, 600))

    # /v1/rag/index
    def rag_index(
        self,
        *,
        index_id: str,
        docs: List[str | Path],
        domain: str = "general",
        layer: Optional[str | Path] = None,
        api_key: Optional[str] = None,
    ) -> Dict[str, Any]:
        payload: Dict[str, Any] = {
            "index_id": index_id,
            "domain": domain,
            "layer_base64": self._layer_base64(layer) if layer else None,
            "api_key": api_key or self.api_key,
            "docs": [self._doc_payload(d) for d in docs],
        }
        return self._post("/v1/rag/index", payload, timeout=max(self.timeout, 600))

    # /v1/rag/query
    def rag_query(
        self,
        *,
        layer: str | Path,
        index_id: str,
        question: str,
        api_key: Optional[str] = None,
        llm_provider: str = "gemini",
        llm_model: Optional[str] = None,
        top_k: int = 5,
        response_depth: str = "balanced",
    ) -> Dict[str, Any]:
        payload = {
            "layer_base64": self._layer_base64(layer),
            "index_id": index_id,
            "question": question,
            "api_key": api_key or self.api_key,
            "llm_provider": llm_provider,
            "llm_model": llm_model,
            "top_k": int(top_k),
            "response_depth": response_depth,
        }
        return self._post("/v1/rag/query", payload, timeout=max(self.timeout, 300))

    # /v1/agent/run
    def agent_run(
        self,
        *,
        config: Dict[str, Any] | str | Path,
        layer_map: Dict[str, str | Path],
        api_key: Optional[str] = None,
        llm_provider: str = "gemini",
        llm_model: Optional[str] = None,
    ) -> Dict[str, Any]:
        cfg: Dict[str, Any]
        if isinstance(config, dict):
            cfg = config
        else:
            p = Path(config)
            text = p.read_text(encoding="utf-8")
            if p.suffix.lower() in (".yml", ".yaml"):
                try:
                    import yaml
                except Exception as exc:
                    raise RuntimeError("pyyaml is required to read YAML agent config") from exc
                cfg = yaml.safe_load(text) or {}
            else:
                cfg = json.loads(text)

        layers = [
            {
                "agent_id": aid,
                "layer_path": str(path),
                "layer_base64": self._layer_base64(path),
            }
            for aid, path in layer_map.items()
        ]
        payload = {
            "config": cfg,
            "layers": layers,
            "api_key": api_key or self.api_key,
            "llm_provider": llm_provider,
            "llm_model": llm_model,
        }
        return self._post("/v1/agent/run", payload, timeout=max(self.timeout, 1800))
