from __future__ import annotations

import argparse
import hashlib
import json
import re
from pathlib import Path


ROOT = Path(__file__).resolve().parents[1]
DEFAULT_SOURCE_DIR = ROOT / "plugins" / "1c" / "rag" / "sources"
DEFAULT_OUTPUT = ROOT / "plugins" / "1c" / "datasets" / "prepared" / "rag_corpus.jsonl"
SUPPORTED_EXTENSIONS = {".md", ".txt", ".bsl", ".os"}
HEADING_RE = re.compile(r"^(#{1,6})\s+(.+)$")
FRONT_MATTER_RE = re.compile(r"^---\n(?P<body>.*?)\n---\n?", re.DOTALL)


def normalize_text(text: str) -> str:
    lines = [line.rstrip() for line in text.replace("\r\n", "\n").split("\n")]
    return "\n".join(lines).strip()


def parse_front_matter(text: str) -> tuple[dict[str, str], str]:
    match = FRONT_MATTER_RE.match(text)
    if not match:
        return {}, text
    metadata: dict[str, str] = {}
    for line in match.group("body").splitlines():
        if ":" not in line:
            continue
        key, value = line.split(":", 1)
        value = value.strip().strip('"').strip("'")
        metadata[key.strip()] = value
    return metadata, text[match.end() :].strip()


def split_paragraphs(text: str) -> list[str]:
    blocks: list[str] = []
    current: list[str] = []
    for line in text.split("\n"):
        stripped = line.strip()
        if not stripped:
            if current:
                blocks.append("\n".join(current).strip())
                current = []
            continue
        if HEADING_RE.match(stripped) and current:
            blocks.append("\n".join(current).strip())
            current = [stripped]
        else:
            current.append(line.rstrip())
    if current:
        blocks.append("\n".join(current).strip())
    return [block for block in blocks if block]


def hard_split_text(text: str, chunk_size: int, overlap: int) -> list[str]:
    chunks: list[str] = []
    start = 0
    while start < len(text):
        end = min(start + chunk_size, len(text))
        chunk = text[start:end].strip()
        if chunk:
            chunks.append(chunk)
        if end == len(text):
            break
        start = end - overlap
    return chunks


def chunk_text(text: str, chunk_size: int, overlap: int) -> list[dict]:
    if chunk_size <= 0:
        raise ValueError("chunk_size must be positive")
    if overlap < 0:
        raise ValueError("overlap must not be negative")
    if overlap >= chunk_size:
        raise ValueError("overlap must be smaller than chunk_size")

    chunks: list[dict] = []
    current: list[str] = []
    current_heading = ""
    last_tail = ""
    headings_seen: list[str] = []

    def flush() -> None:
        nonlocal current, last_tail
        if not current:
            return
        content = "\n\n".join(current).strip()
        if last_tail and not content.startswith(last_tail):
            content = f"{last_tail}\n\n{content}".strip()
        chunks.append(
            {
                "content": content,
                "heading": current_heading,
                "headings": headings_seen[-4:],
            }
        )
        last_tail = content[-overlap:].strip() if overlap else ""
        current = []

    for block in split_paragraphs(text):
        heading_match = HEADING_RE.match(block.split("\n", 1)[0].strip())
        if heading_match:
            current_heading = heading_match.group(2).strip()
            headings_seen.append(current_heading)

        if len(block) > chunk_size:
            flush()
            for part in hard_split_text(block, chunk_size, overlap):
                chunks.append(
                    {
                        "content": part,
                        "heading": current_heading,
                        "headings": headings_seen[-4:],
                    }
                )
                last_tail = part[-overlap:].strip() if overlap else ""
            continue

        candidate = "\n\n".join([*current, block]).strip()
        if current and len(candidate) > chunk_size:
            flush()
        current.append(block)

    flush()
    return chunks


def stable_id(*parts: str) -> str:
    digest = hashlib.sha256("|".join(parts).encode("utf-8")).hexdigest()
    return digest[:16]


def classify_source(path: Path, text: str, front_matter: dict[str, str] | None = None) -> str:
    front_matter = front_matter or {}
    if front_matter.get("source") == "official_1c_its":
        return front_matter.get("source_type") or "official_1c_its"
    name = path.name.lower()
    lowered = text[:4000].lower()
    if path.suffix.lower() in {".bsl", ".os"}:
        return "bsl"
    if "1c bsl module snapshot" in lowered or "```bsl" in lowered or "процедура " in lowered:
        return "bsl"
    if "1c metadata snapshot" in lowered or "metadata snapshot" in lowered:
        return "metadata"
    if "read-only" in lowered or "только чтение" in lowered or "выбрать" in lowered and "запрос" in lowered:
        return "query"
    if "production" in lowered or "резервн" in lowered or "согласован" in lowered or "опасн" in lowered:
        return "safety"
    if "metadata" in name:
        return "metadata"
    if "bsl" in name or "module" in name:
        return "bsl"
    return "docs"


