"""
Supports:
- Text files (.txt, .log, .md) read from local temp files
- Structured documents (.pdf, .pptx, .docx, .doc) via DSS DocumentExtractor with LocalFileDocumentRef
"""

import logging
import os
from dataclasses import dataclass
from typing import Optional

from semantic_models.exceptions import ExecutionError

logger = logging.getLogger(__name__)

# Supported file extensions
STRUCTURED_EXTENSIONS = {".pdf", ".pptx", ".docx", ".doc"}
TEXT_EXTENSIONS = {".txt", ".log", ".md"}

# Default character limit for extracted text
MAX_DOC_CHARS = 8000


@dataclass
class DocumentSection:
    """A section of extracted text from a document."""

    title: Optional[str]
    text: str
    index: int


def extract_document_sections(
    file_path: str,
    max_sections: int = 0,
) -> list[DocumentSection]:
    return _extract_sections_from_file(file_path, max_sections)


def extract_document_text(
    file_path: str,
    max_chars: int = MAX_DOC_CHARS,
) -> str:
    sections = extract_document_sections(file_path, max_sections=0)
    texts = [s.text for s in sections]
    combined = "\n\n".join(texts)
    return combined[:max_chars]


def _read_text_file(file_path: str) -> str:
    try:
        with open(file_path, "r", encoding="utf-8", errors="ignore") as f:
            return f.read()
    except Exception as e:
        raise ExecutionError(f"Failed to read text file: {e}")


def _extract_sections_from_file(
    file_path: str,
    max_sections: int,
) -> list[DocumentSection]:
    ext = os.path.splitext(file_path)[1].lower()

    if ext in TEXT_EXTENSIONS:
        text = _read_text_file(file_path)
        return [DocumentSection(title=None, text=text, index=0)]

    if ext in STRUCTURED_EXTENSIONS:
        return _extract_dss_sections(file_path, max_sections)

    raise ExecutionError(
        f"Unsupported file type '{ext}'. Supported: {', '.join(sorted(TEXT_EXTENSIONS | STRUCTURED_EXTENSIONS))}"
    )


def _flatten_content_dfs(content: list, result: list[str]) -> list[str]:
    for item in content:
        if isinstance(item, dict):
            for field in ("text", "title"):
                text = item.get(field, "")
                if text and isinstance(text, str) and text.strip():
                    result.append(text.strip())

            for child_field in ("content", "children"):
                children = item.get(child_field, [])
                if children and isinstance(children, list):
                    _flatten_content_dfs(children, result)
        elif isinstance(item, str) and item.strip():
            result.append(item.strip())
    return result


def _flatten_section_text(item: dict) -> str:
    texts: list[str] = []
    _flatten_content_dfs([item], texts)
    return "\n\n".join(texts)


def _extract_dss_sections(file_path: str, max_sections: int) -> list[DocumentSection]:
    import dataiku
    from dataikuapi.dss.document_extractor import DocumentExtractor, LocalFileDocumentRef

    client = dataiku.api_client()
    project = client.get_default_project()
    extractor = DocumentExtractor(client, project.project_key)

    logger.info("Extracting sections from document: %s", file_path)

    with open(file_path, "rb") as f:
        doc_ref = LocalFileDocumentRef(f)
        response = extractor.structured_extract(
            document=doc_ref,
            image_handling_mode="IGNORE",
        )

    try:
        response._fail_unless_success()
    except Exception as e:
        raise ExecutionError(f"Document extraction failed: {e}")

    sections: list[DocumentSection] = []

    if hasattr(response, "content") and response.content:
        content = response.content
        if isinstance(content, dict):
            content = [content]

        for i, item in enumerate(content):
            if max_sections > 0 and len(sections) >= max_sections:
                break

            if not isinstance(item, dict):
                continue

            title = item.get("title")
            text = _flatten_section_text(item)

            if text.strip():
                sections.append(
                    DocumentSection(
                        title=title if isinstance(title, str) else None,
                        text=text,
                        index=i,
                    )
                )

    elif hasattr(response, "text_chunks") and response.text_chunks:
        for i, item in enumerate(response.text_chunks):
            if max_sections > 0 and len(sections) >= max_sections:
                break

            text = (item.get("text") or "").strip() if isinstance(item, dict) else str(item).strip()
            if text:
                sections.append(DocumentSection(title=None, text=text, index=i))

    logger.info("Extracted %d sections from document: %s", len(sections), file_path)

    if not sections:
        logger.warning("No sections extracted from document: %s", file_path)

    return sections
