# retoor <retoor@molodetz.nl>

import asyncio
import logging

import httpx

from devplacepy import stealth
from devplacepy.config import (
    DEFAULT_CORRECTION_PROMPT,
    INTERNAL_GATEWAY_URL,
    INTERNAL_MODEL,
)
from devplacepy.database import add_correction_usage, get_table
from devplacepy.services.background import background

logger = logging.getLogger(__name__)

CORRECTABLE_FIELDS: dict[str, tuple[str, ...]] = {
    "posts": ("title", "content"),
    "projects": ("title", "description"),
    "gists": ("title", "description"),
    "comments": ("content",),
    "messages": ("content",),
    "users": ("bio",),
}

CORRECTION_TIMEOUT_SECONDS = 20.0
MAX_GROWTH_FACTOR = 3
PENDING_SCOPE_KEY = "devplace_pending_corrections"


def _usage_from_headers(response_headers) -> dict:
    def _int(name: str) -> int:
        try:
            return int(response_headers.get(name) or 0)
        except (TypeError, ValueError):
            return 0

    def _float(name: str) -> float:
        try:
            return float(response_headers.get(name) or 0.0)
        except (TypeError, ValueError):
            return 0.0

    return {
        "calls": 1,
        "prompt_tokens": _int("X-Gateway-Prompt-Tokens"),
        "completion_tokens": _int("X-Gateway-Completion-Tokens"),
        "total_tokens": _int("X-Gateway-Total-Tokens"),
        "cost_usd": _float("X-Gateway-Cost-USD"),
        "upstream_latency_ms": _float("X-Gateway-Upstream-Latency-Ms"),
        "total_latency_ms": _float("X-Gateway-Total-Latency-Ms"),
    }


def gateway_complete(
    api_key: str,
    system: str,
    text: str,
    timeout: float,
    max_growth_factor: int | None = None,
) -> tuple[str, dict | None]:
    text = text or ""
    if not text.strip():
        return text, None
    payload = {
        "model": INTERNAL_MODEL,
        "messages": [
            {"role": "system", "content": system},
            {"role": "user", "content": text},
        ],
        "temperature": 0.1,
    }
    headers = {"Content-Type": "application/json"}
    if api_key:
        headers["Authorization"] = f"Bearer {api_key}"
    try:
        with stealth.stealth_sync_client(timeout=timeout) as client:
            response = client.post(INTERNAL_GATEWAY_URL, json=payload, headers=headers)
            response.raise_for_status()
            usage = _usage_from_headers(response.headers)
            content = (
                response.json()
                .get("choices", [{}])[0]
                .get("message", {})
                .get("content", "")
                .strip()
            )
    except (httpx.HTTPError, ValueError, KeyError, IndexError) as exc:
        logger.warning("AI gateway completion failed, keeping original: %s", exc)
        return text, None
    if not content:
        return text, usage
    if max_growth_factor and len(content) > len(text) * max_growth_factor + 200:
        logger.warning("AI gateway output too large, keeping original")
        return text, usage
    return content, usage


def correct_text(api_key: str, prompt: str, text: str) -> tuple[str, dict | None]:
    system = (
        "You are a text correction engine. Apply the correction instruction below to "
        "the user's message and return ONLY the resulting text, with no preamble, no "
        "explanation, no quotes and no code fences. Preserve the original meaning, "
        "language, line breaks and markdown. Correction instruction: "
        + (prompt or DEFAULT_CORRECTION_PROMPT).strip()
    )
    return gateway_complete(
        api_key, system, text, CORRECTION_TIMEOUT_SECONDS, MAX_GROWTH_FACTOR
    )


def schedule_correction(
    user: dict | None, table: str, uid: str, request: object = None
) -> None:
    if not user or table not in CORRECTABLE_FIELDS or not uid:
        return
    if not user.get("ai_correction_enabled"):
        return
    api_key = (user.get("api_key") or "").strip()
    if not api_key:
        return
    prompt = (user.get("ai_correction_prompt") or DEFAULT_CORRECTION_PROMPT).strip()
    user_uid = user.get("uid") or ""
    if user.get("ai_correction_sync") and _run_inline_awaited(
        api_key, prompt, table, uid, user_uid, request
    ):
        return
    background.submit(_run_correction, api_key, prompt, table, uid, user_uid)


def _run_inline_awaited(
    api_key: str,
    prompt: str,
    table: str,
    uid: str,
    user_uid: str,
    request: object,
) -> bool:
    scope = getattr(request, "scope", None)
    if scope is None:
        return False
    try:
        loop = asyncio.get_running_loop()
    except RuntimeError:
        return False
    future = loop.run_in_executor(
        None, _run_correction, api_key, prompt, table, uid, user_uid
    )
    scope.setdefault(PENDING_SCOPE_KEY, []).append(future)
    return True


def _run_correction(
    api_key: str, prompt: str, table: str, uid: str, user_uid: str
) -> None:
    fields = CORRECTABLE_FIELDS.get(table)
    if not fields:
        return
    row = get_table(table).find_one(uid=uid)
    if not row:
        return
    updates: dict = {}
    totals = {
        "calls": 0,
        "prompt_tokens": 0,
        "completion_tokens": 0,
        "total_tokens": 0,
        "cost_usd": 0.0,
        "upstream_latency_ms": 0.0,
        "total_latency_ms": 0.0,
    }
    for field in fields:
        original = row.get(field) or ""
        if not original.strip():
            continue
        corrected, usage = correct_text(api_key, prompt, original)
        if usage:
            for key in totals:
                totals[key] += usage[key]
        if corrected and corrected != original:
            updates[field] = corrected
    if updates:
        updates["uid"] = uid
        get_table(table).update(updates, ["uid"])
        if table == "users" and user_uid:
            from devplacepy.utils import clear_user_cache

            clear_user_cache(user_uid)
    if totals["calls"] and user_uid:
        add_correction_usage(user_uid, totals)
