from __future__ import annotations

import re
from dataclasses import dataclass

from django.db.models import Q

from control_plane.authoring.models import (
    DocumentAuthority,
    SourceDocumentVersion,
    SourcePassage,
    Work,
    WorkType,
)
from control_plane.authoring.state_management import json_sha256

STOP_WORDS = {
    "and",
    "are",
    "about",
    "after",
    "again",
    "also",
    "before",
    "being",
    "between",
    "but",
    "could",
    "for",
    "from",
    "has",
    "her",
    "him",
    "his",
    "have",
    "into",
    "its",
    "must",
    "not",
    "scene",
    "she",
    "should",
    "that",
    "the",
    "their",
    "them",
    "then",
    "there",
    "they",
    "this",
    "through",
    "what",
    "when",
    "where",
    "which",
    "while",
    "with",
    "would",
    "was",
    "were",
    "write",
}


@dataclass(frozen=True)
class RankedPassage:
    passage: SourcePassage
    score: float
    reason: str


def query_terms(query: str, limit: int = 16) -> list[str]:
    counts: dict[str, int] = {}
    for token in re.findall(r"[a-zA-Z][a-zA-Z0-9']{2,}", query.lower()):
        if token in STOP_WORDS:
            continue
        counts[token] = counts.get(token, 0) + 1
    ordered = sorted(
        counts.items(), key=lambda item: (-item[1], -len(item[0]), item[0])
    )
    return [token for token, _count in ordered[:limit]]


