"""Bounded document extraction and portable provenance. No application imports or I/O at import."""
from __future__ import annotations

import codecs
import copy
import csv
import hashlib
import io
import json
import os
from pathlib import Path
import re
import subprocess
import sys
import tempfile
from dataclasses import asdict, dataclass, field
from datetime import datetime, timezone
from uuid import uuid4
import zipfile


class DocumentError(ValueError):
    def __init__(self, code):
        self.code = code
        super().__init__(code)


def now():
    return datetime.now(timezone.utc).isoformat()


def digest(value):
    return hashlib.sha256(value if isinstance(value, bytes) else value.encode("utf-8")).hexdigest()


def packed(value):
    return json.dumps(value, ensure_ascii=False, sort_keys=True, allow_nan=False, separators=(",", ":"))


def unique_json_object(pairs):
    """Used at the JSON ingress, before duplicate IDs can be silently overwritten."""
    result = {}
    for key, value in pairs:
        if key in result and result[key] != value:
            raise DocumentError("conflicting_duplicate_json_key")
        result[key] = value
    return result


@dataclass(frozen=True)
class Limits:
    file_bytes: int = 10 * 1024 * 1024
    segment_chars: int = 2000
    extract_chars: int = 100000
    extract_segments: int = 200
    pdf_pages: int = 50
    store_bytes: int = 8 * 1024 * 1024
    sources: int = 16
    parser_seconds: float = 30
    retrieval_seconds: float = 8
    retrieval_windows: int = 8
    parser_memory_bytes: int = 768 * 1024 * 1024
    expanded_bytes: int = 40 * 1024 * 1024
    calls: int = 4
    interpretation_seconds: float = 90
    durable_bytes: int = 1024 * 1024
    source_bytes: int = 64 * 1024
    selections: int = 32
    quote_chars: int = 1200
    context_chars: int = 600
    session_bytes: int = 4 * 1024 * 1024
    session_file_bytes: int = 10 * 1024 * 1024
    events: int = 256
    operations: int = 256

    @classmethod
    def from_env(cls):
        defaults = cls()
        values = {}
        for key in defaults.__dataclass_fields__:
            raw = os.getenv("SOFIEL_DOC_" + key.upper())
            if raw is not None:
                try:
                    value = type(getattr(defaults, key))(raw)
                    if not 0 < value <= getattr(defaults, key):
                        raise ValueError()
                    values[key] = value
                except (ValueError, OverflowError):
                    raise DocumentError("invalid_limit_" + key)
        return cls(**values)


@dataclass
class ExtractionResult:
    status: str
    operation_id: str
    error_code: str | None = None
    source: dict = field(default_factory=dict)
    segments: dict = field(default_factory=dict)
    warnings: list = field(default_factory=list)


def empty_store(instance_id=None):
    return {"schema_version": 1, "instance_id": instance_id or "doc_" + uuid4().hex,
            "revision": 0, "sources": {}, "segments": {}, "knowledge": {}, "operations": {}}


def _require(condition, code):
    if not condition:
        raise DocumentError(code)


def _text(value, limit=4000, empty=False):
    return isinstance(value, str) and len(value) <= limit and (empty or bool(value.strip()))


def _date(value):
    try:
        return isinstance(value, str) and datetime.fromisoformat(value).tzinfo is not None
    except ValueError:
        return False


def validate_store(value, limits=None):
    try:
        if isinstance(value, dict) and value.get("schema_version") == 2:
            from sofiel_selective import validate_v2
            return validate_v2(value, limits)
        return _validate_store(value, limits)
    except DocumentError:
        raise
    except (TypeError, ValueError, KeyError, AttributeError, OverflowError):
        raise DocumentError("invalid_document_structure") from None


