"""Add (ingest) artifacts into the local store."""

from __future__ import annotations

from dataclasses import dataclass
from pathlib import Path
from typing import Any, Dict, List, Tuple
import os
import shutil
import uuid

from .hash import canonicalize_bytes, sha256_hex
from .store import ObjectStore
from ..chunking.registry import get_method
from ..chunking.api import ChunkerContext, ManifestChunk
from ..manifest.schema import DocManifest, CollectionManifest, CollectionEntry
from ..manifest.codec import dump_manifest_bytes, dump_collection_bytes, load_manifest_bytes
from ..index.local import LocalIndex


class AddError(Exception):
    pass


@dataclass(frozen=True)
class AddResult:
    doc_id: str
    manifest_hash: str
    chunk_count: int
    format: str
    chunking: Dict[str, Any]
    path: Path


@dataclass(frozen=True)
class AddBatchResult:
    results: List[AddResult]


@dataclass(frozen=True)
class AddCollectionResult:
    collection_id: str
    documents: List[AddResult]


def _ensure_initialized(store: ObjectStore) -> None:
    if not store.objects_dir().exists():
        raise AddError("store not initialized; run 'gnos init'")


def _infer_format_chunking(path: Path) -> Tuple[str, str]:
    ext = path.suffix.lower()
    if ext in {".md", ".markdown"}:
        return "markdown", "framed-sequential-v1"
    if ext in {".py"}:
        return "python", "framed-sequential-v1"
    if ext in {".json"}:
        return "json", "framed-sequential-v1"
    if ext in {".txt"}:
        return "text", "framed-sequential-v1"
    return "text", "framed-sequential-v1"


def _chunking_object(method: str) -> Dict[str, Any]:
    if method == "identity-v1":
        return {"method": method, "params": {}}
    if method == "sequential-v1":
        return {"method": method, "params": {"chunk_size": 1024 * 1024}}
    if method == "framed-sequential-v1":
        return {"method": method, "params": {"max_payload_bytes": 1024 * 1024}}
    return {"method": method, "params": {}}


def _add_file(
    path: Path,
    store: ObjectStore,
    format: str = "auto",
    chunking: str = "auto",
    chunking_params: Dict[str, Any] | None = None,
    index: LocalIndex | None = None,
    _staging_store: ObjectStore | None = None,
) -> AddResult:
    """Add a single file as a document and store its manifest/chunks."""
    _ensure_initialized(store)
    if format == "auto" or chunking == "auto":
        inferred_format, inferred_method = _infer_format_chunking(path)
        if format == "auto":
            format = inferred_format
        if chunking == "auto":
            chunking = inferred_method
    chunking_obj = _chunking_object(chunking)
    if chunking_params:
        # Merge user-provided overrides into defaults.
        chunking_obj = {"method": chunking_obj["method"], "params": {**chunking_obj.get("params", {}), **chunking_params}}

    raw = path.read_bytes()
    canon = canonicalize_bytes(raw)
    doc_id = sha256_hex(canon)

    method = get_method(chunking_obj["method"])
    ctx = ChunkerContext(path=str(path), format=format)
    split_result = method.split(canon, chunking_obj.get("params", {}), ctx)
    if not isinstance(split_result.params, dict) or not all(isinstance(k, str) for k in split_result.params.keys()):
        raise AddError("chunker returned invalid params (must be dict with string keys)")
    chunking_obj = {"method": chunking_obj["method"], "params": split_result.params}
    chunks = split_result.chunks

    chunk_entries: List[Dict[str, Any]] = []
    reconstruct_chunks: List[ManifestChunk] = []
    target_store = _staging_store or store

    for chunk in chunks:
        if "hash" in chunk.entry:
            raise AddError("chunk entry must not contain 'hash' (computed by gnos)")
        if not all(isinstance(k, str) for k in chunk.entry.keys()):
            raise AddError("chunk entry keys must be strings")

        h = sha256_hex(chunk.data)
        entry = dict(chunk.entry)
        entry["hash"] = h
        chunk_entries.append(entry)
        reconstruct_chunks.append(ManifestChunk(data=chunk.data, entry=entry))

        if not target_store.has(h):
            target_store.put(h, chunk.data)

    # Enforce identity round-trip at add time (no side effects beyond staging).
    reconstructed = method.reconstruct(reconstruct_chunks, chunking_obj.get("params", {}))
    if reconstructed != canon:
        raise AddError("chunker reconstruction does not match document bytes")

    manifest = DocManifest(
        type="doc-manifest",
        version=1,
        doc_id=doc_id,
        format=format,
        chunking=chunking_obj,
        chunks=chunk_entries,
    )
    manifest_bytes = dump_manifest_bytes(manifest)
    manifest_hash = sha256_hex(manifest_bytes)
    if not target_store.has(manifest_hash):
        target_store.put(manifest_hash, manifest_bytes)

    if index is not None:
        index.record_manifest(doc_id, manifest_hash, format, chunking_obj, chunk_entries)

    return AddResult(
        doc_id=doc_id,
        manifest_hash=manifest_hash,
        chunk_count=len(chunk_entries),
        format=format,
        chunking=chunking_obj,
        path=path,
    )


