import os
import tempfile
import asyncio
from typing import List, Optional
from pathlib import Path
import logging
from io import BytesIO
import base64

from fastapi import FastAPI, File, UploadFile, HTTPException, BackgroundTasks
from fastapi.responses import HTMLResponse, JSONResponse
from fastapi.middleware.cors import CORSMiddleware
import uvicorn

# PDF and image processing
import PyPDF2
from pdf2image import convert_from_path, convert_from_bytes
from PIL import Image

# ML and AI
import torch
from transformers import AutoProcessor, AutoModelForVision2Seq

# Try importing anthropic with error handling
try:
    import anthropic
    ANTHROPIC_AVAILABLE = True
except ImportError as e:
    print(f"Warning: Anthropic library not available: {str(e)}")
    ANTHROPIC_AVAILABLE = False

# Try importing docling with fallback
try:
    from docling_core.types.doc import DoclingDocument
    from docling_core.types.doc.document import DocTagsDocument
    DOCLING_AVAILABLE = True
except ImportError:
    print("Warning: docling_core not available. Using fallback markdown generation.")
    DOCLING_AVAILABLE = False

# Environment and configuration
from dotenv import load_dotenv
import aiofiles
from config import config

# Load environment variables
load_dotenv()

# Logging setup
logging.basicConfig(level=logging.INFO)
logger = logging.getLogger(__name__)

app = FastAPI(
    title="PDF Parsing with SmolDocling",
    description="Extract text from PDFs using SmolDocling model and summarize with Anthropic API",
    version="1.0.0"
)

# Add CORS middleware
app.add_middleware(
    CORSMiddleware,
    allow_origins=["*"],
    allow_credentials=True,
    allow_methods=["*"],
    allow_headers=["*"],
)

# Global variables for model
processor = None
model = None
anthropic_client = None
DEVICE = "cuda" if torch.cuda.is_available() else "cpu"