def _validate_store(value, limits=None):
    """Validate, never repair by dropping records. Returns an isolated JSON copy."""
    limits = limits or Limits()
    _require(isinstance(value, dict), "documents_not_object")
    _require(type(value.get("schema_version")) is int and value["schema_version"] == 1,
             "unsupported_documents_version")
    _require(_text(value.get("instance_id"), 100), "invalid_instance_id")
    _require(type(value.get("revision")) is int and value["revision"] >= 0, "invalid_revision")
    for name in ("sources", "segments", "knowledge", "operations"):
        _require(isinstance(value.get(name), dict), "invalid_" + name)
        for key, row in value[name].items():
            _require(_text(key, 150) and isinstance(row, dict), "invalid_record")
            _require(row.get("instance_id") == value["instance_id"], "foreign_instance")
    _require(len(value["sources"]) <= limits.sources, "source_limit")
    for key, source in value["sources"].items():
        _require(source.get("source_id") == key and re.fullmatch(r"[0-9a-f]{64}", source.get("sha256", "")), "invalid_source")
        _require(_text(source.get("name"), 255) and source.get("format") in {"txt", "md", "csv", "docx", "pdf"}, "invalid_source_format")
        _require(_date(source.get("created_at")) and type(source.get("active")) is bool, "invalid_source_metadata")
        _require(source.get("normalization_version") == "unicode-v1" and isinstance(source.get("coverage"), dict), "invalid_coverage")
        _require(isinstance(source.get("declared_authors"), list) and all(_text(x, 200) for x in source["declared_authors"]), "invalid_authors")
        _require(source.get("participant") is None or _text(source["participant"], 200), "invalid_participant")
        cov = source["coverage"]
        for bucket in ("visited", "empty", "failed"):
            _require(isinstance(cov.get(bucket), list) and all(type(x) is int and x >= 0 for x in cov[bucket]), "invalid_coverage_units")
        _require(type(cov.get("total_units")) is int and cov["total_units"] >= 0, "invalid_unit_count")
        _require(type(cov.get("pending_units")) is int and cov["pending_units"] >= 0, "invalid_pending_count")
        cursor = cov.get("cursor")
        _require(cursor is None or (isinstance(cursor, dict) and all(type(cursor.get(x)) is int and cursor[x] >= 0 for x in ("unit", "offset"))), "invalid_cursor")
    ranges = {}
    for key, segment in value["segments"].items():
        _require(segment.get("segment_id") == key and segment.get("source_id") in value["sources"], "dangling_segment")
        text = segment.get("text")
        _require(_text(text, limits.segment_chars, empty=True), "invalid_segment_text")
        start, end = segment.get("char_start"), segment.get("char_end")
        _require(type(start) is int and type(end) is int and 0 <= start <= end and end-start == len(text), "invalid_offsets")
        _require(segment.get("sha256") == digest(text), "invalid_segment_hash")
        _require(isinstance(segment.get("locator"), dict), "invalid_locator")
        unit = segment["locator"].get("unit")
        _require(type(unit) is int and unit >= 0, "invalid_unit")
        ranges.setdefault((segment["source_id"], unit), []).append((start, end))
    for spans in ranges.values():
        spans.sort()
        _require(all(left[1] <= right[0] for left, right in zip(spans, spans[1:])), "overlapping_segments")
    for key, knowledge in value["knowledge"].items():
        _require(knowledge.get("knowledge_id") == key and knowledge.get("source_id") in value["sources"], "invalid_knowledge")
        _require(knowledge.get("operation_id") in value["operations"], "missing_operation")
        _require(_text(knowledge.get("statement"), 2000), "invalid_statement")
        _require(knowledge.get("claim_kind") in {"source_claim", "sofiel_interpretation"}, "invalid_claim_kind")
        _require(knowledge.get("state") in {"candidate", "incorporated", "rejected"}, "invalid_knowledge_state")
        attribution = knowledge.get("attribution", {})
        _require(isinstance(attribution, dict) and attribution.get("role") in {"author", "narrator", "character", "quotation", "unknown"}, "invalid_attribution")
        _require(_text(attribution.get("label", ""), 200, True), "invalid_attribution_label")
        _require(isinstance(knowledge.get("relevance_reasons"), list) and all(_text(r, 1000) for r in knowledge["relevance_reasons"]), "invalid_relevance")
        _require(_date(knowledge.get("created_at")), "invalid_knowledge_date")
        _require(_date(knowledge.get("incorporated_at")) if knowledge["state"] == "incorporated" else knowledge.get("incorporated_at") is None, "invalid_commit_date")
        refs = knowledge.get("support_refs")
        _require(isinstance(refs, list) and 0 < len(refs) <= 8, "missing_support")
        for ref in refs:
            _require(isinstance(ref, dict), "invalid_support")
            seg = value["segments"].get(ref.get("segment_id"), {})
            _require(seg.get("source_id") == knowledge["source_id"], "foreign_support")
            quote = ref.get("quote")
            _require(_text(quote, limits.segment_chars) and quote in seg.get("text", ""), "unsupported_quote")
        _require(isinstance(knowledge.get("memory_refs"), list), "invalid_memory_refs")
        for ref in knowledge["memory_refs"]:
            _require(isinstance(ref, dict) and _text(ref.get("memory_id"), 100) and ref.get("status") in {"resolved", "orphaned"}, "invalid_memory_reference")
        _require(isinstance(knowledge.get("relations"), list), "invalid_relations")
        for relation in knowledge["relations"]:
            _require(isinstance(relation, dict) and relation.get("kind") in {"contradiction", "duplication", "derivation"} and relation.get("target_id") in value["knowledge"], "invalid_relation")
    for key, op in value["operations"].items():
        _require(op.get("operation_id") == key and _date(op.get("created_at")), "invalid_operation")
        for field_name in ("submitted_segment_ids", "analyzed_segment_ids"):
            _require(isinstance(op.get(field_name), list) and all(s in value["segments"] for s in op[field_name]), "invalid_analyzed_segment")
        _require(set(op["analyzed_segment_ids"]).issubset(op["submitted_segment_ids"]), "unsubmitted_analysis")
    try:
        serialized = packed(value)
    except (ValueError, TypeError):
        raise DocumentError("not_finite_json")
    # Match the real, indented memory export, including nesting below documents.
    durable_bytes = len(json.dumps({"documents": value}, ensure_ascii=False, allow_nan=False, indent=2).encode("utf-8"))
    _require(durable_bytes <= limits.store_bytes, "document_store_limit")
    return json.loads(serialized)


