"""Source management for the KB admin panel (Phase 3 / Task #14).

Today the only way to clean up or correct a source is `DELETE FROM
kb_source WHERE id=...` on the shared PG plus an mc rm in MinIO. This
module backs the new sources table in the EspoCRM "ניהול" tab:

  list_sources              — table view with chunk count and last-ingest status
  update_source             — patch the user-editable metadata
                              (title, identifier, dates, source_url)
  delete_source             — hard delete (cascades chunks; drops S3 object)
  start_reingest            — kick off the two-phase re-ingest pipeline:
                              classify → awaiting_review → commit. The
                              classify stage reuses
                              ingest_jobs.process_classify_stage so the
                              user gets fresh AI metadata suggestions in
                              the same review UI as fresh uploads.
  process_reingest_embed_stage — final stage after the user commits:
                              UPDATE kb_source metadata + replace labels +
                              re-chunk and re-embed (the original
                              process_reingest_job logic, minus the
                              "preserve metadata" rule).

The reingest path differs from ingest_source: ingest_source creates a
NEW kb_source row and supersedes the previous version. Reingest keeps
the same row id (so cached _lastSearch chunk_index references in the
browser still point at a valid source), wipes its chunks, and inserts
fresh ones from the same file. As of v0.9.0, hand-edited metadata is
NO LONGER preserved silently — the user reviews AI suggestions and
edits them on the batch review screen before committing, so any
overwrite is explicit. Sources whose AI-extracted metadata was thin
(e.g. title="החלטה" only) get fresh suggestions on every re-process.
"""

from __future__ import annotations

import datetime as _dt
import json
import logging
import re
from typing import Optional

from api.services.kb import ingest as kb_ingest
from api.services.kb import s3 as kb_s3
from api.services.kb import voyage
from api.services.kb.chunker import chunk as chunk_text
from api.services.kb.db import get_pool

logger = logging.getLogger("shira.kb.admin_sources")

# Whitelist of fields the user can edit through PUT /admin/kb/sources/{id}.
# Anything else (kind, topic_id, checksum, source_id, ...) is structural
# and must not be flipped from a metadata edit.
_EDITABLE_FIELDS = ("title", "identifier", "source_url", "published_at",
                    "effective_at", "description")


# ── Listing ─────────────────────────────────────────────────────────────────

async def list_sources(
    *,
    topic_id: Optional[int] = None,
    kind: Optional[str] = None,
    label_id: Optional[int] = None,
    limit: int = 200,
) -> list[dict]:
    """Source rows enriched with chunk_count, last_ingest_status, and labels.

    The last-ingest column is a LATERAL join against kb_ingest_job — picks
    the most recent job row for the source (regardless of status), so the
    UI can flag sources whose last re-ingest failed.

    Labels are aggregated as a JSON array via a correlated subquery —
    keeps the row count clean (no N×labels Cartesian product).
    """
    limit = max(1, min(500, limit))
    where: list[str] = ["s.superseded_by IS NULL"]
    args: list = []
    if topic_id is not None:
        args.append(topic_id)
        where.append(f"s.topic_id = ${len(args)}")
    if kind:
        args.append(kind)
        where.append(f"s.kind = ${len(args)}")
    if label_id is not None:
        args.append(label_id)
        where.append(f"EXISTS (SELECT 1 FROM kb_source_label sl "
                     f"WHERE sl.source_id = s.id AND sl.label_id = ${len(args)})")
    args.append(limit)
    sql = f"""
        SELECT
            s.id,
            s.kind,
            s.title,
            s.identifier,
            s.published_at,
            s.effective_at,
            s.source_url,
            s.original_path,
            s.checksum,
            s.topic_id,
            s.description,
            s.ai_classified_at,
            s.created_at,
            s.updated_at,
            (SELECT COUNT(*) FROM kb_chunk c WHERE c.source_id = s.id) AS chunk_count,
            (
                SELECT json_build_object(
                    'job_id', j.id,
                    'status', j.status,
                    'created_at', j.created_at,
                    'completed_at', j.completed_at,
                    'error_message', j.error_message
                )
                FROM kb_ingest_job j
                WHERE j.source_id = s.id
                ORDER BY j.created_at DESC
                LIMIT 1
            ) AS last_ingest,
            COALESCE(
                (
                    SELECT json_agg(json_build_object(
                        'id', l.id, 'slug', l.slug, 'name', l.name,
                        'topic_id', l.topic_id
                    ) ORDER BY l.name)
                    FROM kb_label l
                    JOIN kb_source_label sl ON sl.label_id = l.id
                    WHERE sl.source_id = s.id
                ),
                '[]'::json
            ) AS labels
        FROM kb_source s
        WHERE {' AND '.join(where)}
        ORDER BY s.kind, s.title
        LIMIT ${len(args)}
    """
    pool = await get_pool()
    async with pool.acquire() as conn:
        rows = await conn.fetch(sql, *args)
    return [_row_to_dict(r) for r in rows]