class PDFProcessor:
    def __init__(self):
        self.max_pages_per_chunk = config.MAX_PAGES_PER_CHUNK
        
    async def pdf_to_images(self, pdf_bytes: bytes) -> List[Image.Image]:
        """Convert PDF bytes to list of PIL Images"""
        try:
            # Convert PDF to images
            images = convert_from_bytes(
                pdf_bytes,
                dpi=config.PDF_DPI,
                fmt='RGB'
            )
            logger.info(f"Converted PDF to {len(images)} images")
            return images
        except Exception as e:
            logger.error(f"Error converting PDF to images: {str(e)}")
            raise HTTPException(status_code=400, detail=f"Error converting PDF: {str(e)}")
    
    def chunk_images(self, images: List[Image.Image]) -> List[List[Image.Image]]:
        """Chunk images into smaller groups for processing"""
        chunks = []
        for i in range(0, len(images), self.max_pages_per_chunk):
            chunk = images[i:i + self.max_pages_per_chunk]
            chunks.append(chunk)
        logger.info(f"Created {len(chunks)} chunks from {len(images)} images")
        return chunks
    
    async def process_image_with_smoldocling(self, image: Image.Image) -> str:
        """Process single image with SmolDocling model"""
        global processor, model
        
        if processor is None or model is None:
            logger.warning("SmolDocling model not available, using basic OCR fallback")
            return await self._basic_ocr_fallback(image)
        
        try:
            # Prepare the input messages
            messages = [
                {
                    "role": "user",
                    "content": [
                        {"type": "image"},
                        {"type": "text", "text": "Convert this page to docling."}
                    ]
                },
            ]
            
            # Process with the model
            prompt = processor.apply_chat_template(messages, add_generation_prompt=True)
            inputs = processor(text=prompt, images=[image], return_tensors="pt")
            inputs = inputs.to(DEVICE)
            
            # Generate output
            with torch.no_grad():
                generated_ids = model.generate(**inputs, max_new_tokens=config.MAX_NEW_TOKENS)
                prompt_length = inputs.input_ids.shape[1]
                trimmed_generated_ids = generated_ids[:, prompt_length:]
                doctags = processor.batch_decode(
                    trimmed_generated_ids,
                    skip_special_tokens=False,
                )[0].lstrip()
            
            # Convert to markdown using docling if available, otherwise return raw doctags
            if DOCLING_AVAILABLE:
                try:
                    doctags_doc = DocTagsDocument.from_doctags_and_image_pairs([doctags], [image])
                    doc = DoclingDocument.load_from_doctags(doctags_doc, document_name="Page")
                    markdown_content = doc.export_to_markdown()
                    return markdown_content
                except Exception as e:
                    logger.warning(f"Docling conversion failed: {str(e)}, using raw doctags")
                    return self._convert_doctags_to_markdown(doctags)
            else:
                # Fallback: convert doctags to basic markdown
                return self._convert_doctags_to_markdown(doctags)
            
        except Exception as e:
            logger.error(f"Error processing image with SmolDocling: {str(e)}")
            return f"Error processing page: {str(e)}"
    
    async def process_pdf_chunk(self, images: List[Image.Image]) -> str:
        """Process a chunk of images"""
        markdown_parts = []
        
        for i, image in enumerate(images):
            try:
                logger.info(f"Processing image {i+1}/{len(images)} in chunk")
                markdown = await self.process_image_with_smoldocling(image)
                markdown_parts.append(f"## Page {i+1}\n\n{markdown}\n\n")
            except Exception as e:
                logger.error(f"Error processing image {i+1}: {str(e)}")
                markdown_parts.append(f"## Page {i+1}\n\nError processing this page: {str(e)}\n\n")
        
        return "".join(markdown_parts)
    
    def _convert_doctags_to_markdown(self, doctags: str) -> str:
        """Fallback method to convert doctags to basic markdown when docling_core is not available"""
        try:
            # Simple conversion of common doctags to markdown
            lines = doctags.split('\n')
            markdown_lines = []
            
            for line in lines:
                line = line.strip()
                if not line:
                    markdown_lines.append('')
                    continue
                    
                # Convert common doctags to markdown
                if line.startswith('<title>') and line.endswith('</title>'):
                    title = line.replace('<title>', '').replace('</title>', '')
                    markdown_lines.append(f'# {title}')
                elif line.startswith('<heading>') and line.endswith('</heading>'):
                    heading = line.replace('<heading>', '').replace('</heading>', '')
                    markdown_lines.append(f'## {heading}')
                elif line.startswith('<text>') and line.endswith('</text>'):
                    text = line.replace('<text>', '').replace('</text>', '')
                    markdown_lines.append(text)
                elif line.startswith('<list>') and line.endswith('</list>'):
                    item = line.replace('<list>', '').replace('</list>', '')
                    markdown_lines.append(f'- {item}')
                elif line.startswith('<table>') and line.endswith('</table>'):
                    table = line.replace('<table>', '').replace('</table>', '')
                    markdown_lines.append(f'| {table} |')
                elif line.startswith('<formula>') and line.endswith('</formula>'):
                    formula = line.replace('<formula>', '').replace('</formula>', '')
                    markdown_lines.append(f'$$\n{formula}\n$$')
                elif line.startswith('<code>') and line.endswith('</code>'):
                    code = line.replace('<code>', '').replace('</code>', '')
                    markdown_lines.append(f'```\n{code}\n```')
                else:
                    # Remove any remaining tags and add as text
                    import re
                    clean_text = re.sub(r'<[^>]+>', '', line)
                    if clean_text.strip():
                        markdown_lines.append(clean_text)
            
            return '\n'.join(markdown_lines)
            
        except Exception as e:
            logger.error(f"Error converting doctags to markdown: {str(e)}")
            return f"**Raw DocTags Output:**\n\n```\n{doctags}\n```"
    
    async def _basic_ocr_fallback(self, image: Image.Image) -> str:
        """Basic OCR fallback when SmolDocling is not available"""
        try:
            # Try to use pytesseract if available
            try:
                import pytesseract
                text = pytesseract.image_to_string(image)
                return f"# Extracted Text (Basic OCR)\n\n{text}"
            except ImportError:
                pass
            
            # If pytesseract is not available, return a placeholder
            return f"""# PDF Page Processed

**Note**: SmolDocling model is not available. This page contains content that would normally be extracted using advanced OCR.

To get full text extraction capabilities, please:
1. Ensure the SmolDocling model loads correctly
2. Check that all dependencies are installed
3. Try using a GPU-enabled environment for better performance

Image dimensions: {image.size[0]} x {image.size[1]} pixels
"""
            
        except Exception as e:
            logger.error(f"Basic OCR fallback failed: {str(e)}")
            return f"# Error Processing Page\n\nFailed to process this page: {str(e)}"
 
