"""
Farnsworth Document Processor - Ingestion Pipeline

Handles:
- Text extraction from various formats
- Smart chunking strategies
- Metadata extraction
- Embedding generation
"""

import asyncio
import hashlib
import re
from dataclasses import dataclass, field
from datetime import datetime
from pathlib import Path
from typing import Optional, Any

from loguru import logger

from farnsworth.rag.embeddings import EmbeddingManager
from farnsworth.rag.hybrid_retriever import Document


@dataclass
class ChunkingConfig:
    """Configuration for document chunking."""
    chunk_size: int = 512  # Target chunk size in tokens
    chunk_overlap: int = 50  # Overlap between chunks
    min_chunk_size: int = 100  # Minimum chunk size
    split_on: list[str] = field(default_factory=lambda: ["\n\n", "\n", ". ", " "])
    preserve_sentences: bool = True


@dataclass
class ProcessedDocument:
    """A processed document with chunks."""
    source_id: str
    source_path: Optional[str]
    chunks: list[Document]
    metadata: dict = field(default_factory=dict)
    processed_at: datetime = field(default_factory=datetime.now)


class DocumentProcessor:
    """
    Document processing pipeline for RAG.

    Features:
    - Multiple format support (txt, md, pdf, code)
    - Smart semantic chunking
    - Overlap handling
    - Metadata extraction
    """

    def __init__(
        self,
        embedding_manager: Optional[EmbeddingManager] = None,
        chunking_config: Optional[ChunkingConfig] = None,
    ):
        self.embedding_manager = embedding_manager
        self.config = chunking_config or ChunkingConfig()

        # Format handlers
        self._handlers = {
            ".txt": self._process_text,
            ".md": self._process_markdown,
            ".py": self._process_code,
            ".js": self._process_code,
            ".ts": self._process_code,
            ".json": self._process_json,
            ".yaml": self._process_yaml,
            ".yml": self._process_yaml,
        }

    async def process_file(
        self,
        file_path: str,
        metadata: Optional[dict] = None,
    ) -> ProcessedDocument:
        """Process a file and return chunks."""
        path = Path(file_path)

        if not path.exists():
            raise FileNotFoundError(f"File not found: {file_path}")

        # Get handler
        suffix = path.suffix.lower()
        handler = self._handlers.get(suffix, self._process_text)

        # Read and process
        content = path.read_text(encoding='utf-8', errors='ignore')
        processed_content, extracted_metadata = handler(content, path)

        # Merge metadata
        final_metadata = {
            "source_path": str(path),
            "filename": path.name,
            "extension": suffix,
            "size_bytes": path.stat().st_size,
            "modified_at": datetime.fromtimestamp(path.stat().st_mtime).isoformat(),
            **(extracted_metadata or {}),
            **(metadata or {}),
        }

        # Chunk the content
        chunks = self._chunk_text(processed_content, str(path), final_metadata)

        # Generate embeddings
        if self.embedding_manager:
            chunks = await self._add_embeddings(chunks)

        return ProcessedDocument(
            source_id=hashlib.md5(str(path).encode()).hexdigest()[:12],
            source_path=str(path),
            chunks=chunks,
            metadata=final_metadata,
        )

    async def process_text(
        self,
        text: str,
        source_id: Optional[str] = None,
        metadata: Optional[dict] = None,
    ) -> ProcessedDocument:
        """Process raw text content."""
        source_id = source_id or hashlib.md5(text[:100].encode()).hexdigest()[:12]

        chunks = self._chunk_text(text, source_id, metadata or {})

        if self.embedding_manager:
            chunks = await self._add_embeddings(chunks)

        return ProcessedDocument(
            source_id=source_id,
            source_path=None,
            chunks=chunks,
            metadata=metadata or {},
        )

    def _chunk_text(
        self,
        text: str,
        source_id: str,
        base_metadata: dict,
    ) -> list[Document]:
        """Chunk text into overlapping segments."""
        if not text.strip():
            return []

        # Estimate tokens (rough: 4 chars per token)
        char_chunk_size = self.config.chunk_size * 4
        char_overlap = self.config.chunk_overlap * 4
        char_min_size = self.config.min_chunk_size * 4

        chunks = []
        current_pos = 0
        chunk_idx = 0

        while current_pos < len(text):
            # Find chunk end
            chunk_end = min(current_pos + char_chunk_size, len(text))

            # Try to split at a natural boundary
            if chunk_end < len(text):
                chunk_end = self._find_split_point(
                    text, current_pos, chunk_end
                )

            chunk_text = text[current_pos:chunk_end].strip()

            if len(chunk_text) >= char_min_size:
                doc = Document(
                    id=f"{source_id}_chunk_{chunk_idx}",
                    content=chunk_text,
                    metadata={
                        **base_metadata,
                        "chunk_index": chunk_idx,
                        "char_start": current_pos,
                        "char_end": chunk_end,
                    },
                )
                chunks.append(doc)
                chunk_idx += 1

            # Move position with overlap
            current_pos = chunk_end - char_overlap
            if current_pos <= 0 or current_pos >= len(text) - char_min_size:
                current_pos = chunk_end

        return chunks

    def _find_split_point(
        self,
        text: str,
        start: int,
        end: int,
    ) -> int:
        """Find a good split point in text."""
        search_text = text[start:end]

        # Try split points in order of preference
        for delimiter in self.config.split_on:
            last_pos = search_text.rfind(delimiter)
            if last_pos > len(search_text) * 0.3:  # At least 30% through
                return start + last_pos + len(delimiter)

        return end

    async def _add_embeddings(self, chunks: list[Document]) -> list[Document]:
        """Add embeddings to chunks."""
        texts = [c.content for c in chunks]
        results = await self.embedding_manager.embed_batch(texts)

        for chunk, result in zip(chunks, results):
            chunk.embedding = result.embedding

        return chunks

    def _process_text(self, content: str, path: Path) -> tuple[str, dict]:
        """Process plain text."""
        return content, {}

    def _process_markdown(self, content: str, path: Path) -> tuple[str, dict]:
        """Process markdown with structure extraction."""
        metadata = {}

        # Extract title (first H1)
        title_match = re.search(r'^#\s+(.+)$', content, re.MULTILINE)
        if title_match:
            metadata["title"] = title_match.group(1)

        # Extract headers for structure
        headers = re.findall(r'^(#{1,6})\s+(.+)$', content, re.MULTILINE)
        metadata["headers"] = [h[1] for h in headers]

        # Remove code blocks for cleaner chunking (preserve as separate chunks)
        code_blocks = re.findall(r'```[\w]*\n(.*?)\n```', content, re.DOTALL)
        metadata["has_code"] = len(code_blocks) > 0

        return content, metadata

    def _process_code(self, content: str, path: Path) -> tuple[str, dict]:
        """Process code files with structure extraction."""
        metadata = {
            "language": path.suffix[1:],
        }

        # Extract function/class definitions
        if path.suffix == ".py":
            functions = re.findall(r'^def\s+(\w+)', content, re.MULTILINE)
            classes = re.findall(r'^class\s+(\w+)', content, re.MULTILINE)
            metadata["functions"] = functions
            metadata["classes"] = classes
        elif path.suffix in (".js", ".ts"):
            functions = re.findall(r'(?:function|const|let|var)\s+(\w+)\s*[=\(]', content)
            classes = re.findall(r'class\s+(\w+)', content)
            metadata["functions"] = functions
            metadata["classes"] = classes

        # Add docstrings/comments as context
        docstrings = re.findall(r'"""(.*?)"""', content, re.DOTALL)
        metadata["has_docstrings"] = len(docstrings) > 0

        return content, metadata

    def _process_json(self, content: str, path: Path) -> tuple[str, dict]:
        """Process JSON files."""
        import json
        metadata = {"format": "json"}

        try:
            data = json.loads(content)
            if isinstance(data, dict):
                metadata["keys"] = list(data.keys())[:10]
            elif isinstance(data, list):
                metadata["item_count"] = len(data)
        except json.JSONDecodeError:
            pass

        return content, metadata

    def _process_yaml(self, content: str, path: Path) -> tuple[str, dict]:
        """Process YAML files."""
        import yaml
        metadata = {"format": "yaml"}

        try:
            data = yaml.safe_load(content)
            if isinstance(data, dict):
                metadata["keys"] = list(data.keys())[:10]
        except yaml.YAMLError:
            pass

        return content, metadata

    async def process_directory(
        self,
        dir_path: str,
        recursive: bool = True,
        patterns: Optional[list[str]] = None,
    ) -> list[ProcessedDocument]:
        """Process all files in a directory."""
        path = Path(dir_path)
        if not path.is_dir():
            raise NotADirectoryError(f"Not a directory: {dir_path}")

        patterns = patterns or ["*.txt", "*.md", "*.py", "*.js"]
        results = []

        for pattern in patterns:
            if recursive:
                files = path.rglob(pattern)
            else:
                files = path.glob(pattern)

            for file_path in files:
                if file_path.is_file():
                    try:
                        doc = await self.process_file(str(file_path))
                        results.append(doc)
                    except Exception as e:
                        logger.error(f"Failed to process {file_path}: {e}")

        return results