async def get_source(source_id: int) -> Optional[dict]:
    pool = await get_pool()
    async with pool.acquire() as conn:
        row = await conn.fetchrow(
            "SELECT * FROM kb_source WHERE id = $1",
            source_id,
        )
    return _row_to_dict(row) if row else None


# ── Editing ─────────────────────────────────────────────────────────────────

async def update_source(
    source_id: int,
    fields: dict,
    *,
    label_ids: Optional[list[int]] = None,
) -> dict:
    """UPDATE only whitelisted columns. Returns the post-update row.

    If `label_ids` is provided (not None), the source's label set is
    REPLACED to exactly those ids — usage_count adjusted both ways.
    Pass an empty list to clear all labels; pass None to leave labels
    untouched.
    """
    from api.services.kb import labels as kb_labels  # local import: cycle

    clean: dict = {}
    for k in _EDITABLE_FIELDS:
        if k not in fields:
            continue
        v = fields[k]
        if k in ("published_at", "effective_at"):
            clean[k] = kb_ingest._coerce_date(v)
        elif v is None or v == "":
            clean[k] = None
        else:
            clean[k] = str(v).strip()
    if "title" in clean and not clean["title"]:
        raise ValueError("title cannot be empty")

    pool = await get_pool()

    if clean:
        set_pieces = []
        args: list = [source_id]
        for k, v in clean.items():
            args.append(v)
            set_pieces.append(f"{k} = ${len(args)}")
        sql = f"""
            UPDATE kb_source
               SET {', '.join(set_pieces)}
             WHERE id = $1
            RETURNING id
        """
        async with pool.acquire() as conn:
            updated = await conn.fetchrow(sql, *args)
        if not updated:
            raise LookupError(f"source {source_id} not found")
    else:
        # Caller passed nothing editable; verify the source still exists
        # so we can return a sensible 404 vs blindly returning empty.
        async with pool.acquire() as conn:
            existing = await conn.fetchrow(
                "SELECT id FROM kb_source WHERE id = $1", source_id
            )
        if not existing:
            raise LookupError(f"source {source_id} not found")

    if label_ids is not None:
        await kb_labels.replace_for_source(source_id, label_ids)

    # Re-read with labels via list_sources([id])-style query — reuse
    # get_source() for simplicity, then fetch labels separately.
    row = await get_source(source_id)
    if row is not None:
        row["labels"] = await kb_labels.get_labels_for_source(source_id)
    return row


# ── Deletion ────────────────────────────────────────────────────────────────

