"""Deterministic local document research runtime for Gate 10A."""

from __future__ import annotations

from contextlib import contextmanager
from dataclasses import dataclass, field
from math import log
import os
from pathlib import Path
import re
import socket
from typing import Iterable, Iterator

from .ledger import canonical_json, stable_hash


_SUPPORTED_SUFFIXES = {".md", ".txt"}
_SENTENCE_SPLIT_RE = re.compile(r"(?<=[.!?])\s+")
_TOKEN_RE = re.compile(r"[a-z0-9]+")
_STOPWORDS = {
    "a",
    "an",
    "and",
    "are",
    "as",
    "at",
    "be",
    "by",
    "does",
    "for",
    "from",
    "how",
    "in",
    "is",
    "it",
    "of",
    "on",
    "or",
    "the",
    "to",
    "uses",
    "what",
    "when",
    "where",
    "which",
    "who",
    "with",
}


def _normalize_whitespace(text: str) -> str:
    return " ".join(text.strip().split())


def _normalize_key(text: str) -> str:
    return _normalize_whitespace(text).casefold()


def _tokens(text: str) -> list[str]:
    return [token for token in _TOKEN_RE.findall(text.casefold()) if token not in _STOPWORDS]


@dataclass(frozen=True, slots=True)
class SourceSpan:
    span_id: str
    doc_id: str
    section_id: str
    sentence_index: int
    start_offset: int
    end_offset: int
    text: str
    source_trace_id: str
    source_hash: str

    def to_dict(self) -> dict[str, object]:
        return {
            "span_id": self.span_id,
            "doc_id": self.doc_id,
            "section_id": self.section_id,
            "sentence_index": self.sentence_index,
            "start_offset": self.start_offset,
            "end_offset": self.end_offset,
            "text": self.text,
            "source_trace_id": self.source_trace_id,
            "source_hash": self.source_hash,
        }

    @classmethod
    def from_dict(cls, data: dict[str, object]) -> "SourceSpan":
        return cls(
            span_id=str(data["span_id"]),
            doc_id=str(data["doc_id"]),
            section_id=str(data["section_id"]),
            sentence_index=int(data["sentence_index"]),
            start_offset=int(data["start_offset"]),
            end_offset=int(data["end_offset"]),
            text=str(data["text"]),
            source_trace_id=str(data["source_trace_id"]),
            source_hash=str(data["source_hash"]),
        )


@dataclass(frozen=True, slots=True)
class DocumentClaim:
    claim_id: str
    doc_id: str
    key: str
    value: str
    normalized_key: str
    normalized_value: str
    entity_tokens: tuple[str, ...]
    attribute_tokens: tuple[str, ...]
    key_tokens: tuple[str, ...]
    span: SourceSpan

    def to_dict(self) -> dict[str, object]:
        return {
            "claim_id": self.claim_id,
            "doc_id": self.doc_id,
            "key": self.key,
            "value": self.value,
            "normalized_key": self.normalized_key,
            "normalized_value": self.normalized_value,
            "entity_tokens": list(self.entity_tokens),
            "attribute_tokens": list(self.attribute_tokens),
            "key_tokens": list(self.key_tokens),
            "span": self.span.to_dict(),
        }

    @classmethod
    def from_dict(cls, data: dict[str, object]) -> "DocumentClaim":
        return cls(
            claim_id=str(data["claim_id"]),
            doc_id=str(data["doc_id"]),
            key=str(data["key"]),
            value=str(data["value"]),
            normalized_key=str(data["normalized_key"]),
            normalized_value=str(data["normalized_value"]),
            entity_tokens=tuple(str(item) for item in data.get("entity_tokens", [])),
            attribute_tokens=tuple(str(item) for item in data.get("attribute_tokens", [])),
            key_tokens=tuple(str(item) for item in data.get("key_tokens", [])),
            span=SourceSpan.from_dict(dict(data["span"])),
        )


