from __future__ import annotations

import argparse
import asyncio
import json
import os
from pathlib import Path
from typing import Any, Dict, Iterable, List, Sequence
from uuid import uuid4

from idxr.pipelines.document_ingest import (
    DocumentIngestContext,
    DocumentIngestPipeline,
    DocumentIngestRequest,
)
from idxr.pipelines.document_ingest.interfaces import (
    DocumentDownloader,
    EmbeddingStrategy,
)
from idxr.pipelines.document_ingest.strategies import (
    LocalFileDocumentDownloader,
    ChromaVectorStoreWriter,
    DefaultMetadataBuilder,
    HttpDocumentDownloader,
    InMemoryChromaClientFactory,
    LocalChromaClientFactory,
    MimeTypeTextExtractor,
    OpenAIEmbeddingStrategy,
    SentenceTransformerEmbeddingStrategy,
    TokenCountingChunker,
)
from idxr.pipelines.document_ingest.strategies.vector_store import (
    ChromaClientFactory,
)
from idxr.pipelines.document_ingest.crawlers import (
    DiscoveredSource,
    crawl_entry_points,
    crawl_local_root,
    load_index_locations,
)


def _parse_metadata(pairs: Iterable[str]) -> Dict[str, Any]:
    metadata: Dict[str, Any] = {}
    for pair in pairs:
        if "=" not in pair:
            raise argparse.ArgumentTypeError(
                f"Malformed metadata entry '{pair}'. Expected key=value."
            )
        key, value = pair.split("=", 1)
        metadata[key.strip()] = value.strip()
    return metadata


def _load_metadata_json(path: str | None) -> Dict[str, Any]:
    if not path:
        return {}
    data_path = Path(path).expanduser()
    if not data_path.exists():
        raise argparse.ArgumentTypeError(f"Metadata file '{data_path}' not found.")
    try:
        raw = json.loads(data_path.read_text(encoding="utf-8"))
    except json.JSONDecodeError as exc:  # pragma: no cover - defensive
        raise argparse.ArgumentTypeError(
            f"Invalid JSON in metadata file '{data_path}': {exc}"
        ) from exc
    if not isinstance(raw, dict):
        raise argparse.ArgumentTypeError(
            f"Metadata file '{data_path}' must contain a JSON object."
        )
    return raw


def _is_local_index(index_list: List[str]) -> bool:
    """Check if any index location appears to be a local file path."""
    for index_location in index_list:
        # If it doesn't start with http:// or https://, assume it's local
        if not index_location.startswith(("http://", "https://")):
            return True
    return False


def _guess_mime_type(file_ext: str) -> str | None:
    """Guess MIME type from file extension."""
    mime_map = {
        ".pdf": "application/pdf",
        ".html": "text/html",
        ".htm": "text/html",
        ".txt": "text/plain",
        ".md": "text/markdown",
        ".json": "application/json",
        ".xml": "application/xml",
    }
    return mime_map.get(file_ext)


def _generate_environment_configs(
    args: argparse.Namespace,
    collection_name: str,
    total_chunks: int,
    embedding_backend: str,
    actual_embedding_model: str,
    chunk_encoder_model: str,
    chunk_size_tokens: int,
    chunk_overlap_tokens: int,
) -> None:
    """
    Generate environment.json and query_config.json files after ingestion.

    Args:
        args: CLI arguments
        collection_name: Name of the ingested collection
        total_chunks: Total number of chunks ingested
        embedding_backend: Name of the embedding backend used
        actual_embedding_model: The actual embedding model name used
        chunk_encoder_model: Tokenizer model used for chunking
        chunk_size_tokens: Maximum tokens per chunk
        chunk_overlap_tokens: Token overlap between chunks
    """
    from .env_generator import generate_environment_configs

    persist_dir = Path(args.persist_dir).expanduser().resolve()

    try:
        query_config_path, environment_path = generate_environment_configs(
            persist_dir=persist_dir,
            collection_name=collection_name,
            total_chunks=total_chunks,
            embedding_backend=embedding_backend,
            embedding_model=actual_embedding_model,
            model_registry_target=args.model_registry_target,
            discriminator_field="collection_name",
            query_config_output=args.query_config_output,
            environment_output=args.environment_output,
            chunk_encoder_model=chunk_encoder_model,
            chunk_size_tokens=chunk_size_tokens,
            chunk_overlap_tokens=chunk_overlap_tokens,
        )

        print(f"\nQuery config written to {query_config_path}")
        print(f"Environment config written to {environment_path}")
        print("\nTo start ChromaSQL server, run:")
        print(f"  poetry run python -m chromasql.cli --client 'local:{persist_dir}'")
        print("\nOr use the environment config programmatically:")
        print("  from idxr.models import CollectionEnvironment")
        print(f"  env = CollectionEnvironment.from_file('{environment_path}')")
    except OSError as exc:
        print(f"Warning: Failed to generate environment configs: {exc}")