async def delete_source(source_id: int) -> dict:
    """Hard-delete a source. kb_chunk rows cascade via FK; the underlying
    S3 object is deleted on a best-effort basis (a stray file is harmless,
    but cluttered).
    """
    src = await get_source(source_id)
    if not src:
        raise LookupError(f"source {source_id} not found")

    pool = await get_pool()
    async with pool.acquire() as conn:
        result = await conn.execute(
            "DELETE FROM kb_source WHERE id = $1", source_id
        )
    deleted = result.endswith(" 1")

    # Best-effort S3 cleanup — never fail the whole delete on a stray
    # object that might already be gone.
    s3_deleted = False
    bucket, key = _parse_s3_uri(src.get("original_path") or "")
    if key:
        try:
            kb_s3._client().delete_object(Bucket=bucket or kb_s3._bucket(), Key=key)
            s3_deleted = True
        except Exception as exc:
            logger.warning("[admin.delete] S3 delete failed for %s: %s", key, exc)

    return {"deleted": deleted, "source_id": source_id, "s3_deleted": s3_deleted}


# ── Re-ingest ───────────────────────────────────────────────────────────────

async def start_reingest(
    source_id: int, requested_by_user: Optional[str]
) -> dict:
    """Create a kb_ingest_job tied to the existing source_id.

    The job enters the same two-phase pipeline as a fresh upload —
    classify (AI metadata extraction) → awaiting_review → embed
    (re-chunk + UPDATE kb_source) — so the user can review suggestions
    on the batch review screen before anything overwrites kb_source.

    Returns {"job_id", "batch_id"}; the route schedules
    `ingest_jobs.process_classify_stage(job_id)` after this, and the
    browser opens `/admin/kb/batch/<batch_id>` to poll the review UI.
    """
    import uuid as _uuid

    src = await get_source(source_id)
    if not src:
        raise LookupError(f"source {source_id} not found")
    if not src.get("original_path"):
        raise ValueError(
            f"source {source_id} has no original_path — re-ingest needs the original file"
        )

    # Pre-flight: confirm the file actually exists in S3 before queuing
    # the job. Saves the user from waiting 30s on a doomed job.
    bucket, key = _parse_s3_uri(src["original_path"])
    if not key:
        raise ValueError(f"unparseable original_path: {src['original_path']}")
    try:
        kb_s3._client().head_object(Bucket=bucket or kb_s3._bucket(), Key=key)
    except Exception as exc:
        raise ValueError(f"original file unavailable: {exc}")

    # Reuse the kb_ingest_job table; mark this row as a re-ingest by
    # pre-filling source_id (uploads start with source_id NULL). The
    # batch_id is auto-generated so the existing review UI works
    # without changes — re-ingests are single-file batches. Existing
    # labels go into metadata_json so the review form can pre-fill
    # them as a fallback when the AI classifier hasn't run yet (or
    # returns empty).
    from api.services.kb import labels as kb_labels  # local import: cycle
    existing_labels = await kb_labels.get_labels_for_source(source_id)
    metadata = {
        "title": src.get("title"),
        "identifier": src.get("identifier"),
        "published_at": _date_to_iso(src.get("published_at")),
        "effective_at": _date_to_iso(src.get("effective_at")),
        "source_url": src.get("source_url"),
        "summary": src.get("description"),
        "labels": [lbl.get("name") for lbl in existing_labels if lbl.get("name")],
        "reingest": True,
    }
    batch_id = _uuid.uuid4().hex
    # Use the actual filename (with extension) — NOT the source title.
    # parse_first_pages dispatches by extension (.pdf/.docx/.txt); a
    # title like "החלטה" without an extension makes the parser bail
    # and return empty text, which then makes the classifier return {}.
    # See incident 2026-04-28: source 141 (title="החלטה") had every
    # re-ingest classify as empty until this was fixed.
    pool = await get_pool()
    async with pool.acquire() as conn:
        row = await conn.fetchrow(
            """
            INSERT INTO kb_ingest_job
                (source_id, original_filename, s3_key, kind, topic_id,
                 metadata_json, status, requested_by_user, batch_id)
            VALUES ($1, $2, $3, $4, $5, $6::jsonb, 'queued', $7, $8)
            RETURNING id
            """,
            source_id,
            _filename_from_key(key),
            key,
            src["kind"],
            src["topic_id"],
            json.dumps(metadata, ensure_ascii=False),
            requested_by_user,
            batch_id,
        )
    job_id = row["id"]
    logger.info(
        "[admin.reingest] queued job_id=%s source_id=%s batch_id=%s",
        job_id, source_id, batch_id,
    )
    return {"job_id": job_id, "batch_id": batch_id}