class SummaryGenerator:
    def __init__(self, api_key: str):
        if not ANTHROPIC_AVAILABLE:
            raise ImportError("Anthropic library is not available")
            
        try:
            # Initialize Anthropic client with explicit parameters
            self.client = anthropic.Anthropic(
                api_key=api_key
            )
            logger.info("Anthropic client created successfully")
        except Exception as e:
            logger.error(f"Failed to initialize Anthropic client: {str(e)}")
            raise e
    
    async def summarize_text(self, text: str) -> str:
        """Generate summary using Anthropic Claude API"""
        try:
            # If text is too long, chunk it
            max_tokens = config.MAX_TOKENS_PER_CHUNK * 2  # Claude's context window
            if len(text.split()) > max_tokens:
                # Split text into chunks and summarize each, then combine
                chunks = self._chunk_text(text, config.MAX_TOKENS_PER_CHUNK)
                chunk_summaries = []
                
                for i, chunk in enumerate(chunks):
                    logger.info(f"Summarizing chunk {i+1}/{len(chunks)}")
                    summary = await self._summarize_chunk(chunk)
                    chunk_summaries.append(summary)
                
                # Combine chunk summaries into final summary
                combined_text = "\n\n".join(chunk_summaries)
                final_summary = await self._summarize_chunk(combined_text, is_final=True)
                return final_summary
            else:
                return await self._summarize_chunk(text)
                
        except Exception as e:
            logger.error(f"Error generating summary: {str(e)}")
            raise HTTPException(status_code=500, detail=f"Error generating summary: {str(e)}")
    
    def _chunk_text(self, text: str, max_tokens: int) -> List[str]:
        """Split text into chunks"""
        words = text.split()
        chunks = []
        current_chunk = []
        
        for word in words:
            current_chunk.append(word)
            if len(current_chunk) >= max_tokens:
                chunks.append(" ".join(current_chunk))
                current_chunk = []
        
        if current_chunk:
            chunks.append(" ".join(current_chunk))
        
        return chunks
    
    async def _summarize_chunk(self, text: str, is_final: bool = False) -> str:
        """Summarize a single chunk of text"""
        if is_final:
            prompt = f"""Please provide a comprehensive final summary of this document based on the following chunk summaries:

{text}

Create a well-structured, detailed summary that captures all the key points, main themes, and important details from the entire document."""
        else:
            prompt = f"""Please provide a detailed summary of the following text, capturing all key points, main themes, and important details:

{text}

Make sure to preserve important information that might be needed for a final comprehensive summary."""
        
        try:
            message = self.client.messages.create(
                model=config.ANTHROPIC_MODEL,
                max_tokens=config.ANTHROPIC_MAX_TOKENS,
                temperature=config.ANTHROPIC_TEMPERATURE,
                messages=[
                    {
                        "role": "user",
                        "content": prompt
                    }
                ]
            )
            return message.content[0].text
        except Exception as e:
            logger.error(f"Error calling Anthropic API: {str(e)}")
            return f"Error generating summary: {str(e)}"

# Initialize processors
pdf_processor = PDFProcessor()
summary_generator = None

