"""
LLM orchestration service: builds prompts, calls providers, streams SSE events.

Workflow (sequential):
  1. [Optional] Edison Scientific literature search → grounded citations
  2. OpenAI / Ollama structuring call → JSON effects list
"""

import asyncio
import hashlib
import json
import re
from itertools import combinations
from typing import AsyncGenerator, TYPE_CHECKING

if TYPE_CHECKING:
    from api.models.requests import SuggestEffectsRequest

# ---------------------------------------------------------------------------
# Fixed system instructions (not user-editable)
# ---------------------------------------------------------------------------
_SYSTEM_INSTRUCTIONS = """\
You are a scientific advisor for experimental design in chemistry and materials science.
Your task is to recommend which model terms (main effects, interactions, quadratic effects) \
are most scientifically justified for an optimal experimental design, based on domain \
knowledge and published literature.

IMPORTANT — effects selection: You must ONLY select effects from the provided \
"Available effects" list. Do not suggest effects that are not in the list. Do not invent \
or modify variable names. Select the minimum necessary set for a well-powered design — \
do not include effects you cannot justify scientifically.

CRITICAL — sources: Only cite references that appear verbatim in the "Literature context" \
section of the user message. If no "Literature context" section is present, you MUST return \
an empty sources array ([]). Never generate, guess, or invent citations. Fabricated \
references are significantly more harmful than an empty list.\
"""

# ---------------------------------------------------------------------------
# Input sanitization
# ---------------------------------------------------------------------------

_CONTROL_CHAR_RE = re.compile(r"[\x00-\x08\x0b\x0c\x0e-\x1f\x7f]")
_MAX_CONTEXT_LENGTH = 2000
_MAX_VARIABLE_NAME_LENGTH = 100


def _sanitize_text(text: str, max_length: int) -> str:
    """Strip control characters and truncate to max_length."""
    cleaned = _CONTROL_CHAR_RE.sub("", text)
    if len(cleaned) > max_length:
        cleaned = cleaned[:max_length] + "…"
    return cleaned


def _sanitize_variable_name(name: str) -> str:
    """Sanitize a variable name for safe prompt interpolation."""
    return _sanitize_text(name, _MAX_VARIABLE_NAME_LENGTH)


# ---------------------------------------------------------------------------
# JSON schema for structured output
# OpenAI strict mode requires: additionalProperties=false on all objects,
# all properties in required.  We use arrays-of-objects for reasoning/confidence
# instead of free-form dicts to satisfy that constraint.
# ---------------------------------------------------------------------------
EFFECTS_SCHEMA: dict = {
    "type": "object",
    "properties": {
        "effects": {
            "type": "array",
            "description": "Selected effect strings taken verbatim from the available effects list.",
            "items": {"type": "string"},
        },
        "reasoning": {
            "type": "array",
            "description": "Scientific reasoning for each selected effect.",
            "items": {
                "type": "object",
                "properties": {
                    "effect": {"type": "string"},
                    "reason": {"type": "string"},
                },
                "required": ["effect", "reason"],
                "additionalProperties": False,
            },
        },
        "confidence": {
            "type": "array",
            "description": "Confidence level for each selected effect.",
            "items": {
                "type": "object",
                "properties": {
                    "effect": {"type": "string"},
                    "level": {"type": "string", "enum": ["high", "medium", "low"]},
                },
                "required": ["effect", "level"],
                "additionalProperties": False,
            },
        },
        "sources": {
            "type": "array",
            "description": (
                "Literature sources supporting these recommendations. "
                "ONLY include citations taken verbatim from the 'Literature context' "
                "section provided in the user message. "
                "If no Literature context section was provided, return an EMPTY array. "
                "Do NOT generate, guess, or hallucinate any citations."
            ),
            "items": {"type": "string"},
        },
        "disclaimer": {
            "type": "string",
            "description": (
                "Standard disclaimer about the AI-generated nature of these recommendations."
            ),
        },
    },
    "required": ["effects", "reasoning", "confidence", "sources", "disclaimer"],
    "additionalProperties": False,
}