def iter_source_files(source_dir: Path) -> list[Path]:
    if not source_dir.exists():
        return []
    return sorted(
        path
        for path in source_dir.rglob("*")
        if path.is_file() and path.suffix.lower() in SUPPORTED_EXTENSIONS
    )


def build_records(source_dir: Path, chunk_size: int, overlap: int) -> list[dict]:
    records: list[dict] = []
    for path in iter_source_files(source_dir):
        relative_path = path.relative_to(source_dir).as_posix()
        text = normalize_text(path.read_text(encoding="utf-8"))
        if not text:
            continue
        front_matter, body = parse_front_matter(text)
        chunk_source_text = body or text

        source_type = classify_source(path, chunk_source_text, front_matter)
        chunks = chunk_text(chunk_source_text, chunk_size=chunk_size, overlap=overlap)
        document_id = stable_id(relative_path, chunk_source_text)
        base_title = front_matter.get("title") or path.stem.replace("_", " ").replace("-", " ")

        for index, chunk in enumerate(chunks):
            content = chunk["content"]
            title = chunk.get("heading") or base_title
            records.append(
                {
                    "id": stable_id(document_id, str(index), content),
                    "document_id": document_id,
                    "source_path": relative_path,
                    "title": title,
                    "chunk_index": index,
                    "content": content,
                    "metadata": {
                        "domain": "1c",
                        "source_type": source_type,
                        "file_type": path.suffix.lower().lstrip("."),
                        "access": front_matter.get("access") or "",
                        "official_source": front_matter.get("source") or "",
                        "url": front_matter.get("url") or "",
                        "platform_family": front_matter.get("platform_family") or "",
                        "platform_version": front_matter.get("platform_version") or "",
                        "platform_doc_id": front_matter.get("platform_doc_id") or "",
                        "doc_book": front_matter.get("doc_book") or "",
                        "doc_bookmark": front_matter.get("doc_bookmark") or "",
                        "doc_content_id": front_matter.get("doc_content_id") or "",
                        "doc_coordinate": front_matter.get("doc_coordinate") or "",
                        "access_blocked": front_matter.get("access_blocked") or "",
                        "access_findings": front_matter.get("access_findings") or "",
                        "source_sha256": front_matter.get("source_sha256") or "",
                        "chunk_size": chunk_size,
                        "overlap": overlap,
                        "heading": chunk.get("heading") or "",
                        "headings": chunk.get("headings") or [],
                    },
                }
            )
    return records


def build_manifest(source_dir: Path, records: list[dict], chunk_size: int, overlap: int) -> dict:
    by_source: dict[str, dict] = {}
    for record in records:
        source_path = record["source_path"]
        source = by_source.setdefault(
            source_path,
            {
                "source_path": source_path,
                "title": record.get("title"),
                "source_type": (record.get("metadata") or {}).get("source_type"),
                "file_type": (record.get("metadata") or {}).get("file_type"),
                "document_id": record.get("document_id"),
                "chunk_count": 0,
                "content_hash": "",
            },
        )
        source["chunk_count"] += 1

    for relative_path, source in by_source.items():
        path = source_dir / relative_path
        text = normalize_text(path.read_text(encoding="utf-8")) if path.exists() else ""
        source["content_hash"] = hashlib.sha256(text.encode("utf-8")).hexdigest()

    return {
        "schema_version": 1,
        "source_dir": str(source_dir),
        "chunk_size": chunk_size,
        "overlap": overlap,
        "source_count": len(by_source),
        "chunk_count": len(records),
        "sources": sorted(by_source.values(), key=lambda item: item["source_path"]),
    }


def write_jsonl(path: Path, records: list[dict]) -> None:
    path.parent.mkdir(parents=True, exist_ok=True)
    with path.open("w", encoding="utf-8") as handle:
        for record in records:
            handle.write(json.dumps(record, ensure_ascii=False) + "\n")


def main() -> int:
    parser = argparse.ArgumentParser(description="Prepare a JSONL RAG corpus for the 1C plugin.")
    parser.add_argument("--source-dir", type=Path, default=DEFAULT_SOURCE_DIR)
    parser.add_argument("--output", type=Path, default=DEFAULT_OUTPUT)
    parser.add_argument("--manifest", type=Path)
    parser.add_argument("--chunk-size", type=int, default=1800)
    parser.add_argument("--overlap", type=int, default=200)
    args = parser.parse_args()

    records = build_records(
        source_dir=args.source_dir,
        chunk_size=args.chunk_size,
        overlap=args.overlap,
    )
    write_jsonl(args.output, records)
    manifest_path = args.manifest or args.output.with_suffix(".manifest.json")
    manifest_path.parent.mkdir(parents=True, exist_ok=True)
    manifest_path.write_text(
        json.dumps(build_manifest(args.source_dir, records, args.chunk_size, args.overlap), ensure_ascii=False, indent=2) + "\n",
        encoding="utf-8",
    )
    print(f"Wrote {len(records)} chunk(s) to {args.output}")
    print(f"Wrote RAG manifest to {manifest_path}")
    return 0


if __name__ == "__main__":
    raise SystemExit(main())