def memory_index(chats):
    """IDs derive only from actually persisted fields; ambiguous duplicates are omitted."""
    indexed, counts = {}, {}
    for chat in chats:
        if not isinstance(chat, dict):
            continue
        fields = {k: chat.get(k, "") for k in ("ts", "user", "sofiel")}
        key = "mem_" + digest(packed(fields))
        counts[key] = counts.get(key, 0) + 1
        indexed[key] = fields
    return {k: v for k, v in indexed.items() if counts[k] == 1}


def reconcile_memory_refs(store, chats):
    index = memory_index(chats)
    for item in store.get("insights", store.get("knowledge", {})).values():
        for ref in item["memory_refs"]:
            ref["status"] = "resolved" if ref["memory_id"] in index else "orphaned"


def _decode(path, encoding):
    # Explicit encoding or BOM/UTF-8 only. No lossy fallback or universal-newline rewriting.
    with open(path, "rb") as stream:
        prefix = stream.read(4)
    encoding = encoding or ("utf-16" if prefix.startswith((codecs.BOM_UTF16_LE, codecs.BOM_UTF16_BE)) else "utf-8-sig")
    decoder = codecs.getincrementaldecoder(encoding)(errors="strict")
    chunks = []
    with open(path, "rb") as stream:
        while chunk := stream.read(65536):
            chunks.append(decoder.decode(chunk))
    chunks.append(decoder.decode(b"", final=True))
    return "".join(chunks), encoding