_DEFAULT_DISCLAIMER = (
    "These recommendations are generated by an AI model and should be reviewed by "
    "a domain expert. Sources should be independently verified before citation."
)

# ---------------------------------------------------------------------------
# Prompt assembly helpers
# ---------------------------------------------------------------------------

def _format_variable(v: dict) -> str:
    vtype = v.get("type", "real")
    name = _sanitize_variable_name(v["name"])
    if vtype == "categorical":
        cats = v.get("categories", [])
        return f"- {name} (categorical): {cats}"
    lo = v.get("lower_bound", "?")
    hi = v.get("upper_bound", "?")
    unit = v.get("unit") or ""
    unit_str = f" {unit}" if unit else ""
    return f"- {name} ({vtype}): {lo}–{hi}{unit_str}"


def _build_effects_lists(
    variables: list[dict],
) -> tuple[list[str], list[str], list[str]]:
    cont_types = {"real", "integer", "discrete", None}
    all_names = [_sanitize_variable_name(v["name"]) for v in variables]
    cont_names = [_sanitize_variable_name(v["name"]) for v in variables if v.get("type") in cont_types]
    interactions = [f"{a}*{b}" for a, b in combinations(all_names, 2)]
    quadratics = [f"{n}**2" for n in cont_names]
    return all_names, interactions, quadratics


def build_user_prompt(
    variables: list[dict],
    system_context: str,
    literature_context: str | None = None,
) -> str:
    main, interactions, quadratics = _build_effects_lists(variables)
    var_lines = "\n".join(_format_variable(v) for v in variables)
    safe_context = _sanitize_text(system_context, _MAX_CONTEXT_LENGTH)

    lit_section = ""
    if literature_context:
        lit_section = (
            f"\nLiterature context (retrieved from scientific database — "
            f"citations included):\n{literature_context}\n"
        )

    interactions_str = ", ".join(interactions) if interactions else "(none — fewer than 2 variables)"
    quadratics_str = ", ".join(quadratics) if quadratics else "(none — no continuous variables)"

    return (
        f"System description: {safe_context}"
        f"{lit_section}\n"
        f"Variable space:\n{var_lines}\n\n"
        f"Available effects to select from:\n"
        f"  Main effects:    {', '.join(main)}\n"
        f"  Interactions:    {interactions_str}\n"
        f"  Quadratic:       {quadratics_str}\n\n"
        f"Recommend which effects are most likely to be scientifically significant "
        f"for this system{' based on the literature context above' if literature_context else ''}."
    )


def build_edison_query(variables: list[dict], system_context: str) -> str:
    var_desc = "; ".join(
        f"{_sanitize_variable_name(v['name'])} ({v.get('type', 'real')})" for v in variables
    )
    safe_context = _sanitize_text(system_context, _MAX_CONTEXT_LENGTH)
    return (
        f"What main effects, interactions, and quadratic/curvature terms are "
        f"scientifically significant for the following experimental system?\n\n"
        f"System: {safe_context}\n\n"
        f"Variables: {var_desc}\n\n"
        f"Specifically: (1) Which variables most strongly affect the target response? "
        f"(2) Which pairwise interactions are likely important? "
        f"(3) Which variables likely show non-linear (quadratic/curvature) behavior? "
        f"Please cite specific literature sources where possible."
    )


# ---------------------------------------------------------------------------
# Provider factory
# ---------------------------------------------------------------------------

def _get_structuring_provider(config: "LLMProviderConfig"):  # noqa: F821
    from api.services.llm_config import resolve_api_key

    provider = config.provider
    if provider == "openai":
        from api.services.providers.openai_provider import OpenAIProvider
        key = resolve_api_key("openai", config.api_key)
        return OpenAIProvider(api_key=key, model=config.model)
    if provider == "ollama":
        from api.services.providers.ollama_provider import OllamaProvider
        base_url = (config.base_url or "http://localhost:11434").rstrip("/")
        # The OpenAI SDK appends /chat/completions to base_url, so it must
        # include Ollama's /v1 prefix (e.g. http://localhost:11434/v1).
        if not base_url.endswith("/v1"):
            base_url = base_url + "/v1"
        return OllamaProvider(model=config.model, base_url=base_url)
    raise ValueError(f"Unknown structuring provider: {provider!r}")


