# retoor <retoor@molodetz.nl>

import asyncio
import logging
import threading
from concurrent.futures import ThreadPoolExecutor

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_setting, get_table
from devplacepy.services.background import background
from devplacepy.services.openai_gateway.usage import parse_usage_headers

logger = logging.getLogger(__name__)

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

CORRECTION_TIMEOUT_SECONDS = 20.0
MAX_GROWTH_FACTOR = 3
PENDING_SCOPE_KEY = "devplace_pending_corrections"

AI_APPLY_EXECUTOR = ThreadPoolExecutor(max_workers=4, thread_name_prefix="ai-apply")

_gateway_client: httpx.Client | None = None
_gateway_client_lock = threading.Lock()


def _client() -> httpx.Client:
    global _gateway_client
    if _gateway_client is None:
        with _gateway_client_lock:
            if _gateway_client is None:
                _gateway_client = stealth.stealth_sync_client(
                    timeout=CORRECTION_TIMEOUT_SECONDS
                )
    return _gateway_client


def _usage_from_headers(response_headers) -> dict:
    parsed = parse_usage_headers(response_headers)
    if parsed is None:
        return {
            "calls": 1,
            "prompt_tokens": 0,
            "completion_tokens": 0,
            "total_tokens": 0,
            "cost_usd": 0.0,
            "upstream_latency_ms": 0.0,
            "total_latency_ms": 0.0,
        }
    return {
        "calls": parsed["calls"],
        "prompt_tokens": parsed["prompt_tokens"],
        "completion_tokens": parsed["completion_tokens"],
        "total_tokens": parsed["total_tokens"],
        "cost_usd": parsed["cost_usd"],
        "upstream_latency_ms": parsed["upstream_latency_ms"],
        "total_latency_ms": parsed["total_latency_ms"],
    }


def gateway_complete(
    api_key: str,
    system: str,
    text: str,
    timeout: float,
    max_growth_factor: int | None = None,
    model: str = INTERNAL_MODEL,
) -> tuple[str, dict | None]:
    text = text or ""
    if not text.strip():
        return text, None
    payload = {
        "model": model or INTERNAL_MODEL,
        "messages": [
            {"role": "system", "content": system},
            {"role": "user", "content": text},
        ],
        "temperature": 0.1,
    }
    headers = {
        "Content-Type": "application/json",
        "X-App-Reference": "devplace-correction-v-1-0-0",
    }
    if api_key:
        headers["Authorization"] = f"Bearer {api_key}"
    try:
        response = _client().post(
            INTERNAL_GATEWAY_URL, json=payload, headers=headers, timeout=timeout
        )
        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 correction_model() -> str:
    return get_setting("correction_model", "") or INTERNAL_MODEL


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,
        model=correction_model(),
    )


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 schedule_pending(
        _run_correction, request, api_key, prompt, table, uid, user_uid
    ):
        return
    background.submit(_run_correction, api_key, prompt, table, uid, user_uid)


def schedule_pending(fn, request: object, *args) -> 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(AI_APPLY_EXECUTOR, fn, *args)
    scope.setdefault(PENDING_SCOPE_KEY, []).append(future)
    return True


def new_usage_totals() -> dict:
    return {
        "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,
    }


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 = new_usage_totals()
    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.get(key, 0)
        if corrected and corrected != original:
            updates[field] = corrected
    if updates:
        updates["uid"] = uid
        get_table(table).update(updates, ["uid"])
        if table == "messages":
            from devplacepy.services.messaging.persist import push_content_revision

            push_content_revision(uid, ai_processed=True)
        if table == "users" and user_uid:
            from devplacepy.utils import clear_user_cache

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