def _units(path, fmt, limits, encoding=None, delimiter=None):
    warnings = []
    if fmt in {"txt", "md", "csv"}:
        text, encoding = _decode(path, encoding)
        if fmt != "csv":
            return [(text, {"kind": "text"})], warnings, encoding
        if delimiter is None:
            try:
                delimiter = csv.Sniffer().sniff(text[:8192], delimiters=",;\t|").delimiter
            except csv.Error:
                raise DocumentError("csv_delimiter_required")
        if delimiter not in {",", ";", "\t", "|"}:
            raise DocumentError("invalid_csv_delimiter")
        # Row canonical text is a JSON list, retaining cells, embedded newlines and header row.
        rows = [(packed(row), {"kind": "csv_row", "row": i+1, "columns": len(row), "delimiter": delimiter})
                for i, row in enumerate(csv.reader(io.StringIO(text, newline=""), delimiter=delimiter, strict=True))]
        return rows, ["CSV: texto canónico por fila como lista JSON; primera fila conservada."], encoding
    if fmt == "docx":
        with zipfile.ZipFile(path) as archive:
            infos = archive.infolist()
            _require(len(infos) <= 10000 and sum(x.file_size for x in infos) <= limits.expanded_bytes, "docx_expansion_limit")
            _require(all(not x.flag_bits & 1 for x in infos), "encrypted_docx")
        from docx import Document
        from docx.table import Table
        document = Document(path)
        units = []
        for i, block in enumerate(document.iter_inner_content()):
            if isinstance(block, Table):
                seen = set()
                for row_num, row in enumerate(block.rows, 1):
                    for col_num, cell in enumerate(row.cells, 1):
                        if cell._tc in seen:
                            continue
                        seen.add(cell._tc)
                        units.append((cell.text, {"kind": "docx_cell", "block": i+1, "row": row_num, "column": col_num}))
            else:
                units.append((block.text, {"kind": "docx_paragraph", "block": i+1}))
        return units, ["DOCX: cuerpo, párrafos y celdas de tablas; no notas, encabezados, cuadros de texto ni tablas anidadas."], "xml-unicode"
    from pypdf import PdfReader
    reader = PdfReader(path)
    _require(not reader.is_encrypted, "encrypted_pdf")
    # Lazy page access: text extraction only occurs for the visited window.
    return reader.pages, ["PDF: texto extraído por página; la estructura/tablas pueden perderse. Sin OCR."], "pdf-extracted"