async def process_reingest_embed_stage(
    job_id: int, fields: dict, label_ids: list[int],
) -> None:
    """Final stage of the two-phase re-ingest pipeline.

    Called from `ingest_jobs.commit_job` after the user reviewed the AI
    suggestions. The job is already in `status='processing',
    processing_stage='embedding'` at this point — commit_job set that
    transition before scheduling us.

    Steps:
      1. UPDATE kb_source with the user-confirmed metadata (kind,
         title, identifier, source_url, dates, description=summary).
      2. Replace the source's labels with `label_ids`.
      3. Re-parse the file → chunk → embed → atomic DELETE+INSERT
         of kb_chunk rows + checksum update.
      4. Mark the job done.

    Failures here are real failures (the user already confirmed). They
    surface as `status='failed'` so the polling UI lets the user retry.
    """
    from api.services.kb import labels as kb_labels  # local import: cycle

    pool = await get_pool()
    async with pool.acquire() as conn:
        job_row = await conn.fetchrow(
            "SELECT * FROM kb_ingest_job WHERE id = $1", job_id,
        )
    if not job_row:
        logger.warning("[reingest-embed] job %s vanished", job_id)
        return

    job = _row_to_dict(job_row)
    source_id = job["source_id"]
    if source_id is None:
        await _fail_job(pool, job_id, "re-ingest job has no source_id")
        return

    try:
        # Step 1: UPDATE kb_source with the user-confirmed metadata.
        # commit_job's `fields` dict mirrors JobCommit: kind, title,
        # identifier, source_url, published_at, effective_at, summary.
        # kb_source.description is the schema name for "summary".
        kind = (fields.get("kind") or job["kind"]).strip()
        title = (fields.get("title") or "").strip()
        if not title:
            raise ValueError("title is empty after commit — aborting re-ingest")

        await _update_source_metadata_for_reingest(
            pool=pool,
            source_id=source_id,
            kind=kind,
            title=title,
            identifier=fields.get("identifier"),
            source_url=fields.get("source_url"),
            published_at=fields.get("published_at"),
            effective_at=fields.get("effective_at"),
            description=fields.get("summary"),
        )

        # Step 2: replace labels (empty list clears all).
        await kb_labels.replace_for_source(source_id, label_ids or [])

        # Step 3: re-chunk + re-embed.
        bucket, key = _parse_s3_uri(job["s3_key"])
        if not key:
            key = job["s3_key"]  # bare key, no s3:// prefix
        data = kb_s3._client().get_object(
            Bucket=bucket or kb_s3._bucket(),
            Key=key,
        )["Body"].read()
        if not data:
            raise kb_ingest.IngestError(f"empty file at {key}")

        text = kb_ingest.parse(data, job.get("original_filename") or key)
        if not text:
            raise kb_ingest.IngestError("no text extracted")

        chunks = chunk_text(text, kind)
        if not chunks:
            raise kb_ingest.IngestError("chunker produced zero chunks")

        # Embed first (slow, network-bound) so the actual DB swap is fast.
        texts = [c["content"] for c in chunks]
        vectors = await voyage.embed(texts, input_type="document")

        import hashlib
        new_checksum = hashlib.sha256(text.encode("utf-8")).hexdigest()

        async with pool.acquire() as conn:
            async with conn.transaction():
                # Drop existing chunks and reinsert in one tx so we never
                # leave the source partially-chunked.
                await conn.execute(
                    "DELETE FROM kb_chunk WHERE source_id = $1", source_id
                )
                rows = [
                    (
                        source_id,
                        c["chunk_index"],
                        c["heading_path"],
                        c["section_ref"],
                        c["content"],
                        vectors[i],
                        c["token_count"],
                        c.get("page_number"),
                    )
                    for i, c in enumerate(chunks)
                ]
                await conn.executemany(
                    """
                    INSERT INTO kb_chunk
                        (source_id, chunk_index, heading_path, section_ref,
                         content, embedding, token_count, page_number)
                    VALUES ($1,$2,$3,$4,$5,$6,$7,$8)
                    """,
                    rows,
                )
                await conn.execute(
                    "UPDATE kb_source SET checksum = $2 WHERE id = $1",
                    source_id, new_checksum,
                )
                await conn.execute(
                    """
                    UPDATE kb_ingest_job
                       SET status = 'done',
                           processing_stage = NULL,
                           chunks_created = $2,
                           completed_at = now()
                     WHERE id = $1
                    """,
                    job_id, len(chunks),
                )
        logger.info(
            "[reingest-embed] done job_id=%s source_id=%s chunks=%d labels=%d",
            job_id, source_id, len(chunks), len(label_ids or []),
        )
    except Exception as exc:
        logger.exception("[reingest-embed] failed job_id=%s", job_id)
        await _fail_job(pool, job_id, str(exc))


