#!/usr/bin/env python3
"""Query the local homelab retrieval index and optionally ask Ollama."""

from __future__ import annotations

import argparse
import json
import math
import os
import pathlib
import re
import sys
import urllib.error
import urllib.request
from collections import Counter


DEFAULT_INDEX_DIR = pathlib.Path(os.environ.get("LAB_AI_KNOWLEDGE_INDEX_DIR", "/data/homelab-ai/index"))
TOKEN_RE = re.compile(r"[A-Za-z0-9_./:-]{2,}")


def tokenise(text: str) -> list[str]:
    return [token.lower() for token in TOKEN_RE.findall(text)]


def load_index(index_dir: pathlib.Path) -> dict[str, object]:
    index_path = index_dir / "index.json"
    if not index_path.is_file():
        raise FileNotFoundError(f"missing index: {index_path}")
    return json.loads(index_path.read_text(encoding="utf-8"))


def score_chunks(index: dict[str, object], query: str, limit: int) -> list[tuple[float, dict[str, object]]]:
    query_terms = tokenise(query)
    if not query_terms:
        return []

    query_counts = Counter(query_terms)
    idf = index.get("idf", {})
    average_length = float(index.get("average_length") or 1.0)
    chunks = index.get("chunks", [])
    results: list[tuple[float, dict[str, object]]] = []
    k1 = 1.5
    b = 0.75

    if not isinstance(idf, dict) or not isinstance(chunks, list):
        return []

    for chunk in chunks:
        if not isinstance(chunk, dict):
            continue
        tokens = chunk.get("tokens", [])
        if not isinstance(tokens, list) or not tokens:
            continue
        term_counts = Counter(str(token) for token in tokens)
        doc_len = len(tokens)
        score = 0.0
        for term, query_count in query_counts.items():
            frequency = term_counts.get(term, 0)
            if frequency == 0:
                continue
            term_idf = float(idf.get(term, 0.0))
            denominator = frequency + k1 * (1 - b + b * doc_len / average_length)
            score += query_count * term_idf * (frequency * (k1 + 1)) / denominator
        if score > 0:
            results.append((score, chunk))

    return sorted(results, key=lambda item: item[0], reverse=True)[:limit]


def render_context(results: list[tuple[float, dict[str, object]]], max_chars: int) -> str:
    sections: list[str] = []
    remaining = max_chars

    for score, chunk in results:
        header = f"[{chunk.get('path')}:{chunk.get('start_line')}] score={score:.3f}"
        text = str(chunk.get("text", "")).strip()
        section = f"{header}\n{text}"
        if len(section) > remaining:
            section = section[: max(0, remaining - 20)].rstrip() + "\n[truncated]"
        sections.append(section)
        remaining -= len(section)
        if remaining <= 0:
            break

    return "\n\n---\n\n".join(sections)


def ask_ollama(question: str, context: str, endpoint: str, model: str, timeout: int) -> str:
    prompt = f"""You are helping operate a personal homelab.
Use only the context below. If the context is insufficient, say what to inspect next.
Be concise, factual, and command-oriented.

Question:
{question}

Context:
{context}
"""
    payload = {
        "model": model,
        "prompt": prompt,
        "stream": False,
        "options": {
            "temperature": 0.1,
            "num_predict": 360,
        },
    }
    request = urllib.request.Request(
        f"{endpoint.rstrip('/')}/api/generate",
        data=json.dumps(payload).encode("utf-8"),
        headers={"Content-Type": "application/json"},
        method="POST",
    )
    with urllib.request.urlopen(request, timeout=timeout) as response:
        result = json.loads(response.read().decode("utf-8"))
    return str(result.get("response", "")).strip()


def main() -> int:
    parser = argparse.ArgumentParser(description=__doc__)
    parser.add_argument("query", nargs="+", help="Question or search text")
    parser.add_argument("--index-dir", type=pathlib.Path, default=DEFAULT_INDEX_DIR)
    parser.add_argument("--limit", type=int, default=5)
    parser.add_argument("--max-chars", type=int, default=9000)
    parser.add_argument("--context-only", action="store_true")
    parser.add_argument("--ask", action="store_true")
    parser.add_argument("--ollama-url", default=os.environ.get("LAB_AI_GATEWAY_URL", "http://127.0.0.1:11434"))
    parser.add_argument("--model", default=os.environ.get("LAB_AI_GATEWAY_MODEL", "qwen2.5:0.5b"))
    parser.add_argument("--timeout", type=int, default=int(os.environ.get("LAB_AI_GATEWAY_TIMEOUT_SECONDS", "20")))
    args = parser.parse_args()

    query = " ".join(args.query)
    try:
        index = load_index(args.index_dir)
    except FileNotFoundError as exc:
        print(str(exc), file=sys.stderr)
        return 1

    results = score_chunks(index, query, args.limit)
    if not results:
        print("no matching homelab knowledge chunks found", file=sys.stderr)
        return 1

    context = render_context(results, args.max_chars)
    if args.context_only:
        print(context)
        return 0

    if args.ask:
        try:
            answer = ask_ollama(query, context, args.ollama_url, args.model, args.timeout)
        except (OSError, TimeoutError, urllib.error.URLError, json.JSONDecodeError) as exc:
            print(f"ollama request failed: {exc}", file=sys.stderr)
            return 1
        print(answer)
        return 0

    print(context)
    return 0


if __name__ == "__main__":
    raise SystemExit(main())