def add_bytes(
    data: bytes,
    store: ObjectStore,
    format: str | None = None,
    chunking: str = "auto",
    chunking_params: Dict[str, Any] | None = None,
    index: LocalIndex | None = None,
    dry_run: bool = False,
    label: str = "<bytes>",
) -> AddResult:
    """Add raw bytes as a document (no filesystem path).

    `format` may be None; in that case `ChunkerContext.format` is None and
    chunkers that require a format must obtain it from explicit params.
    """
    _ensure_initialized(store)

    if chunking == "auto":
        # Without a filename extension, default to the core byte-rooted chunker.
        chunking = "framed-sequential-v1"

    chunking_obj = _chunking_object(chunking)
    if chunking_params:
        chunking_obj = {"method": chunking_obj["method"], "params": {**chunking_obj.get("params", {}), **chunking_params}}

    canon = canonicalize_bytes(data)
    doc_id = sha256_hex(canon)

    method = get_method(chunking_obj["method"])
    ctx = ChunkerContext(path=None, format=format)
    split_result = method.split(canon, chunking_obj.get("params", {}), ctx)
    if not isinstance(split_result.params, dict) or not all(isinstance(k, str) for k in split_result.params.keys()):
        raise AddError("chunker returned invalid params (must be dict with string keys)")
    chunking_obj = {"method": chunking_obj["method"], "params": split_result.params}
    chunks = split_result.chunks

    chunk_entries: List[Dict[str, Any]] = []
    reconstruct_chunks: List[ManifestChunk] = []

    staging_root = store.root / "tmp" / f"add-{uuid.uuid4().hex}"
    staging_store = ObjectStore(staging_root)
    staging_store.init()

    try:
        for chunk in chunks:
            if "hash" in chunk.entry:
                raise AddError("chunk entry must not contain 'hash' (computed by gnos)")
            if not all(isinstance(k, str) for k in chunk.entry.keys()):
                raise AddError("chunk entry keys must be strings")

            h = sha256_hex(chunk.data)
            entry = dict(chunk.entry)
            entry["hash"] = h
            chunk_entries.append(entry)
            reconstruct_chunks.append(ManifestChunk(data=chunk.data, entry=entry))

            if not staging_store.has(h):
                staging_store.put(h, chunk.data)

        reconstructed = method.reconstruct(reconstruct_chunks, chunking_obj.get("params", {}))
        if reconstructed != canon:
            raise AddError("chunker reconstruction does not match document bytes")

        manifest = DocManifest(
            type="doc-manifest",
            version=1,
            doc_id=doc_id,
            format=format or "text",
            chunking=chunking_obj,
            chunks=chunk_entries,
        )
        manifest_bytes = dump_manifest_bytes(manifest)
        manifest_hash = sha256_hex(manifest_bytes)
        if not staging_store.has(manifest_hash):
            staging_store.put(manifest_hash, manifest_bytes)

        if dry_run:
            return AddResult(
                doc_id=doc_id,
                manifest_hash=manifest_hash,
                chunk_count=len(chunk_entries),
                format=format or "text",
                chunking=chunking_obj,
                path=Path(label),
            )

        _commit_staged_objects(staging_store, store)
    finally:
        shutil.rmtree(staging_root, ignore_errors=True)

    if index is not None:
        index.record_manifest(doc_id, manifest_hash, format or "text", chunking_obj, chunk_entries)

    return AddResult(
        doc_id=doc_id,
        manifest_hash=manifest_hash,
        chunk_count=len(chunk_entries),
        format=format or "text",
        chunking=chunking_obj,
        path=Path(label),
    )