@dataclass(frozen=True, slots=True)
class ResearchCitation:
    citation_id: str
    claim_id: str
    doc_id: str
    span_id: str
    source_trace_id: str
    source_hash: str
    quote: str

    def to_dict(self) -> dict[str, object]:
        return {
            "citation_id": self.citation_id,
            "claim_id": self.claim_id,
            "doc_id": self.doc_id,
            "span_id": self.span_id,
            "source_trace_id": self.source_trace_id,
            "source_hash": self.source_hash,
            "quote": self.quote,
        }

    @classmethod
    def from_dict(cls, data: dict[str, object]) -> "ResearchCitation":
        return cls(
            citation_id=str(data["citation_id"]),
            claim_id=str(data["claim_id"]),
            doc_id=str(data["doc_id"]),
            span_id=str(data["span_id"]),
            source_trace_id=str(data["source_trace_id"]),
            source_hash=str(data["source_hash"]),
            quote=str(data["quote"]),
        )


@dataclass(frozen=True, slots=True)
class ResearchMemoryRecord:
    memory_id: str
    question: str
    answer_fragment: str
    citation_ids: tuple[str, ...]
    source_trace_ids: tuple[str, ...]

    def to_dict(self) -> dict[str, object]:
        return {
            "memory_id": self.memory_id,
            "question": self.question,
            "answer_fragment": self.answer_fragment,
            "citation_ids": list(self.citation_ids),
            "source_trace_ids": list(self.source_trace_ids),
        }

    @classmethod
    def from_dict(cls, data: dict[str, object]) -> "ResearchMemoryRecord":
        return cls(
            memory_id=str(data["memory_id"]),
            question=str(data["question"]),
            answer_fragment=str(data["answer_fragment"]),
            citation_ids=tuple(str(item) for item in data.get("citation_ids", [])),
            source_trace_ids=tuple(str(item) for item in data.get("source_trace_ids", [])),
        )


@dataclass(frozen=True, slots=True)
class ResearchAnswer:
    task_id: str
    question: str
    status: str
    answer_text: str
    citations: tuple[ResearchCitation, ...]
    memory_records: tuple[ResearchMemoryRecord, ...]
    contradiction_keys: tuple[str, ...] = ()
    unsupported_reason: str = ""
    source_trace_coverage: float = 1.0
    answer_without_source: bool = False

    def to_dict(self) -> dict[str, object]:
        return {
            "task_id": self.task_id,
            "question": self.question,
            "status": self.status,
            "answer_text": self.answer_text,
            "citations": [citation.to_dict() for citation in self.citations],
            "memory_records": [record.to_dict() for record in self.memory_records],
            "contradiction_keys": list(self.contradiction_keys),
            "unsupported_reason": self.unsupported_reason,
            "source_trace_coverage": self.source_trace_coverage,
            "answer_without_source": self.answer_without_source,
        }

    @classmethod
    def from_dict(cls, data: dict[str, object]) -> "ResearchAnswer":
        return cls(
            task_id=str(data["task_id"]),
            question=str(data["question"]),
            status=str(data["status"]),
            answer_text=str(data["answer_text"]),
            citations=tuple(ResearchCitation.from_dict(dict(item)) for item in data.get("citations", [])),
            memory_records=tuple(ResearchMemoryRecord.from_dict(dict(item)) for item in data.get("memory_records", [])),
            contradiction_keys=tuple(str(item) for item in data.get("contradiction_keys", [])),
            unsupported_reason=str(data.get("unsupported_reason", "")),
            source_trace_coverage=float(data.get("source_trace_coverage", 1.0)),
            answer_without_source=bool(data.get("answer_without_source", False)),
        )


@dataclass(frozen=True, slots=True)
class ResearchTask:
    task_id: str
    question: str
    expected_status: str
    expected_fragments: tuple[str, ...] = ()

    def to_dict(self) -> dict[str, object]:
        return {
            "task_id": self.task_id,
            "question": self.question,
            "expected_status": self.expected_status,
            "expected_fragments": list(self.expected_fragments),
        }