def extract_file(path, *, operation_id=None, instance_id=None, cursor=None, expected_hash=None,
                 encoding=None, delimiter=None, limits=None):
    limits = limits or Limits()
    operation_id = operation_id or "op_" + uuid4().hex
    instance_id = instance_id or "doc_" + uuid4().hex
    result = ExtractionResult("failed", operation_id)
    try:
        path = Path(path)
        _require(path.is_file() and path.stat().st_size <= limits.file_bytes, "file_size_or_missing")
        fmt = path.suffix.lower().lstrip(".")
        _require(fmt in {"txt", "md", "csv", "docx", "pdf"}, "unsupported_format")
        sha = hashlib.sha256()
        with path.open("rb") as stream:
            total = 0
            while chunk := stream.read(65536):
                total += len(chunk)
                _require(total <= limits.file_bytes, "file_size_limit")
                sha.update(chunk)
        file_hash = sha.hexdigest()
        _require(expected_hash is None or file_hash == expected_hash, "source_hash_mismatch")
        source_id = "src_" + file_hash
        units, warnings, chosen_encoding = _units(path, fmt, limits, encoding, delimiter)
        start = cursor or {"unit": 0, "offset": 0}
        _require(isinstance(start, dict) and all(type(start.get(k)) is int and start[k] >= 0 for k in ("unit", "offset")), "invalid_cursor")
        _require(start["unit"] <= len(units), "invalid_cursor")
        coverage = {"total_units": len(units), "visited": [], "empty": [], "failed": [], "pending_units": 0, "cursor": None}
        retained = 0
        for index in range(start["unit"], len(units)):
            if fmt == "pdf" and len(coverage["visited"]) >= limits.pdf_pages:
                coverage["cursor"] = {"unit": index, "offset": 0}
                break
            coverage["visited"].append(index)
            try:
                if fmt == "pdf":
                    text, locator = units[index].extract_text() or "", {"kind": "pdf_page", "page": index+1}
                else:
                    text, locator = units[index]
            except Exception:
                coverage["failed"].append(index)
                continue
            locator = {**locator, "unit": index}
            if not text:
                coverage["empty"].append(index)
                continue
            offset = start["offset"] if index == start["unit"] else 0
            _require(offset <= len(text), "invalid_cursor_offset")
            while offset < len(text):
                available = min(limits.segment_chars, limits.extract_chars-retained)
                if available <= 0 or len(result.segments) >= limits.extract_segments:
                    coverage["cursor"] = {"unit": index, "offset": offset}
                    break
                piece = text[offset:offset+available]
                segment_id = "seg_" + digest(f"{source_id}:{index}:{offset}:{digest(piece)}")
                result.segments[segment_id] = {"segment_id": segment_id, "instance_id": instance_id,
                    "source_id": source_id, "locator": locator, "char_start": offset,
                    "char_end": offset+len(piece), "text": piece, "sha256": digest(piece)}
                offset += len(piece)
                retained += len(piece)
            if coverage["cursor"]:
                break
        if coverage["cursor"]:
            coverage["pending_units"] = len(units)-coverage["cursor"]["unit"]
        if coverage["empty"]:
            warnings.append("Unidades sin texto disponible; en PDF podrían requerir OCR.")
        # The UI may replace a temporary upload while a parser is working. Never
        # attach extracted text to a hash of different bytes.
        check = hashlib.sha256()
        checked_bytes = 0
        with path.open("rb") as stream:
            while chunk := stream.read(65536):
                checked_bytes += len(chunk)
                _require(checked_bytes <= limits.file_bytes, "file_size_limit")
                check.update(chunk)
        _require(check.hexdigest() == file_hash, "source_changed_during_extraction")
        result.status = "partial" if coverage["cursor"] or coverage["failed"] else "complete"
        result.source = {"source_id": source_id, "instance_id": instance_id, "sha256": file_hash,
            "name": path.name[:255], "format": fmt, "created_at": now(), "participant": None,
            "declared_authors": [], "active": True, "availability": "retained_text",
            "normalization_version": "unicode-v1", "encoding": chosen_encoding, "coverage": coverage}
        result.warnings = warnings
    except DocumentError as exc:
        result.error_code = exc.code
    except (UnicodeError, LookupError):
        result.error_code = "encoding_required_or_invalid"
    except ImportError:
        result.error_code = "parser_dependency_missing"
    except Exception:
        result.error_code = "parse_failed"
    return result


def extract_isolated(path, *, limits=None, **kwargs):
    """Turn parser infrastructure failures into results the UI can apply."""
    operation_id = kwargs.get("operation_id") or "op_" + uuid4().hex
    kwargs["operation_id"] = operation_id
    try:
        return _extract_isolated(path, limits=limits, **kwargs)
    except OSError:
        return ExtractionResult("failed", operation_id, "parser_io_failed")
    except (ValueError, TypeError, KeyError):
        return ExtractionResult("failed", operation_id, "parser_result_invalid")


