import logging
import os
import tempfile
import time
import zipfile
from typing import List, Optional

import requests
from fastapi import HTTPException, status
from langchain_core.documents import Document

log = logging.getLogger(__name__)


class MinerULoader:
    """
    MinerU document parser loader supporting both Cloud API and Local API modes.

    Cloud API: Uses MinerU managed service with async task-based processing
    Local API: Uses self-hosted MinerU API with synchronous processing
    """

    def __init__(
        self,
        file_path: str,
        api_mode: str = 'local',
        api_url: str = 'http://localhost:8000',
        api_key: str = '',
        params: dict = None,
        timeout: Optional[int] = 300,
        max_markdown_bytes: Optional[int] = None,
    ):
        self.file_path = file_path
        self.api_mode = api_mode.lower()
        self.api_url = api_url.rstrip('/')
        self.api_key = api_key
        self.timeout = timeout
        self.max_markdown_bytes = max_markdown_bytes

        # Parse params dict with defaults
        self.params = params or {}
        self.enable_ocr = self.params.get('enable_ocr', False)
        self.enable_formula = self.params.get('enable_formula', True)
        self.enable_table = self.params.get('enable_table', True)
        self.language = self.params.get('language', 'en')
        self.model_version = self.params.get('model_version', 'pipeline')

        self.page_ranges = self.params.pop('page_ranges', '')

        # Validate API mode
        if self.api_mode not in ['local', 'cloud']:
            raise ValueError(f"Invalid API mode: {self.api_mode}. Must be 'local' or 'cloud'")

        # Validate Cloud API requirements
        if self.api_mode == 'cloud' and not self.api_key:
            raise ValueError('API key is required for Cloud API mode')

    def load(self) -> List[Document]:
        """
        Main entry point for loading and parsing the document.
        Routes to Cloud or Local API based on api_mode.
        """
        try:
            if self.api_mode == 'cloud':
                return self._load_cloud_api()
            else:
                return self._load_local_api()
        except Exception as e:
            log.error(f'Error loading document with MinerU: {e}')
            raise

    def _load_local_api(self) -> List[Document]:
        """
        Load document using Local API (synchronous).
        Posts file to /file_parse endpoint and gets immediate response.
        """
        log.info('Using MinerU Local API at %s', self.api_url)

        filename = os.path.basename(self.file_path)

        # Build form data for Local API
        form_data = {
            **self.params,
            'return_md': 'true',
        }

        # Page ranges (Local API uses start_page_id and end_page_id)
        if self.page_ranges:
            # For simplicity, if page_ranges is specified, log a warning
            # Full page range parsing would require parsing the string
            log.warning(
                f"Page ranges '{self.page_ranges}' specified but Local API uses different format. "
                'Consider using start_page_id/end_page_id parameters if needed.'
            )

        try:
            with open(self.file_path, 'rb') as f:
                files = {'files': (filename, f, 'application/octet-stream')}

                log.info('Sending file to MinerU Local API: %s', filename)
                log.debug('Local API parameters: %s', form_data)

                response = requests.post(
                    f'{self.api_url}/file_parse',
                    data=form_data,
                    files=files,
                    timeout=self.timeout,
                )
                response.raise_for_status()

        except FileNotFoundError:
            raise HTTPException(status.HTTP_404_NOT_FOUND, detail=f'File not found: {self.file_path}')
        except requests.Timeout:
            raise HTTPException(
                status.HTTP_504_GATEWAY_TIMEOUT,
                detail='MinerU Local API request timed out',
            )
        except requests.HTTPError as e:
            error_detail = f'MinerU Local API request failed: {e}'
            if e.response is not None:
                try:
                    error_data = e.response.json()
                    error_detail += f' - {error_data}'
                except Exception:
                    error_detail += f' - {e.response.text}'
            raise HTTPException(status.HTTP_400_BAD_REQUEST, detail=error_detail)
        except Exception as e:
            raise HTTPException(
                status.HTTP_500_INTERNAL_SERVER_ERROR,
                detail=f'Error calling MinerU Local API: {str(e)}',
            )

        # Parse response
        try:
            result = response.json()
        except ValueError as e:
            raise HTTPException(
                status.HTTP_502_BAD_GATEWAY,
                detail=f'Invalid JSON response from MinerU Local API: {e}',
            )

        # Extract markdown content from response
        if 'results' not in result:
            raise HTTPException(
                status.HTTP_502_BAD_GATEWAY,
                detail="MinerU Local API response missing 'results' field",
            )

        results = result['results']
        if not results:
            raise HTTPException(
                status.HTTP_400_BAD_REQUEST,
                detail='MinerU returned empty results',
            )

        # Get the first (and typically only) result
        file_result = list(results.values())[0]
        markdown_content = file_result.get('md_content', '')

        if not markdown_content:
            raise HTTPException(
                status.HTTP_400_BAD_REQUEST,
                detail='MinerU returned empty markdown content',
            )

        log.info('Successfully parsed document with MinerU Local API: %s', filename)

        # Create metadata
        metadata = {
            'source': filename,
            'api_mode': 'local',
            'backend': result.get('backend', 'unknown'),
            'version': result.get('version', 'unknown'),
        }

        return [Document(page_content=markdown_content, metadata=metadata)]

    def _load_cloud_api(self) -> List[Document]:
        """
        Load document using Cloud API (asynchronous).
        Uses batch upload endpoint to avoid need for public file URLs.
        """
        log.info('Using MinerU Cloud API at %s', self.api_url)

        filename = os.path.basename(self.file_path)

        # Step 1: Request presigned upload URL
        batch_id, upload_url = self._request_upload_url(filename)

        # Step 2: Upload file to presigned URL
        self._upload_to_presigned_url(upload_url)

        # Step 3: Poll for results
        result = self._poll_batch_status(batch_id, filename)

        # Step 4: Download and extract markdown from ZIP
        markdown_content = self._download_and_extract_zip(result['full_zip_url'], filename)

        log.info('Successfully parsed document with MinerU Cloud API: %s', filename)

        # Create metadata
        metadata = {
            'source': filename,
            'api_mode': 'cloud',
            'batch_id': batch_id,
        }

        return [Document(page_content=markdown_content, metadata=metadata)]

    def _request_upload_url(self, filename: str) -> tuple:
        """
        Request presigned upload URL from Cloud API.
        Returns (batch_id, upload_url).
        """
        headers = {
            'Authorization': f'Bearer {self.api_key}',
            'Content-Type': 'application/json',
        }

        # Build request body
        request_body = {
            **self.params,
            'files': [
                {
                    'name': filename,
                    'is_ocr': self.enable_ocr,
                }
            ],
        }

        # Add page ranges if specified
        if self.page_ranges:
            request_body['files'][0]['page_ranges'] = self.page_ranges

        log.info('Requesting upload URL for: %s', filename)
        log.debug('Cloud API request body: %s', request_body)

        try:
            response = requests.post(
                f'{self.api_url}/file-urls/batch',
                headers=headers,
                json=request_body,
                timeout=30,
            )
            response.raise_for_status()
        except requests.HTTPError as e:
            error_detail = f'Failed to request upload URL: {e}'
            if e.response is not None:
                try:
                    error_data = e.response.json()
                    error_detail += f' - {error_data.get("msg", error_data)}'
                except Exception:
                    error_detail += f' - {e.response.text}'
            raise HTTPException(status.HTTP_400_BAD_REQUEST, detail=error_detail)
        except Exception as e:
            raise HTTPException(
                status.HTTP_500_INTERNAL_SERVER_ERROR,
                detail=f'Error requesting upload URL: {str(e)}',
            )

        try:
            result = response.json()
        except ValueError as e:
            raise HTTPException(
                status.HTTP_502_BAD_GATEWAY,
                detail=f'Invalid JSON response: {e}',
            )

        # Check for API error response
        if result.get('code') != 0:
            raise HTTPException(
                status.HTTP_400_BAD_REQUEST,
                detail=f'MinerU Cloud API error: {result.get("msg", "Unknown error")}',
            )

        data = result.get('data', {})
        batch_id = data.get('batch_id')
        file_urls = data.get('file_urls', [])

        if not batch_id or not file_urls:
            raise HTTPException(
                status.HTTP_502_BAD_GATEWAY,
                detail='MinerU Cloud API response missing batch_id or file_urls',
            )

        upload_url = file_urls[0]
        log.info('Received upload URL for batch: %s', batch_id)

        return batch_id, upload_url

    def _upload_to_presigned_url(self, upload_url: str) -> None:
        """
        Upload file to presigned URL (no authentication needed).
        """
        log.info(f'Uploading file to presigned URL')

        try:
            with open(self.file_path, 'rb') as f:
                response = requests.put(
                    upload_url,
                    data=f,
                    timeout=self.timeout,
                )
                response.raise_for_status()
        except FileNotFoundError:
            raise HTTPException(status.HTTP_404_NOT_FOUND, detail=f'File not found: {self.file_path}')
        except requests.Timeout:
            raise HTTPException(
                status.HTTP_504_GATEWAY_TIMEOUT,
                detail='File upload to presigned URL timed out',
            )
        except requests.HTTPError as e:
            raise HTTPException(
                status.HTTP_400_BAD_REQUEST,
                detail=f'Failed to upload file to presigned URL: {e}',
            )
        except Exception as e:
            raise HTTPException(
                status.HTTP_500_INTERNAL_SERVER_ERROR,
                detail=f'Error uploading file: {str(e)}',
            )

        log.info('File uploaded successfully')

    def _poll_batch_status(self, batch_id: str, filename: str) -> dict:
        """
        Poll batch status until completion.
        Returns the result dict for the file.
        """
        headers = {
            'Authorization': f'Bearer {self.api_key}',
        }

        max_iterations = 300  # 10 minutes max (2 seconds per iteration)
        poll_interval = 2  # seconds

        log.info('Polling batch status: %s', batch_id)

        for iteration in range(max_iterations):
            try:
                response = requests.get(
                    f'{self.api_url}/extract-results/batch/{batch_id}',
                    headers=headers,
                    timeout=30,
                )
                response.raise_for_status()
            except requests.HTTPError as e:
                error_detail = f'Failed to poll batch status: {e}'
                if e.response is not None:
                    try:
                        error_data = e.response.json()
                        error_detail += f' - {error_data.get("msg", error_data)}'
                    except Exception:
                        error_detail += f' - {e.response.text}'
                raise HTTPException(status.HTTP_400_BAD_REQUEST, detail=error_detail)
            except Exception as e:
                raise HTTPException(
                    status.HTTP_500_INTERNAL_SERVER_ERROR,
                    detail=f'Error polling batch status: {str(e)}',
                )

            try:
                result = response.json()
            except ValueError as e:
                raise HTTPException(
                    status.HTTP_502_BAD_GATEWAY,
                    detail=f'Invalid JSON response while polling: {e}',
                )

            # Check for API error response
            if result.get('code') != 0:
                raise HTTPException(
                    status.HTTP_400_BAD_REQUEST,
                    detail=f'MinerU Cloud API error: {result.get("msg", "Unknown error")}',
                )

            data = result.get('data', {})
            extract_result = data.get('extract_result', [])

            # Find our file in the batch results
            file_result = None
            for item in extract_result:
                if item.get('file_name') == filename:
                    file_result = item
                    break

            if not file_result:
                raise HTTPException(
                    status.HTTP_502_BAD_GATEWAY,
                    detail=f'File {filename} not found in batch results',
                )

            state = file_result.get('state')

            if state == 'done':
                log.info('Processing complete for %s', filename)
                return file_result
            elif state == 'failed':
                error_msg = file_result.get('err_msg', 'Unknown error')
                raise HTTPException(
                    status.HTTP_400_BAD_REQUEST,
                    detail=f'MinerU processing failed: {error_msg}',
                )
            elif state in ['waiting-file', 'pending', 'running', 'converting']:
                # Still processing
                if iteration % 10 == 0:  # Log every 20 seconds
                    log.info('Processing status: %s (iteration %s/%s)', state, iteration + 1, max_iterations)
                time.sleep(poll_interval)
            else:
                log.warning(f'Unknown state: {state}')
                time.sleep(poll_interval)

        # Timeout
        raise HTTPException(
            status.HTTP_504_GATEWAY_TIMEOUT,
            detail='MinerU processing timed out after 10 minutes',
        )

    def _download_and_extract_zip(self, zip_url: str, filename: str) -> str:
        """
        Download ZIP file from CDN and extract markdown content.
        Returns the markdown content as a string.
        """
        log.info('Downloading results from: %s', zip_url)

        try:
            response = requests.get(zip_url, timeout=60)
            response.raise_for_status()
        except requests.HTTPError as e:
            raise HTTPException(
                status.HTTP_400_BAD_REQUEST,
                detail=f'Failed to download results ZIP: {e}',
            )
        except Exception as e:
            raise HTTPException(
                status.HTTP_500_INTERNAL_SERVER_ERROR,
                detail=f'Error downloading results: {str(e)}',
            )

        # Save ZIP to temporary file before reading.
        tmp_zip_path = None
        markdown_content = None
        try:
            with tempfile.NamedTemporaryFile(delete=False, suffix='.zip') as tmp_zip:
                tmp_zip.write(response.content)
                tmp_zip_path = tmp_zip.name

            with zipfile.ZipFile(tmp_zip_path, 'r') as zip_ref:
                members = zip_ref.infolist()
                all_files = [member.filename for member in members]
                md_members = [member for member in members if member.filename.endswith('.md')]
                read_errors = []

                for member in md_members:
                    log.info('Found markdown file in ZIP: %s', member.filename)
                    try:
                        with zip_ref.open(member, 'r') as f:
                            if self.max_markdown_bytes is None:
                                content = f.read()
                            else:
                                content = f.read(self.max_markdown_bytes + 1)
                                if len(content) > self.max_markdown_bytes:
                                    raise HTTPException(
                                        status.HTTP_502_BAD_GATEWAY,
                                        detail=f'Markdown file in results ZIP is too large: {member.filename}',
                                    )
                            markdown_content = content.decode('utf-8')
                    except UnicodeDecodeError as e:
                        read_errors.append(f'{member.filename}: {e}')
                        log.warning(f'Failed to decode {member.filename}: {e}')
                        continue
                    except HTTPException:
                        raise
                    except Exception as e:
                        read_errors.append(f'{member.filename}: {e}')
                        log.warning(f'Failed to read {member.filename}: {e}')
                        continue
                    if markdown_content:
                        break

                if markdown_content is None:
                    log.error(f'Available files in ZIP: {all_files}')
                    if read_errors:
                        error_msg = f"Found .md files but couldn't read them: {read_errors}"
                    else:
                        error_msg = f'No .md files found in ZIP. Available files: {all_files}'
                    raise HTTPException(
                        status.HTTP_502_BAD_GATEWAY,
                        detail=error_msg,
                    )
        except zipfile.BadZipFile as e:
            raise HTTPException(
                status.HTTP_502_BAD_GATEWAY,
                detail=f'Invalid ZIP file received: {e}',
            )
        except HTTPException:
            raise
        except Exception as e:
            raise HTTPException(
                status.HTTP_500_INTERNAL_SERVER_ERROR,
                detail=f'Error extracting ZIP: {str(e)}',
            )
        finally:
            if tmp_zip_path:
                try:
                    os.unlink(tmp_zip_path)
                except FileNotFoundError:
                    pass
                except Exception as e:
                    log.warning(f'Failed to remove temporary ZIP file {tmp_zip_path}: {e}')

        if not markdown_content:
            raise HTTPException(
                status.HTTP_400_BAD_REQUEST,
                detail='Extracted markdown content is empty',
            )

        log.info('Successfully extracted markdown content (%s characters)', len(markdown_content))
        return markdown_content