@dataclass(frozen=True, slots=True)
class IngestedDocument:
    doc_id: str
    path: str
    source_trace_id: str
    source_hash: str
    spans: tuple[SourceSpan, ...]
    claims: tuple[DocumentClaim, ...]

    def to_dict(self) -> dict[str, object]:
        return {
            "doc_id": self.doc_id,
            "path": self.path,
            "source_trace_id": self.source_trace_id,
            "source_hash": self.source_hash,
            "spans": [span.to_dict() for span in self.spans],
            "claims": [claim.to_dict() for claim in self.claims],
        }


@dataclass(frozen=True, slots=True)
class LocalOnlyNetworkGuardReport:
    enabled: bool
    blocked_targets: tuple[str, ...]
    marker_env_var: str

    def to_dict(self) -> dict[str, object]:
        return {
            "enabled": self.enabled,
            "blocked_targets": list(self.blocked_targets),
            "marker_env_var": self.marker_env_var,
        }


def _split_sentences(text: str) -> list[str]:
    cleaned = _normalize_whitespace(text)
    if not cleaned:
        return []
    return [sentence.strip() for sentence in _SENTENCE_SPLIT_RE.split(cleaned) if sentence.strip()]


def _claim_from_span(doc_id: str, span: SourceSpan) -> DocumentClaim | None:
    if ":" not in span.text:
        return None
    key, raw_value = span.text.split(":", 1)
    key = _normalize_whitespace(key)
    value = _normalize_whitespace(raw_value.strip().rstrip("."))
    if not key or not value:
        return None
    key_tokens = tuple(_tokens(key))
    if not key_tokens:
        return None
    entity_tokens = key_tokens[:2] if len(key_tokens) >= 2 else key_tokens
    attribute_tokens = key_tokens[2:] if len(key_tokens) > 2 else key_tokens
    claim_id = stable_hash({"doc_id": doc_id, "span_id": span.span_id, "key": key, "value": value})[:20]
    return DocumentClaim(
        claim_id=claim_id,
        doc_id=doc_id,
        key=key,
        value=value,
        normalized_key=_normalize_key(key),
        normalized_value=_normalize_key(value),
        entity_tokens=entity_tokens,
        attribute_tokens=attribute_tokens,
        key_tokens=key_tokens,
        span=span,
    )


def ingest_document(path: str | Path, *, root: str | Path | None = None) -> IngestedDocument:
    source_path = Path(path)
    if source_path.suffix.casefold() not in _SUPPORTED_SUFFIXES:
        raise ValueError(f"unsupported document suffix: {source_path.suffix}")
    doc_root = Path(root) if root is not None else source_path.parent
    relative_path = source_path.resolve().relative_to(doc_root.resolve())
    doc_id = relative_path.as_posix()
    text = source_path.read_text(encoding="utf-8")
    source_hash = stable_hash({"doc_id": doc_id, "text": text})
    source_trace_id = stable_hash({"doc_id": doc_id, "source_hash": source_hash})[:20]
    spans: list[SourceSpan] = []
    claims: list[DocumentClaim] = []
    section_id = "root"
    cursor = 0
    sentence_index = 0
    for raw_line in text.splitlines():
        stripped = raw_line.strip()
        if stripped.startswith("#"):
            section_id = _normalize_whitespace(stripped.lstrip("#")) or "root"
        for sentence in _split_sentences(raw_line):
            start_offset = text.find(sentence, cursor)
            if start_offset < 0:
                start_offset = cursor
            end_offset = start_offset + len(sentence)
            cursor = end_offset
            span_id = stable_hash({"doc_id": doc_id, "sentence_index": sentence_index, "text": sentence})[:20]
            span = SourceSpan(
                span_id=span_id,
                doc_id=doc_id,
                section_id=section_id,
                sentence_index=sentence_index,
                start_offset=start_offset,
                end_offset=end_offset,
                text=sentence,
                source_trace_id=source_trace_id,
                source_hash=source_hash,
            )
            spans.append(span)
            claim = _claim_from_span(doc_id, span)
            if claim is not None:
                claims.append(claim)
            sentence_index += 1
    return IngestedDocument(
        doc_id=doc_id,
        path=str(source_path.resolve()),
        source_trace_id=source_trace_id,
        source_hash=source_hash,
        spans=tuple(spans),
        claims=tuple(claims),
    )