def _extract_isolated(path, *, limits=None, **kwargs):
    """One short-lived child, no app/Torch import. communicate timeout always reaps it."""
    limits = limits or Limits()
    request = {"path": str(Path(path).resolve()), "limits": asdict(limits), "kwargs": kwargs}
    with tempfile.TemporaryDirectory(prefix="sofiel-parser-") as directory:
        request_path, response_path = Path(directory)/"request.json", Path(directory)/"response.json"
        request_path.write_text(packed(request), encoding="utf-8")
        process = subprocess.Popen([sys.executable, str(Path(__file__).resolve()), str(request_path), str(response_path)],
            stdout=subprocess.DEVNULL, stderr=subprocess.DEVNULL,
            creationflags=getattr(subprocess, "CREATE_NO_WINDOW", 0))
        try:
            process.communicate(timeout=limits.parser_seconds)
        except subprocess.TimeoutExpired:
            process.kill()
            process.communicate()
            return ExtractionResult("failed", kwargs.get("operation_id", ""), "parser_timeout")
        except OSError:
            process.kill()
            process.communicate()
            raise
        if process.returncode != 0 or not response_path.is_file():
            return ExtractionResult("failed", kwargs.get("operation_id", ""), "parser_process_failed")
        result = ExtractionResult(**json.loads(response_path.read_text(encoding="utf-8")))
        if sys.platform != "linux":
            result.warnings.append("Límite duro de memoria del parser no disponible en este sistema; cuotas y timeout aplicados.")
        return result


def retrieve(store, query, *, max_chars=12000, limit=6, atoms=None):
    """Lexical retrieval across ALL retained segments, with separate attributed knowledge."""
    if not store:
        return {}
    if store.get("schema_version") == 2:
        from sofiel_selective import retrieve_v2
        return retrieve_v2(store, query, max_chars=max_chars, limit=limit, atoms=atoms)
    terms = set(re.findall(r"\w{3,}", str(query).casefold()))
    derived = []
    supported_ids = set()
    for atom in atoms or []:
        provenance = getattr(atom, "evidence_provenance", [])
        if provenance and terms & set(re.findall(r"\w{3,}", atom.concepto.casefold())):
            refs = {sid for event in provenance for sid in event.get("segment_ids", [])}
            live = {sid for sid in refs if sid in store["segments"] and store["sources"][store["segments"][sid]["source_id"]]["active"]}
            if live:
                supported_ids.update(live)
                derived.append({"concept": atom.concepto[:400], "provenance": copy.deepcopy(provenance[-2:]),
                                "support_verified": False})
            if len(derived) >= 3:
                break
    ranked = []
    for segment in store["segments"].values():
        if not store["sources"][segment["source_id"]]["active"]:
            continue
        related = [k for k in store["knowledge"].values() if k["state"] == "incorporated" and any(r["segment_id"] == segment["segment_id"] for r in k["support_refs"])]
        corpus = segment["text"] + " " + " ".join(k["statement"] for k in related)
        score = len(terms & set(re.findall(r"\w{3,}", corpus.casefold())))
        if segment["segment_id"] in supported_ids:
            score += 1
        if score:
            ranked.append((score, segment["segment_id"], segment, related))
    result = {"revision": store["revision"], "instance_id": store["instance_id"], "retrieval": "lexical", "passages": [], "knowledge": []}
    for _, _, segment, knowledge in sorted(ranked, key=lambda row: (-row[0], row[1])):
        source = store["sources"][segment["source_id"]]
        candidate = copy.deepcopy(result)
        candidate["passages"].append({**copy.deepcopy(segment), "source_name": source["name"], "declared_authors": source["declared_authors"]})
        known = {k["knowledge_id"] for k in candidate["knowledge"]}
        candidate["knowledge"].extend(copy.deepcopy(k) for k in knowledge if k["knowledge_id"] not in known)
        if len(packed(candidate)) <= max_chars:
            result = candidate
        if len(result["passages"]) >= limit:
            break
    if derived and len(packed({**result, "derived_atoms": derived})) <= max_chars:
        result["derived_atoms"] = derived
    return result if result["passages"] else {}


if __name__ == "__main__":
    request = json.loads(Path(sys.argv[1]).read_text(encoding="utf-8"))
    limits = Limits(**request["limits"])
    if sys.platform == "linux":
        import resource
        resource.setrlimit(resource.RLIMIT_AS, (limits.parser_memory_bytes, limits.parser_memory_bytes))
    result = extract_file(request["path"], limits=limits, **request["kwargs"])
    Path(sys.argv[2]).write_text(packed(asdict(result)), encoding="utf-8")