"""
MinerU Client - Support both online API and local vLLM backend
"""

import os
import subprocess
import requests
import zipfile
import time
import shutil
from pathlib import Path
from typing import List, Optional, Dict, Union
from urllib.parse import urlparse


class MinerUClient:
    """MinerU client supporting both online API and local vLLM backend"""
    
    def __init__(
        self,
        api_token: Optional[str] = None,
        base_url: str = "https://mineru.net/api/v4",
        use_local: bool = False,
        local_url: str = "http://127.0.0.1:30000"
    ):
        """
        Initialize MinerU client
        
        Args:
            api_token: MinerU API token (required for online mode)
            base_url: MinerU API base URL (default: https://mineru.net/api/v4)
            use_local: Whether to use local vLLM backend (default: False)
            local_url: Local vLLM backend URL (default: http://127.0.0.1:30000)
        """
        self.use_local = use_local
        self.local_url = local_url
        
        if not use_local:
            # Online mode
            self.api_token = api_token or os.environ.get("MINERU_API_TOKEN")
            if not self.api_token:
                raise ValueError(
                    "MINERU_API_TOKEN is required for online mode. "
                    "Set it as environment variable or pass as parameter."
                )
            self.base_url = base_url
            self.headers = {
                "Content-Type": "application/json",
                "Authorization": f"Bearer {self.api_token}"
            }
            self.session = requests.Session()
            self.session.proxies = {'http': None, 'https': None}
            self.session.verify = True
            self.session.timeout = 30
        else:
            # Local mode - check if mineru command is available
            # Use -v instead of --version as it's faster and more reliable
            try:
                result = subprocess.run(
                    ["mineru", "-v"],
                    capture_output=True,
                    check=True,
                    timeout=10,
                    text=True
                )
            except FileNotFoundError:
                raise ValueError(
                    "MinerU command not found. Please install MinerU first:\n"
                    "  conda create -n mineru python=3.12 -y\n"
                    "  conda activate mineru\n"
                    "  uv pip install mineru"
                )
            except subprocess.TimeoutExpired:
                # If version check times out, still allow initialization
                # The actual command will fail later if mineru is not available
                pass
            except subprocess.CalledProcessError:
                raise ValueError(
                    "MinerU command check failed. Please ensure MinerU is properly installed."
                )
    
    def process_file(
        self,
        input_path: Union[str, Path],
        output_path: Union[str, Path],
        is_ocr: bool = True,
        enable_formula: bool = True,
        enable_table: bool = True,
        language: str = "en",
        layout_model: str = "doclayout_yolo"
    ) -> Dict:
        """
        Process a single file using MinerU
        
        Args:
            input_path: Path to input file (PDF, PNG, JPG, etc.)
            output_path: Path to output directory
            is_ocr: Enable OCR recognition
            enable_formula: Enable formula recognition
            enable_table: Enable table recognition
            language: Document language (en, ch, etc.)
            layout_model: Layout model name
            
        Returns:
            Dictionary with processing result
        """
        input_path = Path(input_path)
        output_path = Path(output_path)
        output_path.mkdir(parents=True, exist_ok=True)
        
        if self.use_local:
            return self._process_local(input_path, output_path)
        else:
            return self._process_online(
                input_path, output_path, is_ocr, enable_formula,
                enable_table, language, layout_model
            )
    
    def process_files_batch(
        self,
        file_paths: List[Union[str, Path]],
        output_dir: Union[str, Path],
        is_ocr: bool = True,
        enable_formula: bool = True,
        enable_table: bool = True,
        language: str = "en",
        layout_model: str = "doclayout_yolo",
        max_files_per_batch: int = 200
    ) -> Dict:
        """
        Process multiple files in batch
        
        Args:
            file_paths: List of input file paths
            output_dir: Output directory
            is_ocr: Enable OCR recognition
            enable_formula: Enable formula recognition
            enable_table: Enable table recognition
            language: Document language
            layout_model: Layout model name
            max_files_per_batch: Maximum files per batch
            
        Returns:
            Dictionary with batch processing result
        """
        if self.use_local:
            # Local mode: process files sequentially
            results = []
            for file_path in file_paths:
                result = self.process_file(
                    file_path, output_dir, is_ocr, enable_formula,
                    enable_table, language, layout_model
                )
                results.append(result)
            return {
                'success': all(r.get('success', False) for r in results),
                'processed_files': results,
                'total_count': len(results),
                'success_count': sum(1 for r in results if r.get('success', False))
            }
        else:
            # Online mode: use batch API
            return self._process_batch_online(
                file_paths, output_dir, is_ocr, enable_formula,
                enable_table, language, layout_model, max_files_per_batch
            )
    
    def _process_local(
        self,
        input_path: Path,
        output_path: Path
    ) -> Dict:
        """Process file using local MinerU vLLM backend"""
        try:
            # Build mineru command
            cmd = [
                "mineru",
                "-p", str(input_path),
                "-o", str(output_path),
                "-b", "vlm-http-client",
                "-u", self.local_url
            ]
            
            # Execute command
            result = subprocess.run(
                cmd,
                capture_output=True,
                text=True,
                timeout=600  # 10 minutes timeout
            )
            
            if result.returncode == 0:
                # Find output markdown file (mineru creates files in subdirectories like output/1/vlm/)
                md_files = list(output_path.rglob("*.md"))
                if md_files:
                    # Sort by modification time, get the most recent one
                    md_files.sort(key=lambda p: p.stat().st_mtime, reverse=True)
                    return {
                        'success': True,
                        'input_path': str(input_path),
                        'output_path': str(output_path),
                        'md_file': str(md_files[0]),
                        'message': 'File processed successfully'
                    }
                else:
                    # Provide more debugging info
                    return {
                        'success': False,
                        'error': f'No markdown file found in output directory: {output_path}. '
                                f'Command stdout: {result.stdout[:200] if result.stdout else "None"}'
                    }
            else:
                # Combine both stdout and stderr for better error info
                error_msg = result.stderr if result.stderr else result.stdout
                return {
                    'success': False,
                    'error': f'MinerU command failed (return code: {result.returncode}). '
                            f'Error: {error_msg}'
                }
        except subprocess.TimeoutExpired:
            return {
                'success': False,
                'error': 'Processing timeout (exceeded 10 minutes)'
            }
        except Exception as e:
            return {
                'success': False,
                'error': f'Processing failed: {str(e)}'
            }
    
    def _process_online(
        self,
        input_path: Path,
        output_path: Path,
        is_ocr: bool,
        enable_formula: bool,
        enable_table: bool,
        language: str,
        layout_model: str
    ) -> Dict:
        """Process file using online MinerU API"""
        try:
            # Submit task
            submit_result = self._submit_task_online(
                input_path, is_ocr, enable_formula, enable_table, language, layout_model
            )
            if not submit_result['success']:
                return submit_result
            
            task_id = submit_result['task_id']
            
            # Wait for completion
            max_wait_time = 300  # 5 minutes
            start_time = time.time()
            
            while time.time() - start_time < max_wait_time:
                status_result = self._check_task_status(task_id)
                if not status_result['success']:
                    return status_result
                
                state = status_result['state']
                if state == 'done':
                    zip_url = status_result.get('zip_url')
                    if zip_url:
                        return self._download_and_extract(zip_url, output_path)
                    else:
                        return {
                            'success': False,
                            'error': 'Task completed but no download URL found'
                        }
                elif state == 'failed':
                    return {
                        'success': False,
                        'error': f'Task failed: {status_result.get("error_msg", "Unknown error")}'
                    }
                
                time.sleep(5)
            
            return {
                'success': False,
                'error': f'Task timeout (exceeded {max_wait_time} seconds)'
            }
        except Exception as e:
            return {
                'success': False,
                'error': f'Processing failed: {str(e)}'
            }
    
    def _submit_task_online(
        self,
        file_path: Path,
        is_ocr: bool,
        enable_formula: bool,
        enable_table: bool,
        language: str,
        layout_model: str
    ) -> Dict:
        """Submit task to online API"""
        try:
            # Get upload URL
            response = self.session.post(
                f"{self.base_url}/file-urls",
                headers=self.headers,
                json={
                    "name": file_path.name,
                    "is_ocr": is_ocr,
                    "enable_formula": enable_formula,
                    "enable_table": enable_table,
                    "language": language,
                    "layout_model": layout_model
                }
            )
            
            if response.status_code != 200:
                return {
                    'success': False,
                    'error': f'Request failed with status {response.status_code}'
                }
            
            result = response.json()
            if result.get("code") != 0:
                return {
                    'success': False,
                    'error': result.get("msg", "Unknown error")
                }
            
            # Upload file
            upload_url = result["data"]["file_url"]
            task_id = result["data"]["task_id"]
            
            with open(file_path, 'rb') as f:
                upload_response = self.session.put(upload_url, data=f)
                if upload_response.status_code not in [200, 201]:
                    return {
                        'success': False,
                        'error': f'Upload failed with status {upload_response.status_code}'
                    }
            
            return {
                'success': True,
                'task_id': task_id
            }
        except Exception as e:
            return {
                'success': False,
                'error': f'Submit task failed: {str(e)}'
            }
    
    def _check_task_status(self, task_id: str) -> Dict:
        """Check task status"""
        try:
            response = self.session.get(
                f"{self.base_url}/extract/task/{task_id}",
                headers=self.headers
            )
            
            if response.status_code == 200:
                result = response.json()
                if result.get('code') == 0:
                    data = result.get('data', {})
                    return {
                        'success': True,
                        'state': data.get('state'),
                        'zip_url': data.get('full_zip_url'),
                        'error_msg': data.get('err_msg', '')
                    }
                else:
                    return {
                        'success': False,
                        'error': result.get('msg', 'Get task status failed')
                    }
            else:
                return {
                    'success': False,
                    'error': f'HTTP error: {response.status_code}'
                }
        except Exception as e:
            return {
                'success': False,
                'error': f'Check task status failed: {str(e)}'
            }
    
    def _download_and_extract(self, zip_url: str, output_path: Path) -> Dict:
        """Download and extract result"""
        try:
            # Download ZIP
            response = self.session.get(zip_url, stream=True)
            response.raise_for_status()
            
            zip_path = output_path / f"mineru_result_{int(time.time())}.zip"
            with open(zip_path, 'wb') as f:
                for chunk in response.iter_content(chunk_size=8192):
                    if chunk:
                        f.write(chunk)
            
            # Extract
            extract_dir = output_path / zip_path.stem
            with zipfile.ZipFile(zip_path, 'r') as zip_ref:
                zip_ref.extractall(extract_dir)
            
            # Find markdown file
            md_files = list(extract_dir.rglob("*.md"))
            if md_files:
                return {
                    'success': True,
                    'output_path': str(output_path),
                    'extract_dir': str(extract_dir),
                    'md_file': str(md_files[0]),
                    'message': 'File processed successfully'
                }
            else:
                return {
                    'success': False,
                    'error': 'No markdown file found in result'
                }
        except Exception as e:
            return {
                'success': False,
                'error': f'Download and extract failed: {str(e)}'
            }
    
    def _process_batch_online(
        self,
        file_paths: List[Path],
        output_dir: Path,
        is_ocr: bool,
        enable_formula: bool,
        enable_table: bool,
        language: str,
        layout_model: str,
        max_files_per_batch: int
    ) -> Dict:
        """Process batch files using online API"""
        # Simplified batch processing - submit all files
        # In production, you might want to implement proper batch handling
        results = []
        for file_path in file_paths:
            result = self._process_online(
                file_path, output_dir, is_ocr, enable_formula,
                enable_table, language, layout_model
            )
            results.append(result)
            time.sleep(1)  # Rate limiting
        
        return {
            'success': all(r.get('success', False) for r in results),
            'processed_files': results,
            'total_count': len(results),
            'success_count': sum(1 for r in results if r.get('success', False))
        }