# ---------------------------------------------------------------------------
# SSE helper
# ---------------------------------------------------------------------------

def _sse(data: dict) -> str:
    return f"data: {json.dumps(data)}\n\n"


# ---------------------------------------------------------------------------
# Main streaming orchestrator
# ---------------------------------------------------------------------------

async def suggest_effects_stream(
    variables: list[dict],
    request: "SuggestEffectsRequest",
    edison_cache: dict | None = None,
) -> AsyncGenerator[str, None]:
    """
    Yields Server-Sent Event strings.

    Event shapes:
      {"status": "searching_literature",          "message": "...", "trajectory_url": "..."}
      {"status": "searching_literature_progress", "message": "...", "trajectory_url": "..."}
      {"status": "literature_complete",  "message": "...", "trajectory_url": "...", "cached": bool}
      {"status": "literature_warning",   "message": "...", "trajectory_url": "...", "cached": bool}
      {"status": "literature_error",     "message": "..."}
      {"status": "structuring",          "message": "..."}
      {"status": "complete",             "result": {...}}
      {"status": "error",                "message": "..."}

    edison_cache: per-session mutable dict for two-level caching.
      Pending entry:  {"status": "running",   "task_id": str, "trajectory_url": str}
      Complete entry: {"status": "complete",  "task_id": str, "trajectory_url": str,
                       "answer": str, "formatted_answer": str,
                       "has_successful_answer": bool}
    """
    literature_context: str | None = None

    # ------------------------------------------------------------------
    # Stage 1 (optional): Edison Scientific literature search
    # ------------------------------------------------------------------
    if request.edison_config:
        try:
            from api.services.providers.edison_provider import EdisonLiteratureProvider
            from api.services.llm_config import resolve_api_key
            edison_key = resolve_api_key("edison", request.edison_config.api_key)
            edison = EdisonLiteratureProvider(
                api_key=edison_key,
                job_type=request.edison_config.job_type,
            )
            query = build_edison_query(variables, request.system_context)

            # Stable 16-hex cache key: SHA-256(job_type:query)
            cache_key = hashlib.sha256(
                f"{request.edison_config.job_type}:{query}".encode()
            ).hexdigest()[:16]
            force_refresh = request.edison_config.force_refresh

            # ── Path A: completed cache hit ────────────────────────────────
            cached = None if force_refresh else (edison_cache or {}).get(cache_key)
            if cached and cached.get("status") == "complete":
                traj = cached.get("trajectory_url")
                if cached.get("has_successful_answer"):
                    literature_context = cached["formatted_answer"]
                    yield _sse({
                        "status": "literature_complete",
                        "message": "Using cached Edison results — no new search required.",
                        "trajectory_url": traj,
                        "cached": True,
                    })
                else:
                    yield _sse({
                        "status": "literature_warning",
                        "message": (
                            "Cached Edison result: no confident answer found. "
                            "Proceeding without literature context."
                        ),
                        "trajectory_url": traj,
                        "cached": True,
                    })

            else:
                # ── Path B: submit or resume a task ───────────────────────
                task_id: str | None = None
                trajectory_url: str | None = None
                task_result = None

                # Try to resume a running task from cache (avoids re-billing Edison)
                if not force_refresh and cached and cached.get("status") == "running":
                    try:
                        check = await edison.poll_task(cached["task_id"])
                        # Task is still alive on the platform
                        task_id = cached["task_id"]
                        trajectory_url = edison.get_trajectory_url(task_id)
                        if getattr(check, "status", None) == "success":
                            task_result = check  # Already done while we were offline
                        yield _sse({
                            "status": "searching_literature",
                            "message": "Resuming Edison task — watch progress at the link below",
                            "trajectory_url": trajectory_url,
                        })
                    except Exception:
                        # Task expired or not found — fall through to new submission
                        if edison_cache is not None:
                            edison_cache.pop(cache_key, None)

                # Submit a new task if we couldn't resume
                if task_id is None:
                    task_id = await edison.start_search(query)
                    trajectory_url = edison.get_trajectory_url(task_id)
                    if edison_cache is not None:
                        edison_cache[cache_key] = {
                            "status": "running",
                            "task_id": task_id,
                            "trajectory_url": trajectory_url,
                        }
                    yield _sse({
                        "status": "searching_literature",
                        "message": "Edison task submitted — watch progress at the link below",
                        "trajectory_url": trajectory_url,
                    })

                # Poll every 30 s until done (skip if already resolved during resume)
                if task_result is None:
                    timeout_secs = request.edison_config.timeout_secs or 1200
                    loop_start = asyncio.get_event_loop().time()

                    while True:
                        await asyncio.sleep(30)
                        elapsed = int(asyncio.get_event_loop().time() - loop_start)
                        if elapsed >= timeout_secs:
                            raise asyncio.TimeoutError(
                                f"Edison search timed out after "
                                f"{elapsed // 60}m {elapsed % 60}s"
                            )
                        task_result = await edison.poll_task(task_id)
                        if getattr(task_result, "status", None) == "success":
                            break
                        yield _sse({
                            "status": "searching_literature_progress",
                            "message": (
                                f"Still searching literature… "
                                f"({elapsed // 60}m {elapsed % 60}s elapsed)"
                            ),
                            "trajectory_url": trajectory_url,
                        })

                # Extract results and store completed entry in cache
                answer = getattr(task_result, "answer", "")
                formatted_answer = getattr(task_result, "formatted_answer", "")
                has_answer = getattr(task_result, "has_successful_answer", False)
                if edison_cache is not None:
                    edison_cache[cache_key] = {
                        "status": "complete",
                        "task_id": task_id,
                        "trajectory_url": trajectory_url,
                        "answer": answer,
                        "formatted_answer": formatted_answer,
                        "has_successful_answer": has_answer,
                    }

                if has_answer and formatted_answer:
                    literature_context = formatted_answer
                    yield _sse({
                        "status": "literature_complete",
                        "message": "Literature search complete — passing context to structuring model.",
                        "trajectory_url": trajectory_url,
                        "cached": False,
                    })
                else:
                    yield _sse({
                        "status": "literature_warning",
                        "message": (
                            "Edison did not find a confident answer. "
                            "Proceeding without literature context."
                        ),
                        "trajectory_url": trajectory_url,
                        "cached": False,
                    })

        except Exception as exc:
            yield _sse({
                "status": "literature_error",
                "message": f"Literature search failed ({exc}). Proceeding without.",
            })

    # ------------------------------------------------------------------
    # Stage 2: Structuring model (OpenAI / Ollama)
    # ------------------------------------------------------------------
    provider_label = f"{request.structuring_provider.provider}/{request.structuring_provider.model}"
    yield _sse({
        "status": "structuring",
        "message": f"Extracting structured effects with {provider_label}…",
    })

    try:
        provider = _get_structuring_provider(request.structuring_provider)
        user_prompt = build_user_prompt(variables, request.system_context, literature_context)
        result = await provider.suggest_effects(_SYSTEM_INSTRUCTIONS, user_prompt, EFFECTS_SCHEMA)

        # Ensure required keys are present with safe defaults
        result.setdefault("effects", [])
        result.setdefault("reasoning", [])
        result.setdefault("confidence", [])
        result.setdefault("sources", [])
        result.setdefault("disclaimer", _DEFAULT_DISCLAIMER)

        # Attach raw literature context so the UI can display it
        result["literature_context"] = literature_context

        yield _sse({"status": "complete", "result": result})

    except Exception as exc:
        yield _sse({"status": "error", "message": str(exc)})