@app.on_event("startup")
async def startup_event():
    """Initialize models and clients on startup"""
    global processor, model, summary_generator
    
    logger.info("Loading SmolDocling model...")
    try:
        # Load the SmolDocling model
        model_id = config.MODEL_ID
        
        # Try loading with different approaches
        try:
            # First try: Standard loading
            logger.info("Attempting to load processor...")
            processor = AutoProcessor.from_pretrained(
                model_id,
                trust_remote_code=True,
                use_fast=False
            )
            logger.info("Processor loaded successfully")
            
        except Exception as e:
            logger.warning(f"Standard processor loading failed: {str(e)}")
            # Fallback: Try with explicit trust_remote_code
            try:
                from transformers import AutoTokenizer, AutoImageProcessor
                processor = AutoProcessor.from_pretrained(
                    model_id,
                    trust_remote_code=True,
                    revision="main"
                )
                logger.info("Processor loaded with trust_remote_code=True")
            except Exception as e2:
                logger.error(f"All processor loading attempts failed: {str(e2)}")
                raise e2
        
        # Load the model
        logger.info("Loading model...")
        model = AutoModelForVision2Seq.from_pretrained(
            model_id,
            torch_dtype=torch.bfloat16 if DEVICE == "cuda" else torch.float32,
            trust_remote_code=True,
            _attn_implementation="eager",  # Use eager attention for better compatibility
            device_map="auto" if DEVICE == "cuda" else None,
        )
        
        if DEVICE != "cuda":
            model = model.to(DEVICE)
            
        logger.info(f"Model loaded successfully on {DEVICE}")
        
        # Initialize Anthropic client
        if config.ANTHROPIC_API_KEY and ANTHROPIC_AVAILABLE:
            try:
                summary_generator = SummaryGenerator(config.ANTHROPIC_API_KEY)
                logger.info("Anthropic client initialized successfully")
            except Exception as e:
                logger.error(f"Failed to initialize Anthropic client: {str(e)}")
                logger.warning("Summary generation will not be available due to Anthropic client error.")
                summary_generator = None
        else:
            if not config.ANTHROPIC_API_KEY:
                logger.warning("ANTHROPIC_API_KEY not found. Summary generation will not be available.")
            if not ANTHROPIC_AVAILABLE:
                logger.warning("Anthropic library not available. Summary generation will not be available.")
            summary_generator = None
            
    except Exception as e:
        logger.error(f"Error loading model: {str(e)}")
        logger.error("The application will still work for basic PDF text extraction without the SmolDocling model.")
        # Don't raise the error - let the app start without the model
        processor = None
        model = None

