"""PDF parser using OpenRouter LLM integration for image-to-text transcription."""

import asyncio
import base64
import io
import logging
import os
import time
from pathlib import Path
from typing import Any, Dict, List, Optional, Tuple

try:
    from jinja2 import Environment, FileSystemLoader, TemplateNotFound

    JINJA2_AVAILABLE = True
except ImportError:
    JINJA2_AVAILABLE = False

from .parser import ParseResult, PDFParser

logger = logging.getLogger(__name__)


class LLMParser(PDFParser):
    """PDF parser using OpenRouter LLM integration for image-to-text transcription."""

    def __init__(self, config: Optional[Dict[str, Any]] = None, cache_dir: Path = None):
        """Initialize the LLM parser.

        Args:
            config: Configuration options for LLM parsing.
            cache_dir: Directory to cache parsed markdown files.
        """
        super().__init__(cache_dir)
        self.config = config or {}

        # Configuration with defaults
        self.model = self.config.get("model", "google/gemini-2.5-flash")
        self.openrouter_api_key = self.config.get("openrouter_api_key") or os.getenv("OPENROUTER_API_KEY")
        self.concurrency = self.config.get("concurrency", 5)
        self.dpi = self.config.get("dpi", 150)
        self.max_retries = self.config.get("max_retries", 3)
        self.base_url = "https://openrouter.ai/api/v1/chat/completions"

        # Validate configuration
        if not self.openrouter_api_key:
            raise ValueError(
                "OpenRouter API key is required. Set OPENROUTER_API_KEY environment variable or pass it in config."
            )

        # Setup Jinja2 environment
        self._setup_template_environment()

        # Import dependencies
        try:
            import fitz  # PyMuPDF
            import httpx
            from PIL import Image

            self.httpx = httpx
            self.fitz = fitz
            self.Image = Image
        except ImportError as e:
            raise ImportError(
                f"Required dependencies not available: {e}. Install with: pip install httpx pymupdf pillow"
            )

    def _setup_template_environment(self):
        """Setup Jinja2 template environment."""
        if not JINJA2_AVAILABLE:
            logger.warning("Jinja2 not available, using hardcoded prompts. Install with: pip install jinja2")
            self.jinja_env = None
            return

        try:
            # Templates directory relative to this file
            templates_dir = Path(__file__).parent / "templates"

            if not templates_dir.exists():
                logger.warning(f"Templates directory not found: {templates_dir}. Using hardcoded prompts.")
                self.jinja_env = None
                return

            # Setup Jinja2 environment
            self.jinja_env = Environment(
                loader=FileSystemLoader(str(templates_dir)),
                autoescape=False,  # We're generating plain text prompts
                trim_blocks=True,
                lstrip_blocks=True,
            )

            logger.debug(f"Jinja2 environment initialized with templates from: {templates_dir}")

        except Exception as e:
            logger.warning(f"Failed to setup Jinja2 environment: {e}. Using hardcoded prompts.")
            self.jinja_env = None

    def _render_template(self, template_name: str, **kwargs) -> str:
        """Render a Jinja2 template with the given variables.

        Args:
            template_name: Name of the template file.
            **kwargs: Variables to pass to the template.

        Returns:
            Rendered template string.
        """
        if not self.jinja_env:
            # Fallback to hardcoded prompts if Jinja2 not available
            return self._get_fallback_prompt(template_name, **kwargs)

        try:
            template = self.jinja_env.get_template(template_name)
            return template.render(**kwargs)
        except TemplateNotFound:
            logger.warning(f"Template {template_name} not found, using fallback prompt")
            return self._get_fallback_prompt(template_name, **kwargs)
        except Exception as e:
            logger.warning(f"Failed to render template {template_name}: {e}. Using fallback prompt.")
            return self._get_fallback_prompt(template_name, **kwargs)

    def _get_fallback_prompt(self, template_name: str, **kwargs) -> str:
        """Get fallback hardcoded prompts when templates are not available.

        Args:
            template_name: Name of the template.
            **kwargs: Template variables (ignored for fallback).

        Returns:
            Hardcoded prompt string.
        """
        if template_name == "page_transcription.j2":
            return """Convert this PDF page image to markdown format. Follow these guidelines:

1. Maintain accurate markdown syntax and structure
2. Preserve document hierarchy (headers, lists, tables)
3. Replace images with descriptive text in brackets like [Image: description]
4. Ensure 1:1 transcription accuracy (do not summarize)
5. Maintain proper formatting for tables, lists, and text blocks
6. Use appropriate header levels (# ## ###) based on visual hierarchy
7. Preserve any mathematical formulas or equations in text format
8. Include any footnotes or references found on the page

Return only the markdown content, no additional commentary."""

        elif template_name == "metadata_extraction.j2":
            return """Analyze these first few pages of a PDF document and extract:

1. **Title**: The main title of the document (not headers/footers)
2. **Summary**: A concise 2-3 sentence summary of what this document is about

Look for:
- Main document title (often largest text on first page)
- Abstract, introduction, or executive summary sections
- Key topics and themes

Return your response in this exact format:
TITLE: [extracted title]
SUMMARY: [2-3 sentence summary]

If you cannot determine a clear title or summary, return "TITLE: " or "SUMMARY: " with no content after the colon."""

        else:
            logger.warning(f"Unknown template name for fallback: {template_name}")
            return "Please analyze the provided content."

    async def parse(self, file_path: Path) -> ParseResult:
        """Parse a PDF file using LLM image-to-text transcription.

        Args:
            file_path: Path to the PDF file.

        Returns:
            ParseResult with markdown content and metadata.
        """
        try:
            # Check cache first
            cache_path = None
            if self.cache_dir:
                cache_path = self._get_cache_path(file_path)
                if self._is_cache_valid(file_path, cache_path):
                    logger.debug(f"Loading parsed content from cache: {cache_path}")
                    markdown_content = self._load_from_cache(cache_path)
                    metadata = self._load_metadata_from_cache(cache_path)
                    if markdown_content is not None and metadata:
                        return ParseResult(markdown_content=markdown_content, metadata=metadata)

            logger.debug(f"Parsing PDF with LLM: {file_path}")
            start_time = time.time()

            # Open PDF and extract pages
            doc = self.fitz.open(str(file_path))
            total_pages = doc.page_count
            logger.debug(f"PDF has {total_pages} pages")

            # Convert pages to images
            page_images = await self._convert_pages_to_images(doc)
            doc.close()

            # Process pages concurrently
            page_transcriptions = await self._process_pages_concurrently(page_images)

            # Combine transcriptions into final markdown
            markdown_content = self._combine_page_transcriptions(page_transcriptions)

            # Extract metadata from first 5 pages
            metadata = await self._extract_metadata(file_path, page_images[:5], total_pages)

            # Add processing information
            processing_time = time.time() - start_time
            metadata.update(
                {
                    "processing_timestamp": "N/A",  # Will be set by PDFProcessor
                    "processor_version": "llm",
                    "source_filename": file_path.name,
                    "source_directory": str(file_path.parent),
                    "processing_time_seconds": processing_time,
                    "llm_model": self.model,
                    "page_count": total_pages,
                }
            )

            # Save to cache if enabled
            if cache_path:
                logger.debug(f"Saving parsed content to cache: {cache_path}")
                self._save_to_cache(cache_path, markdown_content)
                self._save_metadata_to_cache(cache_path, metadata)

            logger.debug(f"Successfully parsed PDF with LLM in {processing_time:.2f} seconds")
            return ParseResult(markdown_content=markdown_content, metadata=metadata)

        except Exception as e:
            logger.error(f"Failed to parse PDF with LLM: {e}")
            raise RuntimeError(f"Failed to parse PDF with LLM: {e}") from e

    async def _convert_pages_to_images(self, doc) -> List[Tuple[int, str]]:
        """Convert PDF pages to PNG images.

        Args:
            doc: PyMuPDF document object.

        Returns:
            List of tuples (page_number, base64_encoded_image).
        """
        page_images = []

        for page_num in range(doc.page_count):
            try:
                page = doc[page_num]

                # Convert page to image (150 DPI, RGB)
                mat = self.fitz.Matrix(self.dpi / 72, self.dpi / 72)  # Scale factor for DPI
                pix = page.get_pixmap(matrix=mat, colorspace=self.fitz.csRGB)

                # Convert to PIL Image
                img_data = pix.tobytes("png")
                img = self.Image.open(io.BytesIO(img_data))

                # Convert to base64
                buffered = io.BytesIO()
                img.save(buffered, format="PNG")
                img_base64 = base64.b64encode(buffered.getvalue()).decode("utf-8")

                page_images.append((page_num + 1, img_base64))
                logger.debug(f"Converted page {page_num + 1} to image")

            except Exception as e:
                logger.warning(f"Failed to convert page {page_num + 1} to image: {e}")
                # Add placeholder for failed page
                page_images.append((page_num + 1, None))

        return page_images

    async def _process_pages_concurrently(self, page_images: List[Tuple[int, str]]) -> List[Tuple[int, str]]:
        """Process pages concurrently with configurable concurrency.

        Args:
            page_images: List of tuples (page_number, base64_encoded_image).

        Returns:
            List of tuples (page_number, transcribed_markdown).
        """
        semaphore = asyncio.Semaphore(self.concurrency)

        async def process_single_page(page_data: Tuple[int, str]) -> Tuple[int, str]:
            async with semaphore:
                page_num, img_base64 = page_data
                if img_base64 is None:
                    return (
                        page_num,
                        f"# Page {page_num}\n\n*[Error: Could not process this page]*\n",
                    )

                try:
                    transcription = await self._transcribe_image_to_markdown(img_base64, page_num)
                    return (page_num, transcription)
                except Exception as e:
                    logger.warning(f"Failed to transcribe page {page_num}: {e}")
                    return (
                        page_num,
                        f"# Page {page_num}\n\n*[Error: Could not transcribe this page: {str(e)}]*\n",
                    )

        # Process all pages concurrently
        tasks = [process_single_page(page_data) for page_data in page_images]
        results = await asyncio.gather(*tasks, return_exceptions=True)

        # Handle any exceptions and sort by page number
        transcriptions = []
        for result in results:
            if isinstance(result, Exception):
                logger.error(f"Page processing failed: {result}")
                transcriptions.append((0, "*[Error: Page processing failed]*\n"))
            else:
                transcriptions.append(result)

        # Sort by page number
        transcriptions.sort(key=lambda x: x[0])
        return transcriptions

    async def _transcribe_image_to_markdown(self, img_base64: str, page_num: int) -> str:
        """Transcribe a single image to markdown using OpenRouter API.

        Args:
            img_base64: Base64 encoded image.
            page_num: Page number for context.

        Returns:
            Transcribed markdown content.
        """
        # Render prompt using template
        prompt = self._render_template(
            "page_transcription.j2",
            page_number=page_num,
            image_description_format="description",
            strict_transcription=True,
            preserve_formatting=True,
        )

        headers = {
            "Authorization": f"Bearer {self.openrouter_api_key}",
            "Content-Type": "application/json",
        }

        payload = {
            "model": self.model,
            "messages": [
                {
                    "role": "user",
                    "content": [
                        {"type": "text", "text": prompt},
                        {
                            "type": "image_url",
                            "image_url": {"url": f"data:image/png;base64,{img_base64}"},
                        },
                    ],
                }
            ],
            "max_tokens": 4000,
            "temperature": 0.1,  # Low temperature for consistent transcription
        }

        # Retry logic
        last_exception = None
        for attempt in range(self.max_retries):
            try:
                async with self.httpx.AsyncClient(timeout=60.0) as client:
                    response = await client.post(self.base_url, json=payload, headers=headers)

                    if response.status_code == 200:
                        result = response.json()
                        content = result["choices"][0]["message"]["content"]
                        return f"# Page {page_num}\n\n{content}\n"

                    elif response.status_code == 429:  # Rate limit
                        wait_time = 2**attempt  # Exponential backoff
                        logger.warning(
                            f"Rate limited on page {page_num}, waiting {wait_time}s "
                            f"(attempt {attempt + 1}/{self.max_retries})"
                        )
                        await asyncio.sleep(wait_time)
                        continue

                    elif response.status_code == 400:
                        error_data = response.json()
                        error_msg = error_data.get("error", {}).get("message", "Unknown API error")
                        raise RuntimeError(f"API error: {error_msg}")

                    else:
                        response.raise_for_status()

            except Exception as e:
                last_exception = e
                if attempt < self.max_retries - 1:
                    wait_time = 2**attempt
                    logger.warning(
                        f"API call failed for page {page_num}, retrying in {wait_time}s "
                        f"(attempt {attempt + 1}/{self.max_retries}): {e}"
                    )
                    await asyncio.sleep(wait_time)
                else:
                    logger.error(f"Failed to transcribe page {page_num} after {self.max_retries} attempts: {e}")

        # If all retries failed, return error placeholder
        raise RuntimeError(f"Failed to transcribe page {page_num} after {self.max_retries} attempts: {last_exception}")

    def _combine_page_transcriptions(self, page_transcriptions: List[Tuple[int, str]]) -> str:
        """Combine page transcriptions into final markdown document.

        Args:
            page_transcriptions: List of tuples (page_number, transcribed_markdown).

        Returns:
            Combined markdown content.
        """
        combined_content = []

        for page_num, transcription in page_transcriptions:
            if transcription and transcription.strip():
                combined_content.append(transcription.strip())

        return "\n\n---\n\n".join(combined_content)

    async def _extract_metadata(
        self, file_path: Path, first_pages: List[Tuple[int, str]], total_pages: int
    ) -> Dict[str, Any]:
        """Extract title and summary from first 5 pages using LLM.

        Args:
            file_path: Path to the PDF file.
            first_pages: First 5 pages as (page_number, base64_image) tuples.
            total_pages: Total number of pages in the document.

        Returns:
            Dictionary of extracted metadata.
        """
        metadata = {
            "page_count": total_pages,
            "title": "",
            "summary": "",
        }

        # Try to extract basic metadata using PyMuPDF
        try:
            doc = self.fitz.open(str(file_path))
            doc_metadata = doc.metadata
            if doc_metadata:
                metadata.update(
                    {
                        "author": doc_metadata.get("author", ""),
                        "subject": doc_metadata.get("subject", ""),
                        "creator": doc_metadata.get("creator", ""),
                        "producer": doc_metadata.get("producer", ""),
                        "creationDate": doc_metadata.get("creationDate", ""),
                        "modDate": doc_metadata.get("modDate", ""),
                    }
                )
                # Use PDF title if available
                if doc_metadata.get("title"):
                    metadata["title"] = doc_metadata["title"]
            doc.close()
        except Exception as e:
            logger.warning(f"Failed to extract basic metadata: {e}")

        # Extract title and summary using LLM from first pages
        try:
            # Combine first few pages for context
            pages_for_metadata = [page for page in first_pages if page[1] is not None][:5]

            if pages_for_metadata:
                title, summary = await self._extract_title_and_summary(pages_for_metadata)
                if title:
                    metadata["title"] = title
                if summary:
                    metadata["summary"] = summary

        except Exception as e:
            logger.warning(f"Failed to extract title and summary with LLM: {e}")

        # Fallback title from filename if still empty
        if not metadata["title"]:
            metadata["title"] = file_path.stem.replace("_", " ").replace("-", " ").title()

        return metadata

    async def _extract_title_and_summary(self, pages: List[Tuple[int, str]]) -> Tuple[str, str]:
        """Extract title and summary from first pages using LLM.

        Args:
            pages: List of (page_number, base64_image) tuples.

        Returns:
            Tuple of (title, summary).
        """
        if not pages:
            return "", ""

        # Render prompt using template
        prompt = self._render_template(
            "metadata_extraction.j2",
            summary_length="2-3",
            page_count=len(pages),
            focus_areas=["Document structure and organization", "Main themes and objectives"],
            fallback_instructions=(
                "If the document appears to be academic, look for author names and publication details."
            ),
        )

        # Create message with multiple images
        message_content = [{"type": "text", "text": prompt}]

        for page_num, img_base64 in pages:
            message_content.append({"type": "image_url", "image_url": {"url": f"data:image/png;base64,{img_base64}"}})

        headers = {
            "Authorization": f"Bearer {self.openrouter_api_key}",
            "Content-Type": "application/json",
        }

        payload = {
            "model": self.model,
            "messages": [{"role": "user", "content": message_content}],
            "max_tokens": 500,
            "temperature": 0.1,
        }

        try:
            async with self.httpx.AsyncClient(timeout=60.0) as client:
                response = await client.post(self.base_url, json=payload, headers=headers)

                if response.status_code == 200:
                    result = response.json()
                    content = result["choices"][0]["message"]["content"]

                    # Parse title and summary
                    title = ""
                    summary = ""

                    for line in content.split("\n"):
                        line = line.strip()
                        if line.startswith("TITLE:"):
                            title = line[6:].strip()
                        elif line.startswith("SUMMARY:"):
                            summary = line[8:].strip()

                    return title, summary
                else:
                    logger.warning(f"Failed to extract metadata, status: {response.status_code}")
                    return "", ""

        except Exception as e:
            logger.warning(f"Failed to extract title and summary: {e}")
            return "", ""
