"""Weft memory integration — inject relevant memories into task context at claim time."""

from __future__ import annotations

import logging
from typing import TYPE_CHECKING, Any

if TYPE_CHECKING:
    from loom.config import WeftConfig

logger = logging.getLogger(__name__)


async def init_weft(weft_cfg: WeftConfig) -> tuple[Any, Any]:
    """Create Weft asyncpg pool, run migrations, and warm up embedder.

    Returns (pool, embedder) tuple. All weft imports are lazy so this module
    is importable even without weft-memory installed.
    """
    import asyncpg
    from weft.db.migrations import run_migrations
    from weft.embeddings import get_provider

    db_url = weft_cfg.database_url or None  # None = Weft's default
    pool = await asyncpg.create_pool(db_url)
    await run_migrations(pool)
    embedder = get_provider("fastembed")
    # Warm up the embedding model with a throwaway encode
    await embedder.embed("warmup")
    logger.info("Weft integration initialized (budget_tokens=%d)", weft_cfg.budget_tokens)
    return pool, embedder


async def inject_weft_memories(
    task_dict: dict,
    weft_pool: Any,
    weft_embedder: Any,
    budget_tokens: int = 500,
    project_id: str | None = None,
) -> dict:
    """Query Weft for relevant memories and inject into task context.

    Builds a search query from task title + description + files, embeds it,
    fetches matching memories within the token budget, and adds them to
    ``task_dict["context"]["weft_memories"]``.

    Never raises — all errors are caught and logged so claims always succeed.
    """
    try:
        from weft.context import build_context

        # Build query from task metadata
        parts: list[str] = []
        if title := task_dict.get("title"):
            parts.append(title)
        ctx = task_dict.get("context")
        if isinstance(ctx, dict):
            if desc := ctx.get("description"):
                parts.append(desc)
            if files := ctx.get("files"):
                parts.append(" ".join(files[:10]))

        query = " ".join(parts).strip()
        if not query:
            return task_dict

        # Embed query string into vector
        query_embedding = await weft_embedder.embed(query)

        result = await build_context(
            weft_pool,
            query_embedding,
            budget_tokens=budget_tokens,
            project_id=project_id,
        )

        memories = result.get("memories", [])
        if not memories:
            return task_dict

        # Inject into context dict
        if not isinstance(task_dict.get("context"), dict):
            task_dict["context"] = task_dict.get("context") or {}
        task_dict["context"]["weft_memories"] = memories

    except Exception:
        logger.debug("Weft memory injection skipped", exc_info=True)

    return task_dict


async def route_codebase_notes(
    weft_pool: Any,
    weft_embedder: Any,
    codebase_notes: dict[str, str],
    project_id: str | None = None,
) -> int:
    """Store codebase_notes as Weft file-level memories.

    Each entry in codebase_notes is {filepath: summary}. Stored as
    type=fact, source=ingest, topic=["file:<path>", project_id].

    Returns count of memories stored. Never raises.
    """
    stored = 0
    try:
        from weft.models import MemoryCreate, MemorySource, MemoryType
        from weft.store import store_memory

        for filepath, summary in codebase_notes.items():
            if not summary or not filepath:
                continue
            topic = [f"file:{filepath}"]
            if project_id:
                topic.append(project_id)
            create = MemoryCreate(
                type=MemoryType.fact,
                content=f"{filepath}: {summary}",
                topic=topic,
                source=MemorySource.ingest,
                confidence=0.8,
                project_id=project_id,
            )
            embedding = await weft_embedder.embed(create.content)
            await store_memory(weft_pool, create, embedding=embedding)
            stored += 1

    except Exception:
        logger.warning(
            "Failed to route codebase_notes (%d/%d stored)",
            stored, len(codebase_notes), exc_info=True,
        )

    return stored