@app.get("/")
async def root():
    """Serve the main HTML interface"""
    html_content = """
    <!DOCTYPE html>
    <html>
    <head>
        <title>PDF Parser with SmolDocling</title>
        <style>
            body { font-family: Arial, sans-serif; max-width: 800px; margin: 0 auto; padding: 20px; }
            .upload-area { border: 2px dashed #ccc; padding: 20px; text-align: center; margin: 20px 0; }
            .result-area { margin-top: 20px; padding: 20px; background-color: #f5f5f5; border-radius: 5px; }
            button { background-color: #007bff; color: white; padding: 10px 20px; border: none; border-radius: 5px; cursor: pointer; }
            button:hover { background-color: #0056b3; }
            button:disabled { background-color: #ccc; cursor: not-allowed; }
            .progress { display: none; margin: 10px 0; }
            .error { color: red; }
            .success { color: green; }
        </style>
    </head>
    <body>
        <h1>📄 PDF Parser with SmolDocling</h1>
        <p>Upload a PDF document to extract text and generate a summary using AI.</p>
        
        <div class="upload-area">
            <input type="file" id="pdfFile" accept=".pdf" />
            <br><br>
            <button onclick="uploadPDF()" id="uploadBtn">Process PDF</button>
            <div class="progress" id="progress">Processing... Please wait.</div>
        </div>
        
        <div class="result-area" id="results" style="display: none;">
            <h3>Results:</h3>
            <div id="resultContent"></div>
        </div>

        <script>
        async function uploadPDF() {
            const fileInput = document.getElementById('pdfFile');
            const uploadBtn = document.getElementById('uploadBtn');
            const progress = document.getElementById('progress');
            const results = document.getElementById('results');
            const resultContent = document.getElementById('resultContent');
            
            if (!fileInput.files.length) {
                alert('Please select a PDF file');
                return;
            }
            
            const formData = new FormData();
            formData.append('file', fileInput.files[0]);
            
            uploadBtn.disabled = true;
            progress.style.display = 'block';
            results.style.display = 'none';
            
            try {
                const response = await fetch('/upload-pdf/', {
                    method: 'POST',
                    body: formData
                });
                
                const result = await response.json();
                
                if (response.ok) {
                    resultContent.innerHTML = `
                        <div class="success">✅ PDF processed successfully!</div>
                        <h4>Extracted Text (Markdown):</h4>
                        <pre style="white-space: pre-wrap; background: white; padding: 15px; border-radius: 5px; max-height: 400px; overflow-y: auto;">${result.markdown}</pre>
                        ${result.summary ? `
                        <h4>Summary:</h4>
                        <div style="background: white; padding: 15px; border-radius: 5px; border-left: 4px solid #007bff;">${result.summary}</div>
                        ` : ''}
                        <p><small>Processing time: ${result.processing_time} seconds</small></p>
                    `;
                    results.style.display = 'block';
                } else {
                    resultContent.innerHTML = `<div class="error">❌ Error: ${result.detail}</div>`;
                    results.style.display = 'block';
                }
            } catch (error) {
                resultContent.innerHTML = `<div class="error">❌ Error: ${error.message}</div>`;
                results.style.display = 'block';
            } finally {
                uploadBtn.disabled = false;
                progress.style.display = 'none';
            }
        }
        </script>
    </body>
    </html>
    """
    return HTMLResponse(content=html_content)

@app.get("/health")
async def health_check():
    """Health check endpoint"""
    return {"status": "healthy", "model_loaded": model is not None}

@app.post("/upload-pdf/")
async def upload_pdf(file: UploadFile = File(...)):
    """Upload and process PDF file"""
    import time
    start_time = time.time()
    
    # Validate file
    if not file.filename.lower().endswith('.pdf'):
        raise HTTPException(status_code=400, detail="Only PDF files are allowed")
    
    try:
        # Read PDF content
        pdf_content = await file.read()
        logger.info(f"Received PDF file: {file.filename} ({len(pdf_content)} bytes)")
        
        # Convert PDF to images
        images = await pdf_processor.pdf_to_images(pdf_content)
        
        # Process in chunks if PDF is large
        image_chunks = pdf_processor.chunk_images(images)
        all_markdown = []
        
        for i, chunk in enumerate(image_chunks):
            logger.info(f"Processing chunk {i+1}/{len(image_chunks)} ({len(chunk)} pages)")
            chunk_markdown = await pdf_processor.process_pdf_chunk(chunk)
            all_markdown.append(chunk_markdown)
        
        # Combine all markdown
        full_markdown = "\n".join(all_markdown)
        
        # Generate summary if Anthropic client is available
        summary = None
        if summary_generator:
            try:
                logger.info("Generating summary...")
                summary = await summary_generator.summarize_text(full_markdown)
            except Exception as e:
                logger.error(f"Summary generation failed: {str(e)}")
                summary = f"Summary generation failed: {str(e)}"
        
        processing_time = round(time.time() - start_time, 2)
        
        result = {
            "message": "PDF processed successfully",
            "filename": file.filename,
            "total_pages": len(images),
            "chunks_processed": len(image_chunks),
            "markdown": full_markdown,
            "summary": summary,
            "processing_time": processing_time
        }
        
        logger.info(f"PDF processing completed in {processing_time} seconds")
        return result
        
    except Exception as e:
        logger.error(f"Error processing PDF: {str(e)}")
        raise HTTPException(status_code=500, detail=str(e))

if __name__ == "__main__":
    uvicorn.run(app, host="0.0.0.0", port=7860)