def ingest_documents(root: str | Path) -> tuple[IngestedDocument, ...]:
    source_root = Path(root).resolve()
    documents = []
    for path in sorted(source_root.rglob("*")):
        if path.is_file() and path.suffix.casefold() in _SUPPORTED_SUFFIXES:
            documents.append(ingest_document(path, root=source_root))
    return tuple(documents)


def load_local_documents(
    *,
    doc_dir: str | Path | None = None,
    doc_paths: Iterable[str | Path] = (),
) -> tuple[IngestedDocument, ...]:
    provided_paths = tuple(Path(path).resolve() for path in doc_paths)
    if doc_dir is not None and provided_paths:
        raise ValueError("provide either --doc-dir or --doc-path, not both")
    if doc_dir is None and not provided_paths:
        raise ValueError("local document research requires --doc-dir or --doc-path")

    if doc_dir is not None:
        source_root = Path(doc_dir).resolve()
        if not source_root.is_dir():
            raise ValueError(f"document directory does not exist: {source_root}")
        documents = ingest_documents(source_root)
    else:
        invalid_paths = [path for path in provided_paths if not path.is_file()]
        if invalid_paths:
            raise ValueError(f"document path does not exist: {invalid_paths[0]}")
        unsupported_paths = [path for path in provided_paths if path.suffix.casefold() not in _SUPPORTED_SUFFIXES]
        if unsupported_paths:
            raise ValueError(f"unsupported document suffix: {unsupported_paths[0].suffix}")
        source_root = Path(os.path.commonpath([str(path.parent) for path in provided_paths]))
        documents = tuple(ingest_document(path, root=source_root) for path in sorted(provided_paths))

    if not documents:
        raise ValueError("no supported local documents found")
    return documents


@contextmanager
def local_only_network_guard() -> Iterator[LocalOnlyNetworkGuardReport]:
    marker_env_var = "RA_CPU_DOC_RESEARCH_NO_NETWORK"
    blocked_targets = (
        "socket.socket",
        "socket.create_connection",
        "socket.getaddrinfo",
        "socket.gethostbyname",
        "socket.gethostbyname_ex",
    )

    def _blocked_network_call(*args: object, **kwargs: object) -> object:
        raise RuntimeError("network access is disabled for local document research")

    previous_marker = os.environ.get(marker_env_var)
    originals = {
        "socket": socket.socket,
        "create_connection": socket.create_connection,
        "getaddrinfo": socket.getaddrinfo,
        "gethostbyname": socket.gethostbyname,
        "gethostbyname_ex": socket.gethostbyname_ex,
    }
    os.environ[marker_env_var] = "1"
    socket.socket = _blocked_network_call  # type: ignore[assignment]
    socket.create_connection = _blocked_network_call  # type: ignore[assignment]
    socket.getaddrinfo = _blocked_network_call  # type: ignore[assignment]
    socket.gethostbyname = _blocked_network_call  # type: ignore[assignment]
    socket.gethostbyname_ex = _blocked_network_call  # type: ignore[assignment]
    try:
        yield LocalOnlyNetworkGuardReport(
            enabled=True,
            blocked_targets=blocked_targets,
            marker_env_var=marker_env_var,
        )
    finally:
        socket.socket = originals["socket"]
        socket.create_connection = originals["create_connection"]
        socket.getaddrinfo = originals["getaddrinfo"]
        socket.gethostbyname = originals["gethostbyname"]
        socket.gethostbyname_ex = originals["gethostbyname_ex"]
        if previous_marker is None:
            os.environ.pop(marker_env_var, None)
        else:
            os.environ[marker_env_var] = previous_marker