def retrieve_scene_passages(
    *,
    work: Work,
    query: str,
    authorities: list[str] | None = None,
    pinned_document_keys: list[str] | None = None,
    limit: int = 24,
) -> list[RankedPassage]:
    authorities = authorities or [DocumentAuthority.CANON]
    invalid = sorted(set(authorities) - set(DocumentAuthority.values))
    if invalid:
        raise ValueError(f"unsupported document authorities: {', '.join(invalid)}")
    pinned = {value.strip() for value in (pinned_document_keys or []) if value.strip()}
    terms = query_terms(query)
    visible_work_ids = list(
        Work.objects.filter(series=work.series, work_type=WorkType.SERIES_REFERENCE).values_list(
            "id", flat=True
        )
    )
    visible_work_ids.append(work.id)
    base = SourcePassage.objects.select_related(
        "document_version__document"
    ).filter(
        document_version__document__work_id__in=visible_work_ids,
        document_version__authority__in=authorities,
        document_version__superseded_by__isnull=True,
    )

    candidates: dict[object, SourcePassage] = {}
    if terms:
        term_filter = Q()
        for term in terms:
            term_filter |= Q(content__icontains=term)
            term_filter |= Q(document_version__document__title__icontains=term)
        for passage in base.filter(term_filter)[:4000]:
            candidates[passage.id] = passage
    if pinned:
        for passage in base.filter(document_version__document__logical_key__in=pinned)[:2000]:
            candidates[passage.id] = passage
    if not candidates:
        for passage in base[:300]:
            candidates[passage.id] = passage

    ranked: list[RankedPassage] = []
    lowered_query = query.lower()
    for passage in candidates.values():
        document = passage.document_version.document
        haystack = passage.content.lower()
        identity = f"{document.logical_key} {document.title}".lower()
        score = 0.0
        matched = []
        for term in terms:
            occurrences = haystack.count(term)
            if occurrences:
                score += 1.0 + min(occurrences, 4) * 0.5
                matched.append(term)
            if term in identity:
                score += 4.0
        if document.logical_key in pinned:
            score += 100.0
        if passage.document_version.authority == DocumentAuthority.CANON:
            score += 2.0
        if document.title.lower() in lowered_query:
            score += 5.0
        reason = "pinned" if document.logical_key in pinned else "terms: " + ", ".join(matched[:6])
        ranked.append(RankedPassage(passage=passage, score=score, reason=reason.strip()))
    ranked.sort(
        key=lambda item: (
            -item.score,
            item.passage.document_version.document.logical_key,
            item.passage.ordinal,
        )
    )
    requested_limit = max(1, limit)
    if not pinned:
        if len(authorities) == 1:
            return ranked[:requested_limit]
        return _select_across_authorities(ranked, authorities, requested_limit)

    pinned_ranked = [
        item
        for item in ranked
        if item.passage.document_version.document.logical_key in pinned
    ]
    other_ranked = [
        item
        for item in ranked
        if item.passage.document_version.document.logical_key not in pinned
    ]
    if not other_ranked:
        return ranked[:requested_limit]

    pinned_budget = min(len(pinned_ranked), max(1, requested_limit * 2 // 3))
    found_pinned_keys = {
        item.passage.document_version.document.logical_key for item in pinned_ranked
    }
    per_document_limit = max(
        1,
        (pinned_budget + len(found_pinned_keys) - 1) // max(1, len(found_pinned_keys)),
    )
    selected: list[RankedPassage] = []
    pinned_counts: dict[str, int] = {}
    for item in pinned_ranked:
        document_key = item.passage.document_version.document.logical_key
        if len(selected) >= pinned_budget:
            break
        if pinned_counts.get(document_key, 0) >= per_document_limit:
            continue
        selected.append(item)
        pinned_counts[document_key] = pinned_counts.get(document_key, 0) + 1

    selected.extend(
        _select_across_authorities(
            other_ranked,
            authorities,
            requested_limit - len(selected),
        )
    )
    selected_ids = {item.passage.id for item in selected}
    for item in ranked:
        if len(selected) >= requested_limit:
            break
        if item.passage.id not in selected_ids:
            selected.append(item)
            selected_ids.add(item.passage.id)
    return selected


def _select_across_authorities(
    ranked: list[RankedPassage], authorities: list[str], limit: int
) -> list[RankedPassage]:
    authority_groups = {
        authority: [
            item
            for item in ranked
            if item.passage.document_version.authority == authority
        ]
        for authority in authorities
    }
    offsets = {authority: 0 for authority in authorities}
    selected: list[RankedPassage] = []
    while len(selected) < limit:
        added = False
        for authority in authorities:
            offset = offsets[authority]
            group = authority_groups[authority]
            if offset >= len(group):
                continue
            selected.append(group[offset])
            offsets[authority] += 1
            added = True
            if len(selected) >= limit:
                break
        if not added:
            break
    selected_ids = {item.passage.id for item in selected}
    for item in ranked:
        if len(selected) >= limit:
            break
        if item.passage.id not in selected_ids:
            selected.append(item)
            selected_ids.add(item.passage.id)
    return selected


def build_scene_context_pack(
    *,
    work: Work,
    query: str,
    authorities: list[str] | None = None,
    pinned_document_keys: list[str] | None = None,
    governing_document_keys: list[str] | None = None,
    limit: int = 24,
    max_chars: int = 50000,
) -> tuple[dict, list[RankedPassage]]:
    authorities = authorities or [DocumentAuthority.CANON]
    governing_keys = list(
        dict.fromkeys(
            value.strip() for value in (governing_document_keys or []) if value.strip()
        )
    )
    ranked = retrieve_scene_passages(
        work=work,
        query=query,
        authorities=authorities,
        pinned_document_keys=pinned_document_keys,
        limit=limit,
    )
    citations = []
    rendered = []
    used_chars = 0
    kept: list[RankedPassage] = []
    if governing_keys:
        visible_work_ids = list(
            Work.objects.filter(
                series=work.series,
                work_type=WorkType.SERIES_REFERENCE,
            ).values_list("id", flat=True)
        )
        visible_work_ids.append(work.id)
        versions = list(
            SourceDocumentVersion.objects.select_related("document")
            .filter(
                document__work_id__in=visible_work_ids,
                document__logical_key__in=governing_keys,
                authority__in=authorities,
                superseded_by__isnull=True,
            )
            .order_by("document__logical_key")
        )
        versions_by_key: dict[str, list[SourceDocumentVersion]] = {}
        for version in versions:
            versions_by_key.setdefault(version.document.logical_key, []).append(version)
        missing = [key for key in governing_keys if key not in versions_by_key]
        ambiguous = [key for key, values in versions_by_key.items() if len(values) > 1]
        if missing:
            raise ValueError("governing documents not found: " + ", ".join(missing))
        if ambiguous:
            raise ValueError("governing document keys are ambiguous: " + ", ".join(ambiguous))
        for key in governing_keys:
            version = versions_by_key[key][0]
            document = version.document
            label = f"SRC-{len(citations) + 1:02d}"
            end_line = version.content.count("\n") + 1
            block = (
                f"[{label}] authority={version.authority} source={document.logical_key} "
                f"version={version.version} lines=1-{end_line} scope=governing-document\n"
                f"{version.content}"
            )
            if used_chars + len(block) > max_chars:
                raise ValueError("governing documents exceed the context character budget")
            used_chars += len(block)
            rendered.append(block)
            citations.append(
                {
                    "id": label,
                    "kind": "governing_document",
                    "passage_id": None,
                    "document_version_id": str(version.id),
                    "document_key": document.logical_key,
                    "document_title": document.title,
                    "document_version": version.version,
                    "authority": version.authority,
                    "source_path": version.source_path,
                    "start_line": 1,
                    "end_line": end_line,
                    "start_char": 0,
                    "end_char": len(version.content),
                    "sha256": version.source_sha256,
                    "score": None,
                    "reason": "governing document supplied in full",
                }
            )
    governing_set = set(governing_keys)
    for item in ranked:
        passage = item.passage
        version = passage.document_version
        document = version.document
        if document.logical_key in governing_set:
            continue
        excerpt = passage.content[:2500]
        label = f"SRC-{len(citations) + 1:02d}"
        block = (
            f"[{label}] authority={version.authority} source={document.logical_key} "
            f"version={version.version} lines={passage.start_line}-{passage.end_line}\n{excerpt}"
        )
        if rendered and used_chars + len(block) > max_chars:
            continue
        used_chars += len(block)
        kept.append(item)
        citations.append(
            {
                "id": label,
                "passage_id": str(passage.id),
                "document_key": document.logical_key,
                "document_title": document.title,
                "document_version": version.version,
                "authority": version.authority,
                "source_path": version.source_path,
                "start_line": passage.start_line,
                "end_line": passage.end_line,
                "start_char": passage.start_char,
                "end_char": passage.end_char,
                "sha256": passage.sha256,
                "score": item.score,
                "reason": item.reason,
            }
        )
        rendered.append(block)
    pack = {
        "schema_version": 1,
        "work_id": str(work.id),
        "query": query,
        "authorities": authorities,
        "governing_document_keys": governing_keys,
        "citations": citations,
        "rendered_context": (
            "\n\n".join(rendered)
            if rendered
            else "(No matching approved source passages.)"
        ),
    }
    pack["sha256"] = json_sha256(pack)
    return pack, kept