"""Document loaders for RAG system."""
from __future__ import annotations

import csv
import json
import os
import re
from abc import ABC, abstractmethod
from html.parser import HTMLParser
from typing import Any

from logos_adk.rag import Document, DocumentLoadError


class DocumentLoader(ABC):
    """Abstract base class for document loaders."""

    @abstractmethod
    async def load(self, source: str) -> list[Document]:
        """Load documents from a source path."""
        ...

    @abstractmethod
    def supports(self, source: str) -> bool:
        """Check if this loader supports the given source."""
        ...


class TextLoader(DocumentLoader):
    """Loads plain text files."""

    def supports(self, source: str) -> bool:
        return source.endswith(".txt")

    async def load(self, source: str) -> list[Document]:
        if not os.path.exists(source):
            raise DocumentLoadError(f"File not found: {source}")
        with open(source, "r", encoding="utf-8") as f:
            content = f.read()
        return [Document(content=content, metadata={"source": source})]

    async def load_text(self, text: str, metadata: dict | None = None) -> list[Document]:
        """Load from a raw string."""
        return [Document(content=text, metadata=metadata or {})]


class MarkdownLoader(DocumentLoader):
    """Loads Markdown files with heading extraction."""

    def supports(self, source: str) -> bool:
        return source.endswith(".md")

    async def load(self, source: str) -> list[Document]:
        if not os.path.exists(source):
            raise DocumentLoadError(f"File not found: {source}")
        with open(source, "r", encoding="utf-8") as f:
            content = f.read()
        meta: dict = {"source": source}
        match = re.match(r"^#\s+(.+)$", content, re.MULTILINE)
        if match:
            meta["title"] = match.group(1).strip()
        return [Document(content=content, metadata=meta)]


class PDFLoader(DocumentLoader):
    """Loads PDF files page by page. Requires pypdf."""

    def supports(self, source: str) -> bool:
        return source.endswith(".pdf")

    async def load(self, source: str) -> list[Document]:
        if not os.path.exists(source):
            raise DocumentLoadError(f"File not found: {source}")
        try:
            from pypdf import PdfReader
        except ImportError:
            raise DocumentLoadError(
                "pypdf is required for PDF loading. Install with: pip install logos-adk[rag-pdf]"
            )
        reader = PdfReader(source)
        docs: list[Document] = []
        for i, page in enumerate(reader.pages):
            text = page.extract_text() or ""
            if text.strip():
                docs.append(Document(
                    content=text,
                    metadata={"source": source, "page": i + 1},
                ))
        if not docs:
            docs.append(Document(content="", metadata={"source": source}))
        return docs


class CSVLoader(DocumentLoader):
    """Loads CSV files — each row becomes a Document."""

    def supports(self, source: str) -> bool:
        return source.lower().endswith(".csv")

    async def load(self, source: str) -> list[Document]:
        try:
            with open(source, newline="", encoding="utf-8") as f:
                reader = csv.DictReader(f)
                headers = reader.fieldnames or []
                docs: list[Document] = []
                for i, row in enumerate(reader):
                    content = ", ".join(f"{k}: {v}" for k, v in row.items())
                    docs.append(Document(
                        content=content,
                        metadata={"source": source, "row_index": i, "headers": list(headers)},
                    ))
                return docs
        except FileNotFoundError:
            raise DocumentLoadError(f"CSV file not found: {source}")
        except Exception as exc:
            raise DocumentLoadError(f"Failed to load CSV: {exc}") from exc


class JSONLoader(DocumentLoader):
    """Loads JSON/JSONL files — each object becomes a Document."""

    def __init__(self, content_key: str | None = None) -> None:
        self._content_key = content_key

    def supports(self, source: str) -> bool:
        lower = source.lower()
        return lower.endswith(".json") or lower.endswith(".jsonl")

    async def load(self, source: str) -> list[Document]:
        try:
            with open(source, encoding="utf-8") as f:
                text = f.read()
        except FileNotFoundError:
            raise DocumentLoadError(f"JSON file not found: {source}")

        try:
            if source.lower().endswith(".jsonl"):
                items = [json.loads(line) for line in text.strip().splitlines() if line.strip()]
            else:
                data = json.loads(text)
                items = data if isinstance(data, list) else [data]
        except json.JSONDecodeError as exc:
            raise DocumentLoadError(f"Invalid JSON: {exc}") from exc

        docs: list[Document] = []
        for item in items:
            if self._content_key and isinstance(item, dict):
                content = str(item.get(self._content_key, ""))
                meta = {k: v for k, v in item.items() if k != self._content_key}
            else:
                content = json.dumps(item) if not isinstance(item, str) else item
                meta = {}
            meta["source"] = source
            docs.append(Document(content=content, metadata=meta))
        return docs


class _HTMLStripper(HTMLParser):
    """Simple HTML tag stripper."""

    def __init__(self):
        super().__init__()
        self._parts: list[str] = []
        self._title: str | None = None
        self._in_title = False

    def handle_starttag(self, tag, attrs):
        if tag == "title":
            self._in_title = True

    def handle_endtag(self, tag):
        if tag == "title":
            self._in_title = False

    def handle_data(self, data):
        if self._in_title:
            self._title = data.strip()
        self._parts.append(data)

    def get_text(self) -> str:
        return "".join(self._parts).strip()


class HTMLLoader(DocumentLoader):
    """Loads HTML files — strips tags, extracts title."""

    def supports(self, source: str) -> bool:
        lower = source.lower()
        return lower.endswith(".html") or lower.endswith(".htm")

    async def load(self, source: str) -> list[Document]:
        try:
            with open(source, encoding="utf-8") as f:
                html = f.read()
        except FileNotFoundError:
            raise DocumentLoadError(f"HTML file not found: {source}")

        stripper = _HTMLStripper()
        stripper.feed(html)
        text = stripper.get_text()
        meta: dict[str, Any] = {"source": source}
        if stripper._title:
            meta["title"] = stripper._title
        return [Document(content=text, metadata=meta)]


class URLLoader(DocumentLoader):
    """Loads documents from URLs via httpx."""

    def supports(self, source: str) -> bool:
        return source.startswith("http://") or source.startswith("https://")

    async def load(self, source: str) -> list[Document]:
        import httpx

        try:
            async with httpx.AsyncClient() as client:
                resp = await client.get(source)
                resp.raise_for_status()
        except Exception as exc:
            raise DocumentLoadError(f"Failed to fetch URL {source}: {exc}") from exc

        content_type = resp.headers.get("content-type", "")
        meta: dict[str, Any] = {
            "url": source,
            "status_code": resp.status_code,
            "content_type": content_type,
        }

        if "html" in content_type:
            stripper = _HTMLStripper()
            stripper.feed(resp.text)
            text = stripper.get_text()
            if stripper._title:
                meta["title"] = stripper._title
        else:
            text = resp.text

        return [Document(content=text, metadata=meta)]
