"""
utils.py
-------------
Shared utility functions used across the RAG pipeline.

format_prompt — applies Qwen's official Jinja chat template
clean_final_answer  — fixes common Qwen URL formatting bugs
extract_sources — deterministically appends source URLs to answers
"""
import re
from model_setup import tokenizer


def format_prompt(system: str, user: str) -> str:
  """
  Apply Qwen 2.5-Instruct's official chat template via the HF tokenizer.

  Uses tokenizer.apply_chat_template() — a CPU-side string operation,
  no CUDA involved. Safe to call from anywhere, including outside
  the @spaces.GPU-decorated function.
  """
  messages = [
      {"role": "system", "content": system},
      {"role": "user",   "content": user},
  ]
  return tokenizer.apply_chat_template(
      messages,
      tokenize=False,
      add_generation_prompt=True,
  )


def clean_final_answer(text: str) -> str:
  """
  Fix two output formatting bugs common in quantized Qwen models.

  Bug 1 — Trailing whitespace inside markdown link parentheses.
  Bug 2 — Duplicate citation lines (Qwen repeats URLs to "seem thorough").
  """
  text = re.sub(
    r'\((https?://[^\s\)]+?)[\s\n]*\)',
    lambda m: f'({m.group(1)})',
    text,
  )

  seen_url_sets = set()
  deduped_lines = []
  for line in text.split('\n'):
    urls_on_line = re.findall(r'https?://[^\s\)\]]+', line)
    if urls_on_line:
      key = frozenset(urls_on_line)
      if key in seen_url_sets:
        continue
      seen_url_sets.add(key)
    deduped_lines.append(line)

  return '\n'.join(deduped_lines)


def extract_sources(state: dict) -> str:
  """
  Deterministically pull source URLs from the retrieved context and
  return a formatted Sources section to append to the final answer.

  Guarantees sources appear regardless of model citation behaviour —
  the URLs are already in state from retrieval, this just surfaces them.
  """
  urls_seen = set()
  sources   = []

  web_results = state.get("web_results", "")
  for line in web_results.split("\n"):
    stripped = line.strip()
    if stripped.startswith("URL:"):
      url = stripped.replace("URL:", "").strip()
      if url and url not in urls_seen:
        urls_seen.add(url)
        sources.append(url)

  retrieved_docs = state.get("retrieved_docs", "")
  for match in re.finditer(r'\[Source:\s*(https?://[^\]]+)\]', retrieved_docs):
    url = match.group(1).strip()
    if url and url not in urls_seen:
      urls_seen.add(url)
      sources.append(url)

  if not sources:
    return ""

  lines = ["\n\n---\n**Sources:**"]
  for i, url in enumerate(sources, 1):
    lines.append(f"{i}. {url}")

  return "\n".join(lines)