"""
Async write queue — background processing of saliency-triggered and
cue-triggered write requests.

Cloud LLM calls and URD inserts never block the proxy response path.
"""
from __future__ import annotations

import asyncio
import logging
from dataclasses import dataclass
from typing import Optional

import asyncpg

from .cache import SoasRow
from .cue_scanner import CueTriple
from .db import get_or_create_soas, get_config, create_novel_soas, NOVEL_CONFIRMED_SALIENCY
from .llm_client import (
    ModelConfig, get_model_config, call_llm,
    CONTEXT_DISCOVER_PROMPT_TEMPLATE, parse_context_triples,
)
from .urd_writer import InsertRequest, insert_urd_edge
from . import cache

log = logging.getLogger("festinger.write_queue")


@dataclass
class ContextDiscoverRequest:
    """
    User message text plus candidate tokens (words absent from the English
    dictionary) submitted for LLM-driven concept discovery.

    candidate_tokens: hints for the LLM — it decides which are real domain
      concepts vs typos/noise, and extracts relationship triples from the text.

    agent_name: normalised agent identity from X-Agent-Name header (lowercase).
      Used to look up an agent-specific model in the agent_models table.

    fallback_model: last-resort ModelConfig derived from the intercepted request.
      Used when neither an agent-specific model nor write_model_id is configured.
    """
    user_text: str
    candidate_tokens: list[str]
    agent_name: str = ""
    fallback_model: Optional[ModelConfig] = None


@dataclass
class CueWriteRequest:
    """A directly extracted CueTriple — bypasses LLM, goes straight to URD insert."""
    triple: CueTriple


_queue: asyncio.Queue = asyncio.Queue(maxsize=1000)
_running: bool = False

# Limit concurrent LM Studio / cloud-LLM calls so we don't flood a slow
# local model server.  Cue processing (pure DB ops) is never rate-limited.
# Local models (LM Studio, Ollama) can typically only handle one request at a
# time.  Keep this at 1 to avoid overwhelming a single-GPU inference server.
_LLM_CONCURRENCY = 1
_llm_semaphore: asyncio.Semaphore | None = None


async def enqueue_context_discover(
    user_text: str,
    candidate_tokens: list[str],
    agent_name: str = "",
    fallback_model: Optional[ModelConfig] = None,
) -> None:
    """
    Enqueue a user message for LLM-driven concept discovery and relation extraction.

    agent_name: normalised agent identity (from X-Agent-Name header).
      Festinger looks up an agent-specific model in the agent_models table.

    fallback_model: last-resort ModelConfig from the intercepted request.
      Used when neither an agent-specific model nor write_model_id is configured.
    """
    try:
        _queue.put_nowait(ContextDiscoverRequest(
            user_text=user_text,
            candidate_tokens=candidate_tokens,
            agent_name=agent_name,
            fallback_model=fallback_model,
        ))
    except asyncio.QueueFull:
        log.warning("write queue full — dropping context discover")


async def enqueue_cue(triple: CueTriple) -> None:
    try:
        _queue.put_nowait(CueWriteRequest(triple=triple))
    except asyncio.QueueFull:
        log.warning("write queue full — dropping cue: %s", triple)


async def start_worker(pool: asyncpg.Pool) -> None:
    """Launch background worker. Call once at startup."""
    global _running, _llm_semaphore
    _running = True
    _llm_semaphore = asyncio.Semaphore(_LLM_CONCURRENCY)
    asyncio.create_task(_worker(pool))
    log.info("write queue worker started")


async def _worker(pool: asyncpg.Pool) -> None:
    """
    Drain the queue as fast as possible.

    Cue writes (pure DB ops) are awaited inline — they're fast.
    Concept LLM calls are dispatched as fire-and-forget asyncio tasks so
    a slow LM Studio response never stalls the worker loop.
    """
    while _running:
        try:
            item = await asyncio.wait_for(_queue.get(), timeout=5.0)
        except asyncio.TimeoutError:
            continue

        try:
            if isinstance(item, CueWriteRequest):
                # Fast path: no LLM involved, process inline.
                await _process_cue(pool, item.triple)
            elif isinstance(item, ContextDiscoverRequest):
                # Slow path: fire off without awaiting so the worker stays free.
                asyncio.create_task(
                    _process_context_discover_guarded(
                        pool, item.user_text, item.candidate_tokens,
                        item.agent_name, item.fallback_model,
                    )
                )
        except Exception as e:
            log.exception("write queue worker error: %s", e)
        finally:
            _queue.task_done()


async def _process_context_discover_guarded(
    pool: asyncpg.Pool,
    user_text: str,
    candidate_tokens: list[str],
    agent_name: str = "",
    fallback_model: Optional[ModelConfig] = None,
) -> None:
    """Wrapper that acquires the LLM semaphore before concept discovery."""
    assert _llm_semaphore is not None
    async with _llm_semaphore:
        try:
            await _process_context_discover(pool, user_text, candidate_tokens, agent_name, fallback_model)
        except Exception as e:
            log.exception("context discover task error: %s", e)


async def stop_worker() -> None:
    global _running
    _running = False


# ---------------------------------------------------------------------------
# Processing
# ---------------------------------------------------------------------------

