import re
from pathlib import Path
from typing import Any, Dict, List

from schemas import (
    AITutorDocument,
    Activity,
    DocumentMetadata,
    Figure,
    Question,
    Section,
    SourceReference,
    Table,
    TextBlock,
)


class DocumentNormalizer:

    def normalize(
        self,
        docling_document: Any,
        filename: str,
        raw_json: Dict[str, Any],
    ) -> AITutorDocument:

        document_id = Path(filename).stem

        raw = raw_json

        pages = raw.get("pages", {})
        texts = raw.get("texts", [])
        pictures = raw.get("pictures", [])
        tables = raw.get("tables", [])

        markdown = docling_document.export_to_markdown()

        title = self._detect_title(markdown)

        metadata = DocumentMetadata(
            id=document_id,
            filename=filename,
            format=Path(filename).suffix.lower().replace(".", ""),
            title=title,
            language="English",
            page_count=len(pages),
        )

        sections = self._extract_sections(
            markdown,
            document_id,
        )

        text_blocks = self._extract_text_blocks(
            texts,
            document_id,
        )

        normalized_tables = self._extract_tables(
            docling_document,
            tables,
            document_id,
        )

        figures = self._extract_figures(
            pictures,
            document_id,
        )

        activities = self._extract_activities(
            markdown,
            document_id,
        )

        questions = self._extract_questions(
            markdown,
            document_id,
        )

        statistics = {
            "pages": len(pages),
            "text_elements": len(texts),
            "pictures": len(pictures),
            "tables": len(tables),
            "sections": len(sections),
            "activities": len(activities),
            "questions": len(questions),
        }

        return AITutorDocument(
            document=metadata,
            sections=sections,
            text_blocks=text_blocks,
            tables=normalized_tables,
            figures=figures,
            activities=activities,
            questions=questions,
            statistics=statistics,
            raw_docling=raw_json,
        )

    # --------------------------------------------------
    # TITLE
    # --------------------------------------------------

    def _detect_title(self, markdown: str) -> str:

        lines = [
            line.strip()
            for line in markdown.splitlines()
            if line.strip()
        ]

        for line in lines[:20]:

            clean = re.sub(r"<!--.*?-->", "", line).strip()

            if len(clean) > 10:
                return clean

        return ""

    # --------------------------------------------------
    # SECTIONS
    # --------------------------------------------------

    def _extract_sections(
        self,
        markdown: str,
        document_id: str,
    ) -> List[Section]:

        sections = []

        lines = markdown.splitlines()

        section_pattern = re.compile(
            r"^#{1,6}\s+(.+)$"
        )

        current_id = None

        for index, line in enumerate(lines):

            match = section_pattern.match(line.strip())

            if not match:
                continue

            title = match.group(1).strip()

            section_match = re.match(
                r"^(\d+(?:\.\d+)*)\s+(.+)",
                title,
            )

            if section_match:

                section_number = section_match.group(1)

                section_title = section_match.group(2)

                level = section_number.count(".") + 1

                section_id = (
                    f"{document_id}-section-{section_number}"
                )

                section = Section(
                    id=section_id,
                    title=section_title,
                    level=level,
                )

                sections.append(section)

                current_id = section_id

        return sections

    # --------------------------------------------------
    # TEXT
    # --------------------------------------------------

    def _extract_text_blocks(
        self,
        texts: List[Dict[str, Any]],
        document_id: str,
    ) -> List[TextBlock]:

        result = []

        for index, item in enumerate(texts):

            text = item.get("text", "")

            if not text:
                continue

            text_type = (
                item.get("label")
                or item.get("type")
                or "paragraph"
            )

            result.append(
                TextBlock(
                    id=f"{document_id}-text-{index + 1}",
                    text=text,
                    type=str(text_type),
                )
            )

        return result

    # --------------------------------------------------
    # TABLES
    # --------------------------------------------------

    def _extract_tables(
        self,
        docling_document: Any,
        tables: List[Dict[str, Any]],
        document_id: str,
    ) -> List[Table]:

        result = []

        for index, table in enumerate(tables):

            markdown = None

            try:
                table_obj = None

                table_ref = table.get("self_ref")

                if table_ref:
                    table_obj = docling_document.get_ref(
                        table_ref
                    )

                if table_obj:
                    markdown = table_obj.export_to_markdown()

            except Exception:
                markdown = None

            result.append(
                Table(
                    id=f"{document_id}-table-{index + 1}",
                    markdown=markdown,
                )
            )

        return result

    # --------------------------------------------------
    # FIGURES
    # --------------------------------------------------

    def _extract_figures(
        self,
        pictures: List[Dict[str, Any]],
        document_id: str,
    ) -> List[Figure]:

        result = []

        for index, picture in enumerate(pictures):

            caption = picture.get("caption")

            if isinstance(caption, dict):
                caption = caption.get("text")

            result.append(
                Figure(
                    id=f"{document_id}-figure-{index + 1}",
                    caption=caption,
                )
            )

        return result

    # --------------------------------------------------
    # ACTIVITIES
    # --------------------------------------------------

    def _extract_activities(
        self,
        markdown: str,
        document_id: str,
    ) -> List[Activity]:

        result = []

        pattern = re.compile(
            r"(?:^|\n)##?\s*(Activity\s+\d+(?:\.\d+)?)"
            r"(?::\s*([^\n]+))?",
            re.IGNORECASE,
        )

        matches = list(pattern.finditer(markdown))

        for index, match in enumerate(matches):

            activity_number = match.group(1)
            title = match.group(2) or "Activity"

            start = match.end()

            if index + 1 < len(matches):
                end = matches[index + 1].start()
            else:
                end = len(markdown)

            content = markdown[start:end].strip()

            result.append(
                Activity(
                    id=f"{document_id}-{activity_number.lower().replace(' ', '-')}",
                    title=f"{activity_number}: {title}",
                    content=[content[:5000]],
                )
            )

        return result

    # --------------------------------------------------
    # QUESTIONS
    # --------------------------------------------------

    def _extract_questions(
        self,
        markdown: str,
        document_id: str,
    ) -> List[Question]:

        result = []

        lines = markdown.splitlines()

        question_number = 0

        for line in lines:

            text = line.strip()

            if not text:
                continue

            if re.match(
                r"^\d+[\.\)]\s+.+\?$",
                text,
            ):

                question_number += 1

                result.append(
                    Question(
                        id=f"{document_id}-question-{question_number}",
                        question=re.sub(
                            r"^\d+[\.\)]\s*",
                            "",
                            text,
                        ),
                    )
                )

        return result