"""Lazy archive resolver for archive-backed Normattiva documents."""

from __future__ import annotations

import json
import logging
import re
import tarfile
import tempfile
from pathlib import Path
from urllib.parse import urlparse

try:
    import requests as _requests
    HAS_REQUESTS = True
except ImportError:
    HAS_REQUESTS = False

log = logging.getLogger(__name__)


class ArchiveBackedDocumentResolver:
    """Download, extract, and search archive-backed legal documents on demand."""

    def __init__(self, dataset_repo: str, hf_token: str, cache_dir: str | Path | None = None):
        self.dataset_repo = dataset_repo
        self.hf_token = hf_token
        self.cache_dir = Path(cache_dir) if cache_dir else Path(tempfile.gettempdir()) / "italaw_archive_cache"
        self.cache_dir.mkdir(parents=True, exist_ok=True)
        self._archive_cache: dict[str, Path] = {}
        # Per-archive manifest: source_archive -> { urn -> internal_path }
        self._manifest_index: dict[str, dict] = {}

    def resolve_document_text(self, metadata: dict) -> str | None:
        """Resolve a chunk or document metadata dict to full text from an archive."""
        source_archive = (metadata.get("source_archive") or "").strip()
        urn = (metadata.get("urn") or "").strip()
        if not source_archive:
            return None
        extracted_root = self._ensure_archive_extracted(source_archive)
        if not extracted_root:
            return None
        internal_path = (metadata.get("archive_internal_path") or "").strip()
        return self._extract_text_from_root(extracted_root, urn=urn, internal_path=internal_path)

    def _ensure_archive_extracted(self, source_archive: str) -> Path | None:
        if source_archive in self._archive_cache:
            return self._archive_cache[source_archive]

        local_path = self._resolve_source_to_local_path(source_archive)
        if not local_path:
            return None
        if local_path.is_dir():
            self._archive_cache[source_archive] = local_path
            return local_path

        extracted = self._extract_archive(local_path)
        if extracted:
            self._archive_cache[source_archive] = extracted
            # Build a lightweight manifest index for fast lookup (jsonl files contain URNs)
            try:
                self._index_archive_manifest(source_archive, extracted)
            except Exception:
                pass
        return extracted

    def _resolve_source_to_local_path(self, source_archive: str) -> Path | None:
        source_archive = (source_archive or "").strip()
        if not source_archive:
            return None

        if source_archive.startswith("http://") or source_archive.startswith("https://"):
            return self._download_url(source_archive)

        local_path = Path(source_archive)
        if local_path.exists():
            return local_path

        if source_archive.startswith("sources/") or source_archive.startswith("normattiva/"):
            url = f"https://huggingface.co/datasets/{self.dataset_repo}/resolve/main/{source_archive}?download=true"
            return self._download_url(url)

        return None

    def _download_url(self, url: str) -> Path | None:
        if not HAS_REQUESTS:
            return None
        parsed = urlparse(url)
        filename = Path(parsed.path or "download.bin").name or "download.bin"
        local_path = self.cache_dir / filename
        if local_path.exists() and local_path.stat().st_size > 0:
            return local_path
        try:
            with _requests.get(url, timeout=120, stream=True, headers={"Authorization": f"Bearer {self.hf_token}"}) as resp:
                resp.raise_for_status()
                with open(local_path, "wb") as handle:
                    for chunk in resp.iter_content(chunk_size=1024 * 1024):
                        if chunk:
                            handle.write(chunk)
            return local_path
        except Exception as exc:
            log.warning("Archive download failed for %s: %s", url, exc)
            return None

    def _extract_archive(self, path: Path) -> Path | None:
        extract_dir = self.cache_dir / f"extract_{path.stem}"
        extract_dir.mkdir(parents=True, exist_ok=True)
        try:
            lower = path.name.lower()
            if lower.endswith(".zip"):
                import zipfile
                with zipfile.ZipFile(path) as archive:
                    archive.extractall(extract_dir)
                return extract_dir
            if lower.endswith((".tar", ".tar.gz", ".tgz", ".tar.bz2", ".tar.xz")):
                with tarfile.open(path) as archive:
                    archive.extractall(extract_dir)
                return extract_dir
            if lower.endswith((".tar.zst", ".zst")):
                import pyzstd
                tar_path = extract_dir / path.name.replace(".zst", "")
                with open(path, "rb") as src, open(tar_path, "wb") as dst:
                    pyzstd.decompress_stream(src, dst)
                with tarfile.open(tar_path) as archive:
                    archive.extractall(extract_dir)
                return extract_dir
        except Exception as exc:
            log.warning("Archive extraction failed for %s: %s", path, exc)
            return None
        return path if path.is_dir() else None

    def _extract_text_from_root(self, root: Path, urn: str = "", internal_path: str = "") -> str | None:
        # Fast path: if we have a manifest mapping for this archive, try it first
        try:
            # Source archive key is present if this root was created during extraction
            for key, mapped in self._manifest_index.items():
                if mapped and urn and urn in mapped:
                    candidate = root / mapped[urn]
                    if candidate.exists():
                        text = self._extract_text_from_path(candidate, urn)
                        if text:
                            return text
        except Exception:
            pass

        if internal_path:
            candidate = root / internal_path
            if candidate.exists():
                text = self._extract_text_from_path(candidate, urn)
                if text:
                    return text

        for xml_path in root.rglob("*.xml"):
            text = self._extract_text_from_path(xml_path, urn)
            if text:
                return text
        for jsonl_path in root.rglob("*.jsonl"):
            text = self._extract_text_from_jsonl(jsonl_path, urn)
            if text:
                return text
        return None

    def _extract_text_from_path(self, path: Path, urn: str = "") -> str | None:
        if path.suffix.lower() == ".jsonl":
            return self._extract_text_from_jsonl(path, urn)
        if path.suffix.lower() == ".xml":
            doc = _extract_akn_text(path)
            if doc and (not urn or doc.get("urn") == urn):
                return doc.get("body") or doc.get("text")
        return None

    def _extract_text_from_jsonl(self, path: Path, urn: str = "") -> str | None:
        try:
            with open(path, encoding="utf-8") as handle:
                for line in handle:
                    line = line.strip()
                    if not line:
                        continue
                    try:
                        doc = json.loads(line)
                    except json.JSONDecodeError:
                        continue
                    doc_urn = doc.get("urn") or doc.get("id") or doc.get("url") or ""
                    if urn and doc_urn != urn:
                        continue
                    text = doc.get("text") or doc.get("body") or doc.get("testo") or ""
                    if text:
                        return text
        except Exception as exc:
            log.warning("JSONL archive search failed for %s: %s", path, exc)
        return None

    def _index_archive_manifest(self, source_archive: str, root: Path) -> None:
        """Scan jsonl files under the extracted root and build a mapping urn -> internal_path."""
        mapping: dict = {}
        try:
            for jsonl_path in root.rglob("*.jsonl"):
                rel = jsonl_path.relative_to(root)
                try:
                    with open(jsonl_path, encoding="utf-8") as fh:
                        for line in fh:
                            try:
                                doc = json.loads(line)
                            except Exception:
                                continue
                            doc_urn = doc.get("urn") or doc.get("id") or doc.get("url") or ""
                            if doc_urn and doc_urn not in mapping:
                                mapping[doc_urn] = str(rel)
                except Exception:
                    continue
        except Exception:
            pass
        self._manifest_index[source_archive] = mapping

    def find_internal_path_for_urn(self, source_archive: str, urn: str) -> str | None:
        m = self._manifest_index.get(source_archive, {})
        return m.get(urn)