def build_pipeline_from_args(args: argparse.Namespace) -> DocumentIngestPipeline:
    downloader: DocumentDownloader
    # Determine if we're processing local files or URLs
    uses_local_files = (
        args.file or args.local_root or (args.index and _is_local_index(args.index))
    )

    if uses_local_files:
        downloader = LocalFileDocumentDownloader(
            max_bytes=args.max_download_bytes,
            default_mime_type=args.mime_type,
        )
    else:
        downloader = HttpDocumentDownloader(
            allowed_mime_types=("application/pdf", "application/octet-stream"),
            max_bytes=args.max_download_bytes,
        )
    extractor = MimeTypeTextExtractor()
    chunker = TokenCountingChunker(
        encoder_model=args.chunk_encoder_model,
        default_chunk_size=args.chunk_size_tokens,
        default_overlap=args.chunk_overlap_tokens,
    )

    chroma_client_type = args.chroma_client.lower()
    if chroma_client_type == "memory":
        client_factory: ChromaClientFactory = InMemoryChromaClientFactory()
    else:
        persist_dir = Path(args.persist_dir).expanduser().resolve()
        client_factory = LocalChromaClientFactory(persist_dir=persist_dir)

    vector_writer = ChromaVectorStoreWriter(
        client_factory=client_factory,
        batch_size=args.batch_size,
    )

    embedding_strategies: Dict[str, EmbeddingStrategy] = {}
    if not args.disable_local_embeddings:
        embedding_strategies["sentence_transformer"] = (
            SentenceTransformerEmbeddingStrategy(
                model_name=args.sentence_transformer_model,
                batch_size=args.embedding_batch_size,
            )
        )

    openai_api_key = os.getenv("OPENAI_API_KEY")
    if args.enable_openai and openai_api_key:
        embedding_strategies["openai"] = OpenAIEmbeddingStrategy(
            model_name=args.openai_embedding_model,
            batch_size=args.embedding_batch_size,
        )

    if not embedding_strategies:
        raise SystemExit(
            "No embedding backend available. Enable local embeddings or set "
            "OPENAI_API_KEY/--enable-openai."
        )

    default_backend = args.embedding_backend or next(iter(embedding_strategies))
    if default_backend not in embedding_strategies:
        raise SystemExit(
            f"Embedding backend '{default_backend}' is not configured. "
            f"Available: {list(embedding_strategies)}"
        )

    return DocumentIngestPipeline(
        downloader=downloader,
        extractor=extractor,
        chunker=chunker,
        metadata_builder=DefaultMetadataBuilder(),
        vector_store_writer=vector_writer,
        embedding_strategies=embedding_strategies,
        default_embedding_backend=default_backend,
    )


async def run_pipeline(args: argparse.Namespace) -> None:
    pipeline = build_pipeline_from_args(args)

    base_metadata = {}
    base_metadata.update(_parse_metadata(args.metadata or ()))
    base_metadata.update(_load_metadata_json(args.metadata_json))

    sources = await discover_sources(args)
    if not sources:
        raise SystemExit(
            "No sources discovered. Provide a URL, --file, --index, "
            "--entry-point, or --local-root."
        )

    total = len(sources)
    total_chunks_ingested = 0
    for idx, source in enumerate(sources, start=1):
        metadata = dict(base_metadata)
        metadata.setdefault("experiment", args.experiment_name)
        metadata.setdefault("source_kind", source.kind)
        metadata.update(source.metadata)
        if source.display_url:
            metadata.setdefault("source_url", source.display_url)

        request = DocumentIngestRequest(
            source_url=source.fetch_url,
            collection_name=args.collection,
            metadata=metadata,
            chunk_size_tokens=args.chunk_size_tokens,
            chunk_overlap_tokens=args.chunk_overlap_tokens,
            embedding_backend=args.embedding_backend,
            mime_type=source.mime_type or args.mime_type,
        )
        context = DocumentIngestContext(
            conversation_id=uuid4(),
            default_metadata={"experiment": args.experiment_name},
        )

        result = await pipeline.run(request, context)
        total_chunks_ingested += result.inserted_chunks
        print(
            json.dumps(
                {
                    "index": idx,
                    "total": total,
                    "source": metadata.get("source_url", source.fetch_url),
                    "inserted_chunks": result.inserted_chunks,
                    "collection": result.collection_name,
                },
                indent=2,
            )
        )

    # Generate environment and query config if requested
    if args.generate_environment and args.chroma_client == "persistent":
        # Determine the actual embedding model used based on the backend
        embedding_backend = pipeline._default_embedding_backend
        if embedding_backend == "sentence_transformer":
            actual_embedding_model = args.sentence_transformer_model
        elif embedding_backend == "openai":
            actual_embedding_model = args.openai_embedding_model
        else:
            # Fallback to chunk encoder model (shouldn't happen in practice)
            actual_embedding_model = args.chunk_encoder_model

        _generate_environment_configs(
            args=args,
            collection_name=args.collection,
            total_chunks=total_chunks_ingested,
            embedding_backend=embedding_backend,
            actual_embedding_model=actual_embedding_model,
            chunk_encoder_model=args.chunk_encoder_model,
            chunk_size_tokens=args.chunk_size_tokens,
            chunk_overlap_tokens=args.chunk_overlap_tokens,
        )


