from __future__ import annotations

import os
import re

MAX_CHAT_PAPERS = int(os.getenv("MAX_CHAT_PAPERS", "4"))
CHAT_SUMMARY_MAX_CHARS = int(os.getenv("CHAT_SUMMARY_MAX_CHARS", "14000"))

_SECTION_HEADING = re.compile(r"^##\s+\d+\.", re.MULTILINE)
_QUESTION_STOPWORDS = frozenset(
    {
        "about",
        "also",
        "based",
        "compare",
        "could",
        "does",
        "first",
        "from",
        "have",
        "paper",
        "papers",
        "please",
        "provide",
        "range",
        "second",
        "source",
        "summarize",
        "summary",
        "that",
        "the",
        "their",
        "these",
        "third",
        "this",
        "those",
        "was",
        "were",
        "what",
        "when",
        "which",
        "would",
        "your",
    }
)
_ORDINAL_PAPER = re.compile(
    r"\b(?P<ordinal>first|1st|second|2nd|third|3rd|fourth|4th)\s+paper\b",
    re.I,
)
_ORDINAL_INDEX = {
    "first": 0,
    "1st": 0,
    "second": 1,
    "2nd": 1,
    "third": 2,
    "3rd": 2,
    "fourth": 3,
    "4th": 3,
}


def focus_document_ids_for_question(message: str, document_ids: list[str]) -> list[str]:
    """When the user names a specific paper, answer from that paper only."""
    if len(document_ids) <= 1:
        return document_ids
    match = _ORDINAL_PAPER.search(message or "")
    if not match:
        return document_ids
    index = _ORDINAL_INDEX.get(match.group("ordinal").lower())
    if index is None or index >= len(document_ids):
        return document_ids
    return [document_ids[index]]


def enhance_ordinal_question(message: str, sources: list[dict]) -> str:
    """When the user says 'second paper', inject that paper's title into the prompt."""
    if not _ORDINAL_PAPER.search(message or "") or len(sources) != 1:
        return message
    title = (sources[0].get("title") or "").strip()
    if not title:
        return message
    return f"{message}\n\nAnswer specifically about the paper titled: {title}"


def _question_terms(question: str) -> list[str]:
    terms: list[str] = []
    lowered = (question or "").lower()
    if "ph" in lowered:
        terms.append("ph")
    for token in re.findall(r"[a-z0-9]+", lowered):
        if token in _QUESTION_STOPWORDS or len(token) < 3:
            continue
        terms.append(token)
    return list(dict.fromkeys(terms))


def _extract_section(text: str, section_number: int, section_label: str) -> str:
    pattern = re.compile(
        rf"^##\s*{section_number}\.\s*{re.escape(section_label)}.*?(?=^##\s+\d+\.|\Z)",
        re.MULTILINE | re.DOTALL | re.IGNORECASE,
    )
    match = pattern.search(text)
    return match.group(0).strip() if match else ""


def _matching_lines(text: str, terms: list[str]) -> str:
    if not terms:
        return ""
    lines: list[str] = []
    for line in text.splitlines():
        lowered = line.lower()
        if any(term in lowered for term in terms):
            lines.append(line.rstrip())
    return "\n".join(lines).strip()


def _fit_blocks(blocks: list[str], budget: int) -> str:
    selected: list[str] = []
    used = 0
    for block in blocks:
        block = block.strip()
        if not block:
            continue
        extra = len(block) + (2 if selected else 0)
        if used + extra > budget:
            remaining = budget - used - (2 if selected else 0)
            if remaining > 120:
                selected.append(block[:remaining].rstrip() + "…")
            break
        selected.append(block)
        used += extra
    return "\n\n".join(selected)


def excerpt_summary(text: str, *, question: str = "", n_papers: int = 1) -> str:
    """Prefer identity + key results (+ question matches) instead of blind head truncation."""
    per_paper = max(2500, CHAT_SUMMARY_MAX_CHARS // max(n_papers, 1))
    if len(text) <= per_paper:
        return text

    blocks: list[str] = []
    identity = _extract_section(text, 1, "Paper Identity")
    key_results = _extract_section(text, 5, "Key Results")
    takeaway = _extract_section(text, 10, "Concise Takeaway")
    terms = _question_terms(question)
    matched = _matching_lines(text, terms)

    if identity:
        blocks.append(identity)
    if key_results:
        blocks.append(key_results)
    elif matched:
        blocks.append(matched)
    if takeaway and takeaway not in blocks:
        blocks.append(takeaway)
    if matched and matched not in blocks:
        blocks.append(matched)

    excerpt = _fit_blocks(blocks, per_paper)
    if excerpt:
        return excerpt

    return text[:per_paper] + "\n\n[Additional sections omitted for length.]"


def truncate_summary(text: str, *, n_papers: int = 1, question: str = "") -> str:
    return excerpt_summary(text, question=question, n_papers=n_papers)