from __future__ import annotations

import asyncio
import os
from typing import Sequence

import litellm

from wishful.config import settings
from wishful.llm.prompts import build_messages, strip_code_fences
from wishful.logging import logger


class GenerationError(ImportError):
    """Raised when the LLM call fails or returns empty output."""


_FAKE_MODE = os.getenv("WISHFUL_FAKE_LLM", "0") == "1"


def _fake_response(functions: Sequence[str]) -> str:
    body = []
    for name in functions or ("generated_helper",):
        body.append(
            f"def {name}(*args, **kwargs):\n    \"\"\"Auto-generated placeholder. Replace with real logic.\"\"\"\n    return {{'args': args, 'kwargs': kwargs}}\n"
        )
    return "\n\n".join(body)


def generate_module_code(
    module: str,
    functions: Sequence[str],
    context: str | None,
    type_schemas: dict[str, str] | None = None,
    function_output_types: dict[str, str] | None = None,
    mode: str | None = None,
) -> str:
    """Call the LLM (or fake stub) to generate module source code (sync version)."""

    if _FAKE_MODE:
        return _fake_response(functions)

    response = _call_llm(
        module, functions, context, type_schemas, function_output_types, mode
    )
    content = _extract_content(response)
    return strip_code_fences(content).strip()


async def agenerate_module_code(
    module: str,
    functions: Sequence[str],
    context: str | None,
    type_schemas: dict[str, str] | None = None,
    function_output_types: dict[str, str] | None = None,
    mode: str | None = None,
) -> str:
    """Call the LLM (or fake stub) to generate module source code (async version).
    
    This is the preferred method for explore() and other features that need
    concurrent operations or responsive UI updates.
    """

    if _FAKE_MODE:
        # Small delay to simulate async behavior in fake mode
        await asyncio.sleep(0.01)
        return _fake_response(functions)

    response = await _acall_llm(
        module, functions, context, type_schemas, function_output_types, mode
    )
    content = _extract_content(response)
    return strip_code_fences(content).strip()


def _call_llm(
    module: str,
    functions: Sequence[str],
    context: str | None,
    type_schemas: dict[str, str] | None = None,
    function_output_types: dict[str, str] | None = None,
    mode: str | None = None,
):
    """Synchronous LLM call."""
    messages = build_messages(
        module, functions, context, type_schemas, function_output_types, mode
    )
    _log_llm_call(module, mode, functions, context, type_schemas, function_output_types, messages)
    
    try:
        return litellm.completion(
            model=settings.model,
            messages=messages,
            temperature=settings.temperature,
            max_tokens=settings.max_tokens,
        )
    except Exception as exc:  # pragma: no cover - network path not executed in tests
        raise GenerationError(f"LLM call failed: {exc}") from exc


async def _acall_llm(
    module: str,
    functions: Sequence[str],
    context: str | None,
    type_schemas: dict[str, str] | None = None,
    function_output_types: dict[str, str] | None = None,
    mode: str | None = None,
):
    """Asynchronous LLM call using litellm.acompletion()."""
    messages = build_messages(
        module, functions, context, type_schemas, function_output_types, mode
    )
    _log_llm_call(module, mode, functions, context, type_schemas, function_output_types, messages)
    
    try:
        return await litellm.acompletion(
            model=settings.model,
            messages=messages,
            temperature=settings.temperature,
            max_tokens=settings.max_tokens,
        )
    except Exception as exc:  # pragma: no cover - network path not executed in tests
        raise GenerationError(f"LLM call failed: {exc}") from exc


def _log_llm_call(
    module: str,
    mode: str | None,
    functions: Sequence[str],
    context: str | None,
    type_schemas: dict[str, str] | None,
    function_output_types: dict[str, str] | None,
    messages: list,
) -> None:
    """Log LLM call details."""
    logger.debug(
        "LLM call module={} mode={} model={} temp={} max_tokens={} functions={} context_len={} type_schemas={} output_types={} preview={}",
        module,
        mode,
        settings.model,
        settings.temperature,
        settings.max_tokens,
        list(functions),
        len(context) if context else 0,
        list((type_schemas or {}).keys()),
        list((function_output_types or {}).keys()),
        (context[:500] + "…" if context and len(context) > 500 else (context or "")),
    )

    # Log the actual prompt messages (truncated for safety)
    prompt_text = "\n".join(f"{m['role'].upper()}: {m['content']}" for m in messages)
    if len(prompt_text) > 4000:
        prompt_text = prompt_text[:4000] + "…"
    logger.debug("LLM prompt for {}:\n{}", module, prompt_text)


def _extract_content(response) -> str:
    try:
        content = response["choices"][0]["message"]["content"]
    except Exception as exc:  # pragma: no cover
        raise GenerationError("Unexpected LLM response structure") from exc

    if not content or not content.strip():
        raise GenerationError("LLM returned empty content")
    return content
