"""Client for the WaterScopeAI Hugging Face Space /agent_run endpoint."""

from __future__ import annotations

import concurrent.futures
import os
import re
from functools import lru_cache
from typing import Any

from gradio_client import Client

SPACE_ID = os.getenv("WATERSCOPE_SPACE_ID", "EnvironmentalAI/WaterScopeAI")
TIMEOUT_SEC = int(os.getenv("WATERSCOPE_TIMEOUT_SEC", "300"))

_SPACE_EXECUTOR = concurrent.futures.ThreadPoolExecutor(
    max_workers=2,
    thread_name_prefix="waterscope-space",
)

MODELS = ("base", "dpo", "da_dpo", "tulucore")
APPLICATIONS = (
    "general_qa",
    "document_qa",
    "paper_synthesis",
    "method_recommendation",
    "research_gaps",
    "claim_check",
)

_SYNTHESIS = re.compile(r"\b(summarize|summary|overview|synthesize|main finding)\b", re.I)
_METHODS = re.compile(r"\b(method|approach|model|algorithm|recommend)\b", re.I)
_GAPS = re.compile(r"\b(gap|limitation|future work|open question|unclear)\b", re.I)
_CLAIM = re.compile(r"\b(does the paper|is it true|claim|support|evidence that)\b", re.I)


def infer_application(question: str, *, explicit: str | None = None) -> str:
    if explicit and explicit != "auto":
        if explicit not in APPLICATIONS:
            raise ValueError(f"Unknown application {explicit!r}")
        return explicit
    text = question.lower()
    if _SYNTHESIS.search(text):
        return "paper_synthesis"
    if _GAPS.search(text):
        return "research_gaps"
    if _CLAIM.search(text):
        return "claim_check"
    if _METHODS.search(text):
        return "method_recommendation"
    return "document_qa"


@lru_cache(maxsize=1)
def _client() -> Client:
    token = os.getenv("HF_TOKEN") or os.getenv("HUGGINGFACE_HUB_TOKEN")
    kwargs: dict[str, Any] = {"httpx_kwargs": {"timeout": TIMEOUT_SEC}}
    if token:
        kwargs["hf_token"] = token
    return Client(SPACE_ID, **kwargs)


def run_agent(
    *,
    context: str,
    question: str,
    model: str | None = None,
    application: str = "document_qa",
    max_new_tokens: int | None = None,
) -> tuple[str, str]:
    """Call /agent_run. Returns (answer, diagnostics)."""
    selected_model = model or os.getenv("WATERSCOPE_MODEL", "dpo")
    if selected_model not in MODELS:
        raise ValueError(f"Unknown model {selected_model!r}. Choose from: {MODELS}")
    app_type = infer_application(question, explicit=application if application != "auto" else None)
    tokens = max_new_tokens or int(os.getenv("WATERSCOPE_MAX_TOKENS", "1024"))

    def _predict() -> tuple[str, str]:
        response = _client().predict(
            selected_model,
            app_type,
            context,
            question,
            tokens,
            api_name="/agent_run",
        )
        answer = response[0] if response else ""
        diagnostics = response[1] if len(response) > 1 else ""
        return answer, diagnostics

    # Run in a worker thread so Gradio asyncio does not clash with MCP's loop.
    answer, diagnostics = _SPACE_EXECUTOR.submit(_predict).result(timeout=TIMEOUT_SEC)
    return answer, diagnostics


def agent_answer(
    context: str,
    question: str,
    *,
    application: str = "auto",
    model: str | None = None,
    max_new_tokens: int | None = None,
) -> tuple[str, str, str]:
    """Returns (answer, application_used, diagnostics)."""
    app_type = infer_application(question, explicit=application if application != "auto" else None)
    answer, diagnostics = run_agent(
        context=context,
        question=question,
        model=model,
        application=app_type,
        max_new_tokens=max_new_tokens,
    )
    return answer, app_type, diagnostics