async def _process_cue(pool: asyncpg.Pool, triple: CueTriple) -> None:
    """Insert a cue-extracted triple directly into URD."""
    subj_row = await get_or_create_soas(pool, triple.subject)
    parent_row = await get_or_create_soas(pool, triple.parent)
    dim_row = await get_or_create_soas(pool, triple.dimension)

    req = InsertRequest(
        concept_id=subj_row.id,
        parent_id=parent_row.id,
        dim_id=dim_row.id,
        is_isa=triple.is_isa,
        confidence=triple.confidence,
        source="inferred",
    )
    collision = await insert_urd_edge(pool, req)
    if collision:
        log.info("cue triple collision: %s", collision)


async def _resolve_discover_model(
    pool: asyncpg.Pool,
    agent_name: str,
    fallback_model: Optional[ModelConfig],
) -> Optional[ModelConfig]:
    """
    Resolve which LLM to use for context discovery.

    Priority:
      1. Agent-specific model  — agent_models table, keyed by agent_name
      2. Global default        — write_model_id config key
      3. Request mirror        — fallback_model (same provider/model Agent0 used)
    """
    if agent_name:
        async with pool.acquire() as conn:
            row = await conn.fetchrow(
                """
                SELECT m.provider, m.model_name, m.api_key, m.base_url
                FROM agent_models am
                JOIN models m ON m.id = am.model_id
                WHERE am.agent_name = $1
                """,
                agent_name,
            )
        if row:
            log.debug(
                "context discover: agent=%s → provider=%s model=%s",
                agent_name, row["provider"], row["model_name"],
            )
            return ModelConfig(
                provider=row["provider"],
                model_name=row["model_name"],
                api_key=row["api_key"],
                base_url=row["base_url"] or "",
            )

    write_model_id = await get_config(pool, "write_model_id")
    if write_model_id:
        m = await get_model_config(pool, write_model_id)
        if m:
            return m
        log.warning("write_model_id=%s not found in models table", write_model_id)

    if fallback_model:
        log.debug(
            "context discover: mirroring request model provider=%s model=%s",
            fallback_model.provider, fallback_model.model_name,
        )
        return fallback_model

    return None


async def _process_context_discover(
    pool: asyncpg.Pool,
    user_text: str,
    candidate_tokens: list[str],
    agent_name: str = "",
    fallback_model: Optional[ModelConfig] = None,
) -> None:
    """
    Ask the LLM to evaluate candidate tokens (dictionary misses) and extract
    relationship triples for those it judges to be real domain concepts.

    Model selection priority (see _resolve_discover_model):
      1. Agent-specific model (agent_models table)
      2. write_model_id config (global default)
      3. Request mirror (fallback_model)

    Saliency feedback loop:
      - Confirmed concepts (triples inserted) → saliency raised to NOVEL_CONFIRMED_SALIENCY
        so they surface as recollection hits on subsequent turns.
      - Rejected concepts (no triples) → saliency stays low (NOVEL_INITIAL_SALIENCY),
        effectively hiding typos and noise from the recollection engine.
    """
    model = await _resolve_discover_model(pool, agent_name, fallback_model)
    if not model:
        log.debug("no model resolved for context discover (agent=%r) — skipping", agent_name)
        return

    seed_dims = ["type", "membership", "runs-on", "tech", "owned-by", "geography"]
    prompt = CONTEXT_DISCOVER_PROMPT_TEMPLATE.format(
        candidates=", ".join(candidate_tokens),
        context=user_text[:2000],
        dimensions=", ".join(seed_dims),
    )

    try:
        response = await call_llm(model, prompt)
    except Exception as e:
        log.warning("LLM call failed for context discover: %s", e)
        return

    triples = parse_context_triples(response)
    if not triples:
        log.debug("LLM found no domain concepts among candidates: %s", candidate_tokens)
        return

    MIN_CONFIDENCE = 0.6
    candidate_set = set(candidate_tokens)
    confirmed: set[str] = set()   # candidates confirmed as real concepts

    for t in triples:
        if not t.concept or not t.parent or not t.dimension:
            continue
        if t.concept not in candidate_set:
            log.debug("context discover: ignoring concept not in candidate list: %r", t.concept)
            continue
        if t.confidence < MIN_CONFIDENCE:
            log.debug("context discover: low confidence triple skipped: %s", t)
            continue

        subj_row = await get_or_create_soas(pool, t.concept)
        parent_row = await get_or_create_soas(pool, t.parent)
        dim_row = await get_or_create_soas(pool, t.dimension)

        req = InsertRequest(
            concept_id=subj_row.id,
            parent_id=parent_row.id,
            dim_id=dim_row.id,
            is_isa=t.is_isa,
            confidence=t.confidence,
            source="context_llm",
        )
        collision = await insert_urd_edge(pool, req)
        if not collision:
            confirmed.add(t.concept)

    # Raise saliency for confirmed concepts so recollection fires on next encounter.
    for concept in confirmed:
        async with pool.acquire() as conn:
            await conn.execute(
                "UPDATE soas SET saliency = $1 WHERE token = $2 AND saliency < $1",
                NOVEL_CONFIRMED_SALIENCY, concept,
            )
        row = cache.soas_by_token.get(concept)
        if row and row.saliency < NOVEL_CONFIRMED_SALIENCY:
            row.saliency = NOVEL_CONFIRMED_SALIENCY

    rejected = candidate_set - confirmed
    log.info(
        "context discover: candidates=%s confirmed=%s rejected=%s",
        candidate_tokens, sorted(confirmed), sorted(rejected),
    )