async def _update_source_metadata_for_reingest(
    *, pool, source_id: int, kind: str, title: str,
    identifier: Optional[str], source_url: Optional[str],
    published_at: Optional[str], effective_at: Optional[str],
    description: Optional[str],
) -> None:
    """UPDATE kb_source after a re-ingest commit. Bypasses _EDITABLE_FIELDS
    so we can also update `kind` (the user explicitly chose it on the
    review screen — e.g. caselaw uploaded as circular by mistake).
    """
    pub = kb_ingest._coerce_date(published_at) if published_at else None
    eff = kb_ingest._coerce_date(effective_at) if effective_at else None
    async with pool.acquire() as conn:
        await conn.execute(
            """
            UPDATE kb_source
               SET kind = $2,
                   title = $3,
                   identifier = $4,
                   source_url = $5,
                   published_at = $6,
                   effective_at = $7,
                   description = $8
             WHERE id = $1
            """,
            source_id,
            kind,
            title,
            (identifier or None),
            (source_url or None),
            pub,
            eff,
            (description or None),
        )


# ── helpers ─────────────────────────────────────────────────────────────────

async def _fail_job(pool, job_id: int, msg: str) -> None:
    async with pool.acquire() as conn:
        await conn.execute(
            """
            UPDATE kb_ingest_job
               SET status = 'failed',
                   error_message = $2,
                   completed_at = now()
             WHERE id = $1
            """,
            job_id, msg[:4000],
        )


def _row_to_dict(row) -> Optional[dict]:
    if row is None:
        return None
    out: dict = {}
    for k, v in dict(row).items():
        if isinstance(v, _dt.datetime):
            out[k] = v.isoformat()
        elif isinstance(v, _dt.date):
            out[k] = v.isoformat()
        elif k == "metadata_json" and isinstance(v, str):
            try:
                out[k] = json.loads(v)
            except Exception:
                out[k] = {}
        elif k == "last_ingest" and isinstance(v, str):
            try:
                out[k] = json.loads(v)
            except Exception:
                out[k] = None
        elif k == "labels" and isinstance(v, str):
            try:
                out[k] = json.loads(v)
            except Exception:
                out[k] = []
        else:
            out[k] = v
    return out


_S3_URI = re.compile(r"^s3://([^/]+)/(.+)$")

def _parse_s3_uri(uri: str) -> tuple[str, str]:
    """('bucket', 'key/path') or ('', '') if not an s3:// URI."""
    if not uri:
        return ("", "")
    m = _S3_URI.match(uri)
    if not m:
        return ("", "")
    return (m.group(1), m.group(2))


def _filename_from_key(key: str) -> str:
    return key.rsplit("/", 1)[-1] if key else "source"


def _date_to_iso(value) -> Optional[str]:
    if value is None:
        return None
    if isinstance(value, str):
        return value[:10] if value else None
    if isinstance(value, _dt.date):
        return value.isoformat()
    return None