async def discover_sources(args: argparse.Namespace) -> List[DiscoveredSource]:
    sources: list[DiscoveredSource] = []

    if args.url:
        sources.append(
            DiscoveredSource(
                fetch_url=args.url,
                kind="direct",
                display_url=args.url,
                metadata={"source_kind": "direct"},
            )
        )

    if args.file:
        file_path = Path(args.file).expanduser()
        sources.append(
            DiscoveredSource(
                fetch_url=str(file_path),
                kind="file",
                display_url=str(file_path),
                metadata={"source_kind": "file", "original_path": str(file_path)},
                mime_type=args.mime_type,
            )
        )

    if args.index:
        for index_location in args.index:
            index_urls = await load_index_locations(
                index_location, timeout=args.crawler_timeout
            )
            for url in index_urls:
                # Auto-detect MIME type for local files based on extension
                mime_type = None
                if not url.startswith(("http://", "https://")):
                    file_ext = Path(url).suffix.lower()
                    mime_type = _guess_mime_type(file_ext)

                sources.append(
                    DiscoveredSource(
                        fetch_url=url,
                        kind="index",
                        display_url=url,
                        metadata={
                            "source_kind": "index",
                            "index_source": index_location,
                        },
                        mime_type=mime_type,
                    )
                )

    if args.entry_point:
        entry_urls = await crawl_entry_points(
            args.entry_point,
            max_pages=args.max_pages,
            max_depth=args.max_depth,
            same_domain=not args.no_same_domain,
            timeout=args.crawler_timeout,
        )
        for url in entry_urls:
            sources.append(
                DiscoveredSource(
                    fetch_url=url,
                    kind="crawl",
                    display_url=url,
                    metadata={"source_kind": "crawl"},
                )
            )

    if args.local_root:
        local_sources = crawl_local_root(
            Path(args.local_root),
            url_prefix=args.url_prefix,
            include_exts=args.local_include_ext,
            max_files=args.local_max_files,
        )
        sources.extend(local_sources)

    deduped: list[DiscoveredSource] = []
    seen: set[str] = set()
    for source in sources:
        if source.fetch_url in seen:
            continue
        deduped.append(source)
        seen.add(source.fetch_url)
    return deduped