class LocalDocumentResearchRuntime:
    """Deterministic source-backed answers over local documents only."""

    def __init__(self, documents: Iterable[IngestedDocument]) -> None:
        self.documents = tuple(documents)
        self.claims = tuple(claim for document in self.documents for claim in document.claims)
        self.claims_by_key: dict[str, list[DocumentClaim]] = {}
        self.df_by_token: dict[str, int] = {}
        for claim in self.claims:
            self.claims_by_key.setdefault(claim.normalized_key, []).append(claim)
            for token in set(claim.key_tokens):
                self.df_by_token[token] = self.df_by_token.get(token, 0) + 1
        self.total_claims = max(len(self.claims), 1)

    @classmethod
    def from_directory(cls, root: str | Path) -> "LocalDocumentResearchRuntime":
        return cls(ingest_documents(root))

    def _token_weight(self, token: str) -> float:
        document_frequency = self.df_by_token.get(token, 0)
        return 1.0 + log((1.0 + self.total_claims) / (1.0 + document_frequency))

    def _score_claim(self, claim: DocumentClaim, query_tokens: set[str]) -> float:
        overlap = query_tokens.intersection(claim.key_tokens)
        if not overlap:
            return 0.0
        if not query_tokens.intersection(claim.attribute_tokens):
            return 0.0
        overlap_weight = sum(self._token_weight(token) for token in overlap)
        query_weight = sum(self._token_weight(token) for token in query_tokens)
        if query_weight <= 0:
            return 0.0
        return overlap_weight / query_weight

    def _summary_candidates(self, query_tokens: set[str]) -> list[DocumentClaim]:
        candidates = [
            claim
            for claim in self.claims
            if query_tokens.intersection(claim.entity_tokens) == set(claim.entity_tokens)
        ]
        seen_keys: set[str] = set()
        selected: list[DocumentClaim] = []
        for claim in candidates:
            if claim.normalized_key in seen_keys:
                continue
            conflicting = {item.normalized_value for item in self.claims_by_key.get(claim.normalized_key, [])}
            if len(conflicting) > 1:
                continue
            seen_keys.add(claim.normalized_key)
            selected.append(claim)
        return selected[:3]

    def _citations_from_claims(self, claims: Iterable[DocumentClaim]) -> tuple[ResearchCitation, ...]:
        citations = []
        for claim in claims:
            citations.append(
                ResearchCitation(
                    citation_id=stable_hash({"claim_id": claim.claim_id, "span_id": claim.span.span_id})[:20],
                    claim_id=claim.claim_id,
                    doc_id=claim.doc_id,
                    span_id=claim.span.span_id,
                    source_trace_id=claim.span.source_trace_id,
                    source_hash=claim.span.source_hash,
                    quote=claim.span.text,
                )
            )
        return tuple(citations)

    def _memory_records(self, question: str, answer_fragments: Iterable[str], citations: tuple[ResearchCitation, ...]) -> tuple[ResearchMemoryRecord, ...]:
        source_trace_ids = tuple(dict.fromkeys(citation.source_trace_id for citation in citations))
        citation_ids = tuple(citation.citation_id for citation in citations)
        records = []
        for fragment in answer_fragments:
            fragment_text = _normalize_whitespace(fragment)
            if not fragment_text:
                continue
            records.append(
                ResearchMemoryRecord(
                    memory_id=stable_hash({"question": question, "fragment": fragment_text, "citations": citation_ids})[:20],
                    question=question,
                    answer_fragment=fragment_text,
                    citation_ids=citation_ids,
                    source_trace_ids=source_trace_ids,
                )
            )
        return tuple(records)

    def answer_question(self, *, task_id: str, question: str) -> ResearchAnswer:
        query_tokens = set(_tokens(question))
        lowered = question.casefold().strip()
        if lowered.startswith("summarize ") or lowered.startswith("summary "):
            summary_claims = self._summary_candidates(query_tokens)
            if not summary_claims:
                return ResearchAnswer(
                    task_id=task_id,
                    question=question,
                    status="unsupported",
                    answer_text="Unsupported claim: no local source covers this summary request.",
                    citations=(),
                    memory_records=(),
                    unsupported_reason="no_summary_support",
                    source_trace_coverage=1.0,
                    answer_without_source=False,
                )
            citations = self._citations_from_claims(summary_claims)
            answer_lines = [f"{claim.key}: {claim.value}." for claim in summary_claims]
            answer_text = " ".join(answer_lines)
            return ResearchAnswer(
                task_id=task_id,
                question=question,
                status="supported",
                answer_text=answer_text,
                citations=citations,
                memory_records=self._memory_records(question, answer_lines, citations),
                source_trace_coverage=1.0 if citations else 0.0,
                answer_without_source=not citations,
            )

        scored_claims = sorted(
            ((self._score_claim(claim, query_tokens), claim) for claim in self.claims),
            key=lambda item: (item[0], item[1].claim_id),
            reverse=True,
        )
        top_score, top_claim = scored_claims[0] if scored_claims else (0.0, None)
        if top_claim is None or top_score < 0.55:
            return ResearchAnswer(
                task_id=task_id,
                question=question,
                status="unsupported",
                answer_text="Unsupported claim: no local source supports this answer.",
                citations=(),
                memory_records=(),
                unsupported_reason="no_supporting_claim",
                source_trace_coverage=1.0,
                answer_without_source=False,
            )

        sibling_claims = self.claims_by_key.get(top_claim.normalized_key, [top_claim])
        distinct_values = {claim.normalized_value for claim in sibling_claims}
        if len(distinct_values) > 1:
            citations = self._citations_from_claims(sibling_claims)
            return ResearchAnswer(
                task_id=task_id,
                question=question,
                status="contradicted",
                answer_text="Unsupported claim: local sources disagree on this fact.",
                citations=citations,
                memory_records=self._memory_records(question, [top_claim.key], citations),
                contradiction_keys=(top_claim.key,),
                unsupported_reason="contradictory_sources",
                source_trace_coverage=1.0 if citations else 0.0,
                answer_without_source=False,
            )

        citations = self._citations_from_claims(sibling_claims)
        answer_fragment = f"{top_claim.key}: {top_claim.value}."
        return ResearchAnswer(
            task_id=task_id,
            question=question,
            status="supported",
            answer_text=answer_fragment,
            citations=citations,
            memory_records=self._memory_records(question, [answer_fragment], citations),
            source_trace_coverage=1.0 if citations else 0.0,
            answer_without_source=not citations,
        )


def run_local_document_research_task(
    *,
    task_id: str,
    question: str,
    doc_dir: str | Path | None = None,
    doc_paths: Iterable[str | Path] = (),
) -> tuple[ResearchAnswer, tuple[IngestedDocument, ...], LocalOnlyNetworkGuardReport]:
    documents = load_local_documents(doc_dir=doc_dir, doc_paths=doc_paths)
    runtime = LocalDocumentResearchRuntime(documents)
    with local_only_network_guard() as network_guard:
        answer = runtime.answer_question(task_id=task_id, question=question)
    return answer, documents, network_guard


def write_jsonl(path: str | Path, rows: Iterable[dict[str, object]]) -> None:
    target = Path(path)
    target.write_text("".join(f"{canonical_json(row)}\n" for row in rows), encoding="utf-8")
