from langchain_community.document_loaders import PyPDFLoader, DirectoryLoader
from langchain_text_splitters import RecursiveCharacterTextSplitter
from langchain_community.embeddings import HuggingFaceEmbeddings
from langchain_core.documents import Document
from datasets import load_dataset
from typing import List, Optional, Dict, Any
import csv
import json
from pathlib import Path


# Extract Data From PDF file(s)
def load_pdf_file(data_path: str) -> List[Document]:
    path = Path(data_path)
    if not path.exists():
        raise FileNotFoundError(f"PDF path not found: {data_path}")

    if path.is_file():
        if path.suffix.lower() != ".pdf":
            raise ValueError(
                "When DATA_SOURCE=pdf and path is a file, it must be a .pdf file"
            )
        return PyPDFLoader(str(path)).load()

    loader = DirectoryLoader(str(path), glob="*.pdf", loader_cls=PyPDFLoader)
    return loader.load()


# Load Student Q&A Dataset from Hugging Face
def load_hf_dataset(dataset_name: str) -> List[Document]:
    dataset = load_dataset(dataset_name, split="train")
    documents = []
    for row in dataset:
        # Build a rich text blob from all available fields in the row
        content_parts = []
        for key, value in row.items():
            if value:
                content_parts.append(f"{key.capitalize()}: {value}")
        page_content = "\n".join(content_parts)
        documents.append(
            Document(page_content=page_content, metadata={"source": dataset_name})
        )
    return documents


def _parse_columns(columns: Optional[str]) -> Optional[List[str]]:
    if not columns:
        return None
    parsed = [col.strip() for col in columns.split(",") if col.strip()]
    return parsed or None


def _row_to_document(
    row: Dict[str, Any],
    source_name: str,
    text_columns: Optional[List[str]] = None,
) -> Optional[Document]:
    content_parts: List[str] = []

    if text_columns:
        for key in text_columns:
            value = row.get(key)
            if value is not None and str(value).strip():
                content_parts.append(f"{key.capitalize()}: {value}")
    else:
        for key, value in row.items():
            if value is not None and str(value).strip():
                content_parts.append(f"{key.capitalize()}: {value}")

    if not content_parts:
        return None

    return Document(
        page_content="\n".join(content_parts),
        metadata={"source": source_name},
    )


def load_local_dataset(
    local_path: str, text_columns: Optional[str] = None
) -> List[Document]:
    path = Path(local_path)
    if not path.exists():
        raise FileNotFoundError(f"Dataset file not found: {local_path}")

    parsed_text_columns = _parse_columns(text_columns)
    ext = path.suffix.lower()
    documents: List[Document] = []

    if ext == ".csv":
        with path.open("r", encoding="utf-8") as f:
            reader = csv.DictReader(f)
            for row in reader:
                doc = _row_to_document(
                    row, source_name=str(path.name), text_columns=parsed_text_columns
                )
                if doc:
                    documents.append(doc)
        return documents

    if ext == ".json":
        with path.open("r", encoding="utf-8") as f:
            payload = json.load(f)

        if isinstance(payload, dict):
            payload = [payload]
        if not isinstance(payload, list):
            raise ValueError("JSON dataset must be an object or a list of objects")

        for row in payload:
            if isinstance(row, dict):
                doc = _row_to_document(
                    row, source_name=str(path.name), text_columns=parsed_text_columns
                )
                if doc:
                    documents.append(doc)
        return documents

    if ext == ".jsonl":
        with path.open("r", encoding="utf-8") as f:
            for line in f:
                line = line.strip()
                if not line:
                    continue
                row = json.loads(line)
                if isinstance(row, dict):
                    doc = _row_to_document(
                        row,
                        source_name=str(path.name),
                        text_columns=parsed_text_columns,
                    )
                    if doc:
                        documents.append(doc)
        return documents

    if ext in {".txt", ".md"}:
        text = path.read_text(encoding="utf-8")
        blocks = [block.strip() for block in text.split("\n\n") if block.strip()]
        for block in blocks:
            documents.append(
                Document(page_content=block, metadata={"source": str(path.name)})
            )
        return documents

    raise ValueError(
        "Unsupported dataset format. Use .csv, .json, .jsonl, .txt, or .md"
    )


def load_dataset_by_config(
    data_source: str,
    hf_dataset_name: Optional[str] = None,
    local_dataset_path: Optional[str] = None,
    text_columns: Optional[str] = None,
) -> List[Document]:
    source = (data_source or "hf").strip().lower()

    if source == "hf":
        if not hf_dataset_name:
            raise ValueError("HF_DATASET_NAME is required when DATA_SOURCE=hf")
        return load_hf_dataset(hf_dataset_name)

    if source == "local":
        if not local_dataset_path:
            raise ValueError("LOCAL_DATASET_PATH is required when DATA_SOURCE=local")
        return load_local_dataset(local_dataset_path, text_columns=text_columns)

    if source == "pdf":
        if not local_dataset_path:
            raise ValueError("LOCAL_DATASET_PATH is required when DATA_SOURCE=pdf")
        return load_pdf_file(local_dataset_path)

    raise ValueError("DATA_SOURCE must be one of: hf, local, pdf")


def filter_to_minimal_docs(docs: List[Document]) -> List[Document]:
    """
    Given a list of Document objects, return a new list of Document objects
    containing only 'source' in metadata and the original page_content.
    """
    minimal_docs: List[Document] = []
    for doc in docs:
        src = doc.metadata.get("source")
        minimal_docs.append(
            Document(page_content=doc.page_content, metadata={"source": src})
        )
    return minimal_docs


# Split the Data into Text Chunks
def text_split(extracted_data):
    text_splitter = RecursiveCharacterTextSplitter(chunk_size=500, chunk_overlap=20)
    text_chunks = text_splitter.split_documents(extracted_data)
    return text_chunks


# Download the Embeddings from HuggingFace
def download_hugging_face_embeddings():
    embeddings = HuggingFaceEmbeddings(
        model_name="sentence-transformers/all-MiniLM-L6-v2"
    )  # this model return 384 dimensions
    return embeddings