"""원문 위치를 보존하는 SQLite 코퍼스 저장소.

검색용 전처리 텍스트와 증거 표시용 원문을 분리한다. 현재 수령한 XLSX처럼
페이지 정보가 없는 데이터도 ``coordinate_scope=episode`` 로 정직하게 기록하며,
향후 PDF/DOCX 재추출 시 동일 스키마에 page/paragraph/char offset을 채울 수 있다.
"""

from __future__ import annotations

import hashlib
import json
import sqlite3
from dataclasses import dataclass, field
from pathlib import Path
from typing import Iterable, Iterator


#: v2 — 참조 lemma/요소 특징 캐시 컬럼 추가 (#5). 기존 v1 DB 는 ALTER TABLE 로
#: 자동 승격되며, 캐시가 비어 있으면 질의 시 계산 후 백필된다.
SCHEMA_VERSION = 2


def text_sha256(text: str) -> str:
    return hashlib.sha256(text.encode("utf-8")).hexdigest()


def stable_id(prefix: str, *parts: str) -> str:
    payload = "\x1f".join(parts).encode("utf-8")
    return f"{prefix}-{hashlib.sha256(payload).hexdigest()[:20]}"


@dataclass(frozen=True)
class DocumentRecord:
    document_id: str
    title: str
    source_path: str | None = None
    source_sha256: str | None = None
    metadata: dict = field(default_factory=dict)


@dataclass(frozen=True)
class SegmentRecord:
    segment_id: str
    document_id: str
    text: str
    ordinal: str
    coordinate_scope: str = "episode"
    page_number: int | None = None
    paragraph_number: int | None = None
    char_start: int | None = None
    char_end: int | None = None
    source_locator: str | None = None
    metadata: dict = field(default_factory=dict)
    #: 인덱싱 시점에 계산해 둔 참조 lemma 열 (#5). None 이면 미계산 상태이며,
    #: 질의 경로가 계산 후 백필한다. 형태소 분석을 요청마다 반복하지 않기 위한 캐시.
    lemmas: list[str] | None = None
    #: 인물/모티프 등 최소 요소 특징. ExtractedElements 를 dict 로 직렬화한 형태.
    elements: dict | None = None

    @property
    def text_sha256(self) -> str:
        return text_sha256(self.text)