def parse_args(argv: list[str] | None = None) -> argparse.Namespace:
    parser = argparse.ArgumentParser(
        description="Run the document ingest pipeline against a single document.",
    )
    parser.add_argument(
        "url",
        nargs="?",
        help="Public URL to a document.",
    )
    parser.add_argument(
        "--file",
        help="Path to a local document to ingest.",
    )
    parser.add_argument(
        "--collection",
        default="document_ingest_local",
        help="Chroma collection name to target (default: %(default)s).",
    )
    parser.add_argument(
        "--experiment-name",
        default="local_cli",
        help="Value inserted into chunk metadata for tracing runs.",
    )
    parser.add_argument(
        "--chroma-client",
        choices=["memory", "persistent"],
        default="memory",
        help="Use in-memory or persistent Chroma client (default: %(default)s).",
    )
    parser.add_argument(
        "--persist-dir",
        default="./.chroma-cli",
        help="Directory for persistent client (default: %(default)s).",
    )
    parser.add_argument(
        "--chunk-size-tokens",
        type=int,
        default=800,
        help="Maximum tokens per chunk (default: %(default)s).",
    )
    parser.add_argument(
        "--chunk-overlap-tokens",
        type=int,
        default=200,
        help="Token overlap between consecutive chunks (default: %(default)s).",
    )
    parser.add_argument(
        "--chunk-encoder-model",
        default="text-embedding-3-small",
        help="Tokenizer model used for chunking (default: %(default)s).",
    )
    parser.add_argument(
        "--embedding-backend",
        choices=["sentence_transformer", "openai"],
        help="Force a specific embedding backend. Defaults to the first available.",
    )
    parser.add_argument(
        "--sentence-transformer-model",
        default="sentence-transformers/all-MiniLM-L6-v2",
        help="SentenceTransformer model used for local embeddings.",
    )
    parser.add_argument(
        "--openai-embedding-model",
        default="text-embedding-3-small",
        help="OpenAI embedding model when --enable-openai is active.",
    )
    parser.add_argument(
        "--disable-local-embeddings",
        action="store_true",
        help="Skip creating a local sentence-transformer backend.",
    )
    parser.add_argument(
        "--enable-openai",
        action="store_true",
        help="Include the OpenAI embedding backend (requires OPENAI_API_KEY).",
    )
    parser.add_argument(
        "--embedding-batch-size",
        type=int,
        default=16,
        help="Batch size used by embedding backends (default: %(default)s).",
    )
    parser.add_argument(
        "--batch-size",
        type=int,
        default=64,
        help="Batch size for Chroma upserts (default: %(default)s).",
    )
    parser.add_argument(
        "--metadata",
        action="append",
        help="Custom metadata key=value pairs (can be repeated).",
    )
    parser.add_argument(
        "--metadata-json",
        help="Path to a JSON file merged into chunk metadata.",
    )
    parser.add_argument(
        "--mime-type",
        default="application/pdf",
        help="Default MIME type for downloaded documents (default: %(default)s).",
    )
    parser.add_argument(
        "--max-download-bytes",
        type=int,
        default=15 * 1024 * 1024,
        help="Maximum bytes allowed for downloads or local reads (default: 15MB).",
    )
    parser.add_argument(
        "--index",
        action="append",
        help="URL or path to a sitemap / newline index (repeatable).",
    )
    parser.add_argument(
        "--entry-point",
        action="append",
        help="Seed URL for recursive crawl (repeatable).",
    )
    parser.add_argument(
        "--max-pages",
        type=int,
        default=50,
        help="Maximum pages to crawl from entry points (default: %(default)s).",
    )
    parser.add_argument(
        "--max-depth",
        type=int,
        default=2,
        help="Maximum link depth during entry-point crawl (default: %(default)s).",
    )
    parser.add_argument(
        "--no-same-domain",
        action="store_true",
        help="Allow entry-point crawl to leave the seed domains.",
    )
    parser.add_argument(
        "--crawler-timeout",
        type=float,
        default=15.0,
        help="Timeout in seconds for crawler HTTP requests (default: %(default)s).",
    )
    parser.add_argument(
        "--local-root",
        help="Path to a local website mirror to ingest recursively.",
    )
    parser.add_argument(
        "--local-include-ext",
        action="append",
        help="File extensions (without dot) to include from --local-root.",
    )
    parser.add_argument(
        "--local-max-files",
        type=int,
        default=500,
        help="Maximum number of local files to ingest (default: %(default)s).",
    )
    parser.add_argument(
        "--url-prefix",
        default="https://local.test",
        help="URL prefix used when rewriting local filesystem paths.",
    )
    parser.add_argument(
        "--generate-environment",
        action="store_true",
        help="Generate environment.json after ingestion for use with ChromaSQL.",
    )
    parser.add_argument(
        "--environment-output",
        type=Path,
        help=(
            "Path where environment.json will be written "
            "(default: persist-dir/environment.json)."
        ),
    )
    parser.add_argument(
        "--query-config-output",
        type=Path,
        help=(
            "Path where query_config.json will be written "
            "(default: persist-dir/query_config.json)."
        ),
    )
    parser.add_argument(
        "--model-registry-target",
        default="idxr.pipelines.document_ingest.registry:MODEL_REGISTRY",
        help="Import path to model registry (default: %(default)s).",
    )
    args = parser.parse_args(argv)

    if not any([args.url, args.file, args.index, args.entry_point, args.local_root]):
        parser.error(
            "Provide at least one source: URL, --file, --index, --entry-point, or --local-root."
        )
    return args


def main(argv: Sequence[str]) -> int:
    args = parse_args(list(argv))
    try:
        asyncio.run(run_pipeline(args))
        return 0
    except KeyboardInterrupt:  # pragma: no cover - user cancel
        return 130
    except Exception:
        return 1


if __name__ == "__main__":  # pragma: no cover
    import sys

    raise SystemExit(main(sys.argv[1:]))