def _extract_akn_text(xml_path: Path) -> dict | None:
    import xml.etree.ElementTree as ET

    def _name(tag: str) -> str:
        return tag.rsplit("}", 1)[-1]

    try:
        root = ET.parse(xml_path).getroot()
    except Exception:
        return None

    urn = ""
    title = ""
    body_text = ""
    pub_date = ""

    for elem in root.iter():
        elem_name = _name(elem.tag)
        if elem_name in {"FRBRuri", "FRBRthis"} and not urn:
            urn = elem.attrib.get("value", "")
        elif elem_name == "FRBRdate" and not pub_date:
            pub_date = elem.attrib.get("date", "")
        elif elem_name in {"docTitle", "shortTitle", "title", "heading"} and not title:
            text = " ".join(part.strip() for part in elem.itertext() if part.strip())
            if len(text) > 8:
                title = text

    for elem in root.iter():
        if _name(elem.tag) in {"body", "article", "paragraph", "section", "chapter"}:
            text = " ".join(part.strip() for part in elem.itertext() if part.strip())
            if len(text) > len(body_text):
                body_text = text

    if not body_text:
        return None

    return {
        "title": title or xml_path.stem,
        "urn": urn or xml_path.stem,
        "source_url": str(xml_path),
        "pub_date": pub_date,
        "body": body_text,
    }