# retoor <retoor@molodetz.nl>

from __future__ import annotations

import json
import logging
import re
from typing import Callable

from devplacepy import stealth
from devplacepy.config import INTERNAL_GATEWAY_URL, INTERNAL_MODEL

from .phases import PHASE_PLANNING, PHASE_SEARCHING

logger = logging.getLogger(__name__)

ENHANCE_TIMEOUT_SECONDS = 90.0
ENHANCE_MAX_TOKENS = 400
MAX_SUBQUERIES = 6
MAX_FOLLOWUP_QUERIES = 3
MAX_COVERED_TITLES = 40

PLANNER_PROMPT = (
    "You plan web research. Given a research question, produce a JSON object with one "
    "key 'queries': an array of 3 to 6 concise, diverse web search queries that together "
    "cover the question from multiple angles. Return ONLY the JSON object, no prose."
)
FOLLOWUP_PROMPT = (
    "You plan additional web research to fill coverage gaps. Given a research QUESTION "
    "and the TITLES of sources already gathered, decide whether more searching would "
    "surface information the current sources do not yet cover. Produce a JSON object "
    "with one key 'queries': an array of 0 to 3 concise web search queries targeting "
    "what is missing. Return an empty array when the sources already cover the question "
    "well. Return ONLY the JSON object, no prose."
)


def _fallback(query: str) -> list[str]:
    base = query.strip()
    variants = [
        base,
        f"{base} overview",
        f"{base} latest",
        f"{base} explained",
    ]
    seen: list[str] = []
    for variant in variants:
        if variant and variant not in seen:
            seen.append(variant)
    return seen[:MAX_SUBQUERIES]


def _parse(text: str) -> list[str]:
    match = re.search(r"\{.*\}", text, re.DOTALL)
    if not match:
        return []
    try:
        payload = json.loads(match.group())
    except (ValueError, TypeError):
        return []
    queries = payload.get("queries")
    if not isinstance(queries, list):
        return []
    cleaned = [str(item).strip() for item in queries if str(item).strip()]
    return cleaned[:MAX_SUBQUERIES]


def _noop(frame: dict) -> None:
    return None


async def plan_queries(
    query: str,
    api_key: str,
    emit: Callable[[dict], None] = _noop,
    model: str = INTERNAL_MODEL,
) -> list[str]:
    emit(
        {
            "type": "substep",
            "phase": PHASE_PLANNING,
            "message": "Drafting diverse search angles",
        }
    )
    payload = {
        "model": model or INTERNAL_MODEL,
        "messages": [
            {"role": "system", "content": PLANNER_PROMPT},
            {"role": "user", "content": query},
        ],
        "max_tokens": ENHANCE_MAX_TOKENS,
        "temperature": 0.3,
    }
    headers = {
        "Authorization": f"Bearer {api_key}",
        "Content-Type": "application/json",
        "X-App-Reference": "devplace-deepsearch-v-1-0-0",
    }
    try:
        async with stealth.stealth_async_client(timeout=ENHANCE_TIMEOUT_SECONDS) as client:
            response = await client.post(
                INTERNAL_GATEWAY_URL, json=payload, headers=headers
            )
        if response.status_code >= 400:
            raise RuntimeError(f"planner gateway returned {response.status_code}")
        data = response.json()
        content = (
            (data.get("choices") or [{}])[0].get("message", {}).get("content") or ""
        )
        parsed = _parse(content)
        if parsed:
            emit(
                {
                    "type": "substep",
                    "phase": PHASE_PLANNING,
                    "message": f"Planned {len(parsed)} search angles",
                }
            )
            return parsed
    except Exception as exc:
        logger.warning("deepsearch query planner failed, using fallback: %s", exc)
    fallback = _fallback(query)
    emit(
        {
            "type": "substep",
            "phase": PHASE_PLANNING,
            "message": f"Planner unavailable, using {len(fallback)} heuristic angles",
        }
    )
    return fallback


async def plan_followup_queries(
    query: str,
    covered_titles: list[str],
    api_key: str,
    emit: Callable[[dict], None] = _noop,
    model: str = INTERNAL_MODEL,
) -> list[str]:
    if not covered_titles:
        return []
    titles_block = "\n".join(f"- {title}" for title in covered_titles[:MAX_COVERED_TITLES])
    payload = {
        "model": model or INTERNAL_MODEL,
        "messages": [
            {"role": "system", "content": FOLLOWUP_PROMPT},
            {
                "role": "user",
                "content": f"QUESTION: {query}\n\nALREADY COVERED (source titles):\n{titles_block}",
            },
        ],
        "max_tokens": ENHANCE_MAX_TOKENS,
        "temperature": 0.3,
    }
    headers = {
        "Authorization": f"Bearer {api_key}",
        "Content-Type": "application/json",
        "X-App-Reference": "devplace-deepsearch-v-1-0-0",
    }
    try:
        async with stealth.stealth_async_client(timeout=ENHANCE_TIMEOUT_SECONDS) as client:
            response = await client.post(
                INTERNAL_GATEWAY_URL, json=payload, headers=headers
            )
        if response.status_code >= 400:
            raise RuntimeError(f"followup planner gateway returned {response.status_code}")
        data = response.json()
        content = (
            (data.get("choices") or [{}])[0].get("message", {}).get("content") or ""
        )
        parsed = _parse(content)[:MAX_FOLLOWUP_QUERIES]
        if parsed:
            emit(
                {
                    "type": "substep",
                    "phase": PHASE_SEARCHING,
                    "message": f"Planned {len(parsed)} follow-up queries to fill gaps",
                }
            )
        return parsed
    except Exception as exc:
        logger.info("deepsearch followup planner unavailable, skipping refinement: %s", exc)
        return []