def add_path(
    path: Path,
    store: ObjectStore,
    format: str = "auto",
    chunking: str = "auto",
    chunking_params: Dict[str, Any] | None = None,
    recursive: bool = True,
    index: LocalIndex | None = None,
    dry_run: bool = False,
) -> AddBatchResult:
    """Add a file or directory (recursively) and return per-file results."""
    _ensure_initialized(store)
    results, staging_root, staging_store = _stage_path(
        path,
        store,
        format=format,
        chunking=chunking,
        chunking_params=chunking_params,
        recursive=recursive,
    )
    if dry_run:
        shutil.rmtree(staging_root, ignore_errors=True)
        return AddBatchResult(results=results)
    _commit_staged_objects(staging_store, store)
    shutil.rmtree(staging_root, ignore_errors=True)
    _update_index_from_results(results, store, index)
    return AddBatchResult(results=results)


def _stage_path(
    path: Path,
    store: ObjectStore,
    format: str,
    chunking: str,
    chunking_params: Dict[str, Any] | None,
    recursive: bool,
) -> tuple[List[AddResult], Path, ObjectStore]:
    """Stage add results into a temporary store (no side effects on main store)."""
    results: List[AddResult] = []
    staging_root = store.root / "tmp" / f"add-{uuid.uuid4().hex}"
    staging_store = ObjectStore(staging_root)
    staging_store.init()

    try:
        if path.is_file():
            results.append(
                _add_file(
                    path,
                    store,
                    format=format,
                    chunking=chunking,
                    chunking_params=chunking_params,
                    index=None,
                    _staging_store=staging_store,
                )
            )
        else:
            if not path.is_dir():
                raise AddError("path is neither file nor directory")

            pattern = "**/*" if recursive else "*"
            for p in sorted(path.glob(pattern)):
                if p.is_file():
                    results.append(
                        _add_file(
                            p,
                            store,
                            format=format,
                            chunking=chunking,
                            chunking_params=chunking_params,
                            index=None,
                            _staging_store=staging_store,
                        )
                    )
    except Exception:
        shutil.rmtree(staging_root, ignore_errors=True)
        raise

    return results, staging_root, staging_store


def _commit_staged_objects(staging_store: ObjectStore, store: ObjectStore) -> None:
    """Commit staged objects into the main store."""
    for obj in staging_store.objects_dir().rglob("*"):
        if not obj.is_file():
            continue
        obj_hash = obj.name
        dest = store._path_for(obj_hash)
        dest.parent.mkdir(parents=True, exist_ok=True)
        if not dest.exists():
            os.replace(obj, dest)
        else:
            obj.unlink(missing_ok=True)


def _update_index_from_results(results: List[AddResult], store: ObjectStore, index: LocalIndex | None) -> None:
    if index is None:
        return
    # Recompute chunk entries from stored manifests to keep index consistent.
    for result in results:
        manifest_bytes = store.get(result.manifest_hash)
        if manifest_bytes is None:
            continue
        manifest = load_manifest_bytes(manifest_bytes)
        index.record_manifest(
            manifest.doc_id,
            result.manifest_hash,
            manifest.format,
            manifest.chunking,
            manifest.chunks,
        )


def _sort_collection_entries(entries: List[CollectionEntry]) -> List[CollectionEntry]:
    # Deterministic ordering: path, then doc_id.
    return sorted(entries, key=lambda e: (e.path or "", e.doc_id))


def add_collection(
    path: Path,
    store: ObjectStore,
    format: str = "auto",
    chunking: str = "auto",
    chunking_params: Dict[str, Any] | None = None,
    recursive: bool = True,
    index: LocalIndex | None = None,
    dry_run: bool = False,
) -> AddCollectionResult:
    """Add a path and emit a collection manifest for the resulting documents."""
    _ensure_initialized(store)
    results, staging_root, staging_store = _stage_path(
        path,
        store,
        format=format,
        chunking=chunking,
        chunking_params=chunking_params,
        recursive=recursive,
    )

    if path.is_dir():
        root = path
    else:
        root = path.parent

    entries: List[CollectionEntry] = []
    for result in results:
        rel = result.path.relative_to(root).as_posix()
        entries.append(CollectionEntry(doc_id=result.doc_id, path=rel))

    documents = _sort_collection_entries(entries)
    collection = CollectionManifest(type="collection-manifest", version=1, documents=documents)
    collection_bytes = dump_collection_bytes(collection)
    collection_id = sha256_hex(collection_bytes)
    if not dry_run:
        if not staging_store.has(collection_id):
            staging_store.put(collection_id, collection_bytes)

    if dry_run:
        shutil.rmtree(staging_root, ignore_errors=True)
        return AddCollectionResult(collection_id=collection_id, documents=results)

    _commit_staged_objects(staging_store, store)
    shutil.rmtree(staging_root, ignore_errors=True)
    _update_index_from_results(results, store, index)

    return AddCollectionResult(collection_id=collection_id, documents=results)
