"""
Celery tasks for downloading and indexing document sources.

Redis key  "index_job"  stores current job state as JSON:
  { status, source, progress, message, started_at, finished_at, docs_total, docs_done }
"""
import bz2
import json
import os
import time
import requests
import xml.etree.ElementTree as ET
from datetime import datetime, timezone
from pathlib import Path

import redis
from app.celery_app import celery_app

from app.config import settings
from app.models.db import SessionLocal, Document, Fingerprint

JOB_KEY = "index_job"

WIKIPEDIA_SOURCES = {
    "wikipedia_ru": {
        "url": "https://dumps.wikimedia.org/ruwiki/latest/ruwiki-latest-pages-articles.xml.bz2",
        "meta_url": "https://dumps.wikimedia.org/ruwiki/latest/ruwiki-latest-md5sums.txt",
        "lang": "ru",
        "label": "Wikipedia RU",
    },
    "wikipedia_en": {
        "url": "https://dumps.wikimedia.org/enwiki/latest/enwiki-latest-pages-articles.xml.bz2",
        "meta_url": "https://dumps.wikimedia.org/enwiki/latest/enwiki-latest-md5sums.txt",
        "lang": "en",
        "label": "Wikipedia EN",
    },
}


# ── Redis helpers ──────────────────────────────────────────────────────────

def _r() -> redis.Redis:
    return redis.from_url(settings.REDIS_URL)


def _set_status(data: dict):
    _r().set(JOB_KEY, json.dumps(data, default=str))


def _get_status() -> dict:
    raw = _r().get(JOB_KEY)
    if not raw:
        return {"status": "idle"}
    return json.loads(raw)


def _pub(message: str, progress: int = 0, **extra):
    state = _get_status()
    state.update({"message": message, "progress": progress, **extra})
    _set_status(state)
    _r().publish("index_job_progress", json.dumps(state, default=str))


# ── Text helpers ───────────────────────────────────────────────────────────

import re


def _clean_wiki(text: str) -> str:
    text = re.sub(r'\[\[([^\|]*\|)?([^\]]*)\]\]', r'\2', text)
    text = re.sub(r'\[http[^\s]+ ([^\]]*)\]', r'\1', text)
    text = re.sub(r'\{\{[^}]*\}\}', '', text)
    text = re.sub(r'\[\[Category:[^\]]*\]\]', '', text)
    text = re.sub(r'<ref[^>]*>.*?</ref>', '', text, flags=re.DOTALL)
    text = re.sub(r'<[^>]+>', '', text)
    text = re.sub(r'\n{3,}', '\n\n', text)
    text = re.sub(r' +', ' ', text)
    return text.strip()


def _normalize(text: str) -> str:
    text = re.sub(r'[\x00-\x08\x0b-\x0c\x0e-\x1f\x7f]', '', text)
    text = re.sub(r'[—–‒―]', '-', text)
    text = re.sub(r'[ \t]+', ' ', text)
    text = re.sub(r'\n{3,}', '\n\n', text)
    return text.strip()


# ── Download helper ────────────────────────────────────────────────────────

def _download_file(url: str, dest: Path, label: str) -> Path:
    """Stream-download url → dest, publishing progress."""
    dest.parent.mkdir(parents=True, exist_ok=True)
    r = requests.get(url, stream=True, timeout=30)
    r.raise_for_status()
    total = int(r.headers.get("Content-Length", 0))
    done = 0
    last_pct = -1

    with open(dest, "wb") as f:
        for chunk in r.iter_content(chunk_size=1 << 20):  # 1 MB
            f.write(chunk)
            done += len(chunk)
            pct = int(done / total * 40) if total else 0  # 0-40% of overall
            if pct != last_pct:
                _pub(f"Скачивание {label}: {done >> 20} MB", progress=pct)
                last_pct = pct

    return dest


# ── Indexing helper ────────────────────────────────────────────────────────

