import logging

import httpx

from axon.config import settings
from axon.models import SearchResult

logger = logging.getLogger(__name__)

_SYSTEM_PROMPT = (
    "You are a precise research assistant. Answer the user's query based only on the "
    "provided search results. Cite sources using [1], [2] notation. Be concise and factual."
)


def _headers() -> dict:
    h = {"Content-Type": "application/json"}
    if settings.litellm_api_key:
        h["Authorization"] = f"Bearer {settings.litellm_api_key}"
    return h


def _context_block(results: list[SearchResult], limit: int = 20) -> str:
    parts = []
    for i, r in enumerate(results[:limit], 1):
        parts.append(f"[{i}] {r.title}\nURL: {r.url}\n{r.content}")
    return "\n\n".join(parts)


async def _chat(messages: list[dict], max_tokens: int = 1024, timeout: float = 30.0) -> str:
    payload = {
        "model": settings.synthesis_model,
        "messages": messages,
        "max_tokens": max_tokens,
    }
    async with httpx.AsyncClient() as client:
        response = await client.post(
            f"{settings.litellm_base_url}/chat/completions",
            json=payload,
            headers=_headers(),
            timeout=timeout,
        )
        response.raise_for_status()
        return response.json()["choices"][0]["message"]["content"]


async def synthesize_answer(query: str, results: list[SearchResult]) -> str:
    if not results:
        return "No search results available to synthesize an answer."
    try:
        return await _chat(
            [
                {"role": "system", "content": _SYSTEM_PROMPT},
                {"role": "user", "content": f"Query: {query}\n\nSearch Results:\n{_context_block(results, 5)}"},
            ]
        )
    except Exception as exc:
        logger.error("Answer synthesis failed: %s", exc)
        return f"Answer synthesis failed: {exc}"


async def generate_sub_queries(query: str, count: int = 3) -> list[str]:
    try:
        text = await _chat(
            [
                {
                    "role": "system",
                    "content": (
                        f"Generate exactly {count} distinct search sub-queries to research the given topic. "
                        "Each sub-query explores a different angle. "
                        "Return ONLY the queries, one per line, no numbering or explanation."
                    ),
                },
                {"role": "user", "content": query},
            ],
            max_tokens=256,
        )
        queries = [q.strip() for q in text.strip().splitlines() if q.strip()]
        return queries[:count]
    except Exception as exc:
        logger.error("Sub-query generation failed: %s", exc)
        return [query, f"{query} overview", f"{query} recent developments"][:count]


async def synthesize_report(
    query: str,
    results: list[SearchResult],
    output_format: str = "markdown",
) -> str:
    if not results:
        return "No search results available for report generation."

    format_instruction = (
        "Format as valid JSON with keys: summary (string), key_findings (list of strings), conclusions (string)."
        if output_format == "json"
        else "Format as clear markdown with sections: ## Summary, ## Key Findings, ## Conclusion."
    )

    try:
        return await _chat(
            [
                {"role": "system", "content": _SYSTEM_PROMPT},
                {
                    "role": "user",
                    "content": (
                        f"Topic: {query}\n\nSearch Results:\n{_context_block(results)}\n\n"
                        f"Write a comprehensive research report. {format_instruction} "
                        "Cite sources using [N] notation."
                    ),
                },
            ],
            max_tokens=2048,
            timeout=60.0,
        )
    except Exception as exc:
        logger.error("Report synthesis failed: %s", exc)
        return f"Report synthesis failed: {exc}"