class CorpusStore:
    """SQLite 기반 원문/세그먼트 저장소.

    연결은 호출 단위로 열어 멀티프로세스 API에서도 안전하게 사용한다. WAL은
    읽기 중 증분 적재를 허용한다.
    """

    def __init__(self, path: str | Path):
        self.path = Path(path)

    def _connect(self) -> sqlite3.Connection:
        self.path.parent.mkdir(parents=True, exist_ok=True)
        con = sqlite3.connect(self.path)
        con.row_factory = sqlite3.Row
        con.execute("PRAGMA journal_mode=WAL")
        con.execute("PRAGMA foreign_keys=ON")
        return con

    def initialize(self) -> None:
        with self._connect() as con:
            con.executescript(
                """
                CREATE TABLE IF NOT EXISTS corpus_meta (
                    key TEXT PRIMARY KEY,
                    value TEXT NOT NULL
                );
                CREATE TABLE IF NOT EXISTS documents (
                    document_id TEXT PRIMARY KEY,
                    title TEXT NOT NULL,
                    source_path TEXT,
                    source_sha256 TEXT,
                    metadata_json TEXT NOT NULL DEFAULT '{}',
                    created_at TEXT NOT NULL DEFAULT CURRENT_TIMESTAMP,
                    updated_at TEXT NOT NULL DEFAULT CURRENT_TIMESTAMP
                );
                CREATE TABLE IF NOT EXISTS segments (
                    segment_id TEXT PRIMARY KEY,
                    document_id TEXT NOT NULL REFERENCES documents(document_id) ON DELETE CASCADE,
                    ordinal TEXT NOT NULL,
                    text TEXT NOT NULL,
                    text_sha256 TEXT NOT NULL,
                    coordinate_scope TEXT NOT NULL,
                    page_number INTEGER,
                    paragraph_number INTEGER,
                    char_start INTEGER,
                    char_end INTEGER,
                    source_locator TEXT,
                    metadata_json TEXT NOT NULL DEFAULT '{}',
                    created_at TEXT NOT NULL DEFAULT CURRENT_TIMESTAMP,
                    UNIQUE(document_id, text_sha256)
                );
                CREATE INDEX IF NOT EXISTS idx_segments_document ON segments(document_id);
                CREATE INDEX IF NOT EXISTS idx_segments_hash ON segments(text_sha256);
                """
            )
            self._migrate(con)
            con.execute(
                "INSERT OR REPLACE INTO corpus_meta(key, value) VALUES('schema_version', ?)",
                (str(SCHEMA_VERSION),),
            )

    @staticmethod
    def _migrate(con: sqlite3.Connection) -> None:
        """기존 DB 를 파괴 없이 승격. v1 → v2 는 컬럼 추가만 필요하다."""
        existing = {row["name"] for row in con.execute("PRAGMA table_info(segments)")}
        for column in ("lemmas_json", "elements_json"):
            if column not in existing:
                con.execute(f"ALTER TABLE segments ADD COLUMN {column} TEXT")

    def upsert_document(self, record: DocumentRecord) -> None:
        self.upsert_documents([record])

    def upsert_documents(self, records: Iterable[DocumentRecord]) -> None:
        self.initialize()
        with self._connect() as con:
            con.executemany(
                """
                INSERT INTO documents(document_id,title,source_path,source_sha256,metadata_json)
                VALUES(?,?,?,?,?)
                ON CONFLICT(document_id) DO UPDATE SET
                  title=excluded.title,
                  source_path=COALESCE(excluded.source_path, documents.source_path),
                  source_sha256=COALESCE(excluded.source_sha256, documents.source_sha256),
                  metadata_json=excluded.metadata_json,
                  updated_at=CURRENT_TIMESTAMP
                """,
                [(
                    record.document_id, record.title, record.source_path,
                    record.source_sha256,
                    json.dumps(record.metadata, ensure_ascii=False, sort_keys=True),
                ) for record in records],
            )

    def add_segments(self, records: Iterable[SegmentRecord]) -> tuple[int, int]:
        """세그먼트를 추가하고 (inserted, duplicates) 반환."""
        self.initialize()
        inserted = duplicates = 0
        with self._connect() as con:
            for r in records:
                before = con.total_changes
                con.execute(
                    """
                    INSERT OR IGNORE INTO segments(
                      segment_id,document_id,ordinal,text,text_sha256,coordinate_scope,
                      page_number,paragraph_number,char_start,char_end,source_locator,
                      metadata_json,lemmas_json,elements_json
                    ) VALUES(?,?,?,?,?,?,?,?,?,?,?,?,?,?)
                    """,
                    (
                        r.segment_id,
                        r.document_id,
                        r.ordinal,
                        r.text,
                        r.text_sha256,
                        r.coordinate_scope,
                        r.page_number,
                        r.paragraph_number,
                        r.char_start,
                        r.char_end,
                        r.source_locator,
                        json.dumps(r.metadata, ensure_ascii=False, sort_keys=True),
                        json.dumps(r.lemmas, ensure_ascii=False) if r.lemmas is not None else None,
                        json.dumps(r.elements, ensure_ascii=False, sort_keys=True)
                        if r.elements is not None else None,
                    ),
                )
                if con.total_changes > before:
                    inserted += 1
                else:
                    duplicates += 1
        return inserted, duplicates

    def iter_segments(self) -> Iterator[SegmentRecord]:
        if not self.path.exists():
            return
        with self._connect() as con:
            rows = con.execute(
                """
                SELECT s.*, d.title
                FROM segments s JOIN documents d USING(document_id)
                ORDER BY s.document_id, s.ordinal, s.segment_id
                """
            )
            for row in rows:
                yield self._row_to_segment(row)

    @staticmethod
    def _row_to_segment(row: sqlite3.Row) -> SegmentRecord:
        metadata = json.loads(row["metadata_json"] or "{}")
        metadata.setdefault("document_title", row["title"])
        keys = row.keys()
        lemmas = None
        elements = None
        # v1 DB 를 그대로 읽는 경로에서도 죽지 않도록 컬럼 존재를 확인한다.
        if "lemmas_json" in keys and row["lemmas_json"]:
            lemmas = json.loads(row["lemmas_json"])
        if "elements_json" in keys and row["elements_json"]:
            elements = json.loads(row["elements_json"])
        return SegmentRecord(
            segment_id=row["segment_id"],
            document_id=row["document_id"],
            text=row["text"],
            ordinal=row["ordinal"],
            coordinate_scope=row["coordinate_scope"],
            page_number=row["page_number"],
            paragraph_number=row["paragraph_number"],
            char_start=row["char_start"],
            char_end=row["char_end"],
            source_locator=row["source_locator"],
            metadata=metadata,
            lemmas=lemmas,
            elements=elements,
        )

    def update_segment_features(
        self, features: Iterable[tuple[str, list[str], dict]]
    ) -> int:
        """(segment_id, lemmas, elements) 를 백필한다. 반환값은 갱신된 행 수."""
        rows = [
            (
                json.dumps(lemmas, ensure_ascii=False),
                json.dumps(elements, ensure_ascii=False, sort_keys=True),
                segment_id,
            )
            for segment_id, lemmas, elements in features
        ]
        if not rows:
            return 0
        self.initialize()
        with self._connect() as con:
            con.executemany(
                "UPDATE segments SET lemmas_json=?, elements_json=? WHERE segment_id=?",
                rows,
            )
            return con.total_changes

    def count_missing_features(self) -> int:
        if not self.path.exists():
            return 0
        self.initialize()
        with self._connect() as con:
            return int(con.execute(
                "SELECT COUNT(*) FROM segments WHERE lemmas_json IS NULL"
            ).fetchone()[0])

    def get_segments(self, segment_ids: Iterable[str]) -> dict[str, SegmentRecord]:
        ids = list(dict.fromkeys(segment_ids))
        if not ids or not self.path.exists():
            return {}
        result: dict[str, SegmentRecord] = {}
        with self._connect() as con:
            for start in range(0, len(ids), 500):
                batch = ids[start : start + 500]
                marks = ",".join("?" for _ in batch)
                rows = con.execute(
                    f"""SELECT s.*, d.title FROM segments s
                    JOIN documents d USING(document_id)
                    WHERE segment_id IN ({marks})""",
                    batch,
                )
                for row in rows:
                    result[row["segment_id"]] = self._row_to_segment(row)
        return result

    def stats(self) -> dict[str, int]:
        if not self.path.exists():
            return {"documents": 0, "segments": 0, "characters": 0}
        with self._connect() as con:
            documents = con.execute("SELECT COUNT(*) FROM documents").fetchone()[0]
            segments, characters = con.execute(
                "SELECT COUNT(*), COALESCE(SUM(LENGTH(text)),0) FROM segments"
            ).fetchone()
        return {"documents": documents, "segments": segments, "characters": characters}

    def document_count(self) -> int:
        if not self.path.exists():
            return 0
        with self._connect() as con:
            return int(con.execute("SELECT COUNT(*) FROM documents").fetchone()[0])

    def list_documents(self) -> list[dict]:
        if not self.path.exists():
            return []
        with self._connect() as con:
            rows = con.execute(
                """SELECT d.document_id, d.title, d.source_path,
                COUNT(s.segment_id) AS segment_count,
                COALESCE(SUM(LENGTH(s.text)),0) AS characters
                FROM documents d LEFT JOIN segments s USING(document_id)
                GROUP BY d.document_id ORDER BY d.title"""
            ).fetchall()
        return [dict(row) for row in rows]

    def document_source_groups(self) -> dict[str, str]:
        """문서별 익명 작성자 그룹을 반환한다.

        ``author_group`` 은 수령 XLSX 적재기의 기존 필드이고, ``source_group`` 은
        새 수집 원문·학습 데이터에서 쓰는 공통 필드다. 둘 다 가명 값만 허용한다.
        영속 인덱스가 자기 작성자의 후보를 검색 단계에서 제외할 때 사용한다.
        """
        if not self.path.exists():
            return {}
        with self._connect() as con:
            rows = con.execute(
                "SELECT document_id, metadata_json FROM documents"
            ).fetchall()
        groups: dict[str, str] = {}
        for row in rows:
            metadata = json.loads(row["metadata_json"] or "{}")
            group = str(
                metadata.get("source_group") or metadata.get("author_group") or ""
            ).strip()
            if group:
                groups[str(row["document_id"])] = group
        return groups

    def documents_with_metadata(self) -> list[DocumentRecord]:
        """배치·감사용 문서 목록. 원문 세그먼트는 포함하지 않는다."""
        if not self.path.exists():
            return []
        with self._connect() as con:
            rows = con.execute(
                """SELECT document_id, title, source_path, source_sha256, metadata_json
                FROM documents ORDER BY document_id"""
            ).fetchall()
        return [DocumentRecord(
            document_id=str(row["document_id"]),
            title=str(row["title"]),
            source_path=row["source_path"],
            source_sha256=row["source_sha256"],
            metadata=json.loads(row["metadata_json"] or "{}"),
        ) for row in rows]

    def delete_document(self, document_id: str) -> bool:
        if not self.path.exists():
            return False
        with self._connect() as con:
            before = con.total_changes
            con.execute("DELETE FROM documents WHERE document_id=?", (document_id,))
            return con.total_changes > before