def _index_wikipedia_bz2(bz2_path: Path, source: str, lang: str, update_only: bool) -> int:
    """
    Parse Wikipedia bz2 dump and index documents.
    update_only=True → skip ext_ids already in DB.
    Returns count of newly indexed docs.
    """
    from app.core.fingerprint import WinNowing, MinHashLSHIndex
    from app.core.embeddings import EmbeddingIndex
    from app.core.fulltext import FulltextSearch

    db = SessionLocal()
    winnowing = WinNowing()
    minhash = MinHashLSHIndex()
    emb = EmbeddingIndex()
    fulltext = FulltextSearch()

    # Pre-load known ext_ids for fast skip
    known: set[str] = set()
    if update_only:
        for (ext_id,) in db.query(Document.ext_id).filter(Document.source == source).yield_per(10_000):
            if ext_id:
                known.add(ext_id)

    ns = {'wiki': 'http://www.mediawiki.org/xml/export-0.10/'}

    embed_batch: list[tuple[str, str]] = []
    BATCH = 32
    count = 0
    parsed = 0

    with bz2.BZ2File(str(bz2_path), 'r') as f:
        context = ET.iterparse(f, events=('end',))
        for event, elem in context:
            if not elem.tag.endswith('page'):
                continue

            title_el = elem.find('wiki:title', ns) or elem.find('.//title')
            id_el = elem.find('wiki:id', ns) or elem.find('.//id')
            rev = elem.find('.//revision')
            text_el = rev.find('wiki:text', ns) if rev is not None else None
            if rev is None:
                text_el = elem.find('.//text')

            if title_el is None or text_el is None:
                elem.clear(); continue

            title = (title_el.text or '').strip()
            page_id = (id_el.text or '').strip() if id_el is not None else ''
            raw_text = text_el.text or ''
            ext_id = f"wiki_{page_id}"

            if raw_text.startswith('#REDIRECT') or title.startswith('Wikipedia:'):
                elem.clear(); continue

            text = _normalize(_clean_wiki(raw_text))
            if len(text.split()) < 50:
                elem.clear(); continue

            parsed += 1

            if update_only and ext_id in known:
                elem.clear(); continue

            url = f"https://{lang}.wikipedia.org/wiki/{title.replace(' ', '_')}"

            # PostgreSQL
            doc = Document(source=source, ext_id=ext_id, title=title, url=url, lang=lang)
            db.add(doc)
            db.flush()

            # Fingerprints
            fp = winnowing.fingerprint(text)
            for pos, h in enumerate(list(fp)[:500]):
                db.add(Fingerprint(doc_id=doc.id, hash=format(h, '016x'), position=pos))

            # MinHash
            try:
                minhash.add_document(str(doc.id), text)
            except ValueError:
                pass

            # Elasticsearch
            fulltext.index_document(str(doc.id), title, text, source, lang, url)

            embed_batch.append((str(doc.id), text))
            if len(embed_batch) >= BATCH:
                emb.batch_add(embed_batch, batch_size=BATCH)
                embed_batch.clear()

            count += 1
            if count % 500 == 0:
                db.commit()
                pct = 40 + min(int(count / 100000 * 55), 55)  # 40-95%
                _pub(f"Проиндексировано: {count} документов", progress=pct, docs_done=count)

            elem.clear()

    if embed_batch:
        emb.batch_add(embed_batch, batch_size=BATCH)

    db.commit()
    db.close()
    emb.save_index()

    return count


# ── Celery tasks ───────────────────────────────────────────────────────────

def _run_indexing(source_key: str, update_only: bool):
    src = WIKIPEDIA_SOURCES[source_key]
    label = src["label"]
    lang = src["lang"]

    _set_status({
        "status": "running",
        "source": source_key,
        "label": label,
        "progress": 0,
        "message": "Запуск…",
        "started_at": datetime.now(timezone.utc).isoformat(),
        "docs_done": 0,
    })

    data_dir = Path(settings.RAW_DATA_DIR) / source_key
    bz2_path = data_dir / "dump.xml.bz2"

    # --- Download if needed ---
    if not bz2_path.exists() or not update_only:
        _pub(f"Скачивание дампа {label}…", progress=1)
        try:
            _download_file(src["url"], bz2_path, label)
        except Exception as e:
            _set_status({"status": "error", "message": f"Ошибка скачивания: {e}"})
            return

    _pub("Индексация…", progress=41)

    try:
        count = _index_wikipedia_bz2(bz2_path, source_key, lang, update_only)
    except Exception as e:
        _set_status({"status": "error", "message": f"Ошибка индексации: {e}"})
        return

    action = "обновлено" if update_only else "проиндексировано"
    _set_status({
        "status": "done",
        "source": source_key,
        "label": label,
        "progress": 100,
        "message": f"Готово — {action} {count} документов",
        "finished_at": datetime.now(timezone.utc).isoformat(),
        "docs_done": count,
    })
    _r().publish("index_job_progress", _r().get(JOB_KEY))


@celery_app.task(name="tasks.download_and_index", bind=True)
def download_and_index(self, source_key: str):
    _run_indexing(source_key, update_only=False)


@celery_app.task(name="tasks.update_index", bind=True)
def update_index(self, source_key: str):
    _run_indexing(source_key, update_only=True)


def get_job_status() -> dict:
    return _get_status()