import fitz  # PyMuPDF
import base64
import os
from typing import Dict, List, Tuple, Optional, Union
import ollama  # Import the Ollama library
import tempfile
import logging

from pathlib import Path

class Section:
    """Represents a document section with type, content, and bounding box."""
    
    def __init__(self, section_type: str, content: str, bbox: Tuple[float, float, float, float]):
        self.type = section_type
        self.content = content
        self.bbox = bbox
    
    def to_dict(self) -> Dict:
        return {
            'type': self.type,
            'content': self.content,
            'bbox': self.bbox
        }

def get_available_models() -> List[str]:
    """
    Get a list of available models from Ollama.
    
    Returns:
        List[str]: A list of available model names.
    """
    try:
        models_response = client.list()
        logger.info(f"Ollama models response: {models_response}")
        
        if "models" in models_response and isinstance(models_response["models"], list):
            model_names = [model["model"] for model in models_response["models"]]
            return model_names
        else:
            logger.warning("No models found in Ollama response")
            return []
    except Exception as e:
        logger.error(f"Error fetching Ollama models: {e}")
        return []

def _extract_text_blocks(page: fitz.Page, show_text: bool) -> List[Section]:
    """Extract text blocks from a PDF page."""
    sections = []
    blocks = page.get_text("blocks")
    
    for block in blocks:
        if block[6] == 0:  # Text block
            rect = fitz.Rect(block[:4])
            content = block[4] if block[4] is not None else ''
            
            if content.strip():  # Only add non-empty text
                section = Section('text', content, (rect.x0, rect.y0, rect.x1, rect.y1))
                sections.append(section)
                
                if show_text:
                    page.draw_rect(rect, color=(1, 0, 0), width=1.5)
    
    return sections

def _extract_image_data(page: fitz.Page, img_index: int) -> Optional[bytes]:
    """Extract image data as bytes from PDF page."""
    try:
        img_list = page.get_images(full=True)
        if img_index < len(img_list):
            img_info = img_list[img_index]
            xref = img_info[0]
            pix = fitz.Pixmap(page.parent, xref)
            
            # Convert CMYK to RGB if needed
            if pix.n - pix.alpha < 4:
                img_data = pix.tobytes("png")
            else:
                pix1 = fitz.Pixmap(fitz.csRGB, pix)
                img_data = pix1.tobytes("png")
                pix1 = None
            
            pix = None
            return img_data
    except Exception as e:
        logger.warning(f"Failed to extract image data: {e}")
    return None

def _process_image_with_ai(img_data: bytes, model: str = "qwen2.5vl:7b") -> str:
    """Process image with AI to extract text or analyze diagrams."""
    try:
        import base64
        
        # Encode image to base64
        img_base64 = base64.b64encode(img_data).decode()
        
        # Create prompt for image analysis
        prompt = """Extract all text from this image exactly as it appears. Do not analyze, describe, or interpret the image. Only output the readable text you see, preserving the original formatting and layout as much as possible. If there is no text, respond with "[No text found]"."""

        # Send to Ollama with image
        response = client.chat(
            model=model,
            messages=[
                {
                    'role': 'user',
                    'content': prompt,
                    'images': [img_base64]
                }
            ]
        )
        
        if response and 'message' in response and 'content' in response['message']:
            return response['message']['content']
        else:
            logger.error("Invalid response format from Ollama vision model")
            return "[Image: Could not analyze]"
            
    except Exception as e:
        logger.error(f"Error processing image with AI: {e}")
        return "[Image: Analysis failed]"

def _extract_images(page: fitz.Page, show_images: bool, process_with_ai: bool = False) -> List[Section]:
    """Extract images from a PDF page and optionally process with AI."""
    sections = []
    
    img_list = page.get_images(full=True)
    for img_index, img in enumerate(img_list):
        try:
            bbox = page.get_image_bbox(img)
            if bbox:
                content = "[Image]"
                
                # Process image with AI if enabled
                if process_with_ai:
                    logger.info(f"Processing image {img_index + 1} with AI...")
                    img_data = _extract_image_data(page, img_index)
                    if img_data:
                        ai_analysis = _process_image_with_ai(img_data)
                        content = f"**[Image Analysis]**\n\n{ai_analysis}"
                
                section = Section('image', content, (bbox.x0, bbox.y0, bbox.x1, bbox.y1))
                sections.append(section)
                
                if show_images:
                    page.draw_rect(bbox, color=(0, 0, 1), width=1.5)
        except Exception as e:
            logger.warning(f"Failed to extract image: {e}")
            continue
    
    return sections

# Removed table extraction functions - focusing on text and images only

def process_pdf(pdf_path: str, show_text: bool, show_images: bool) -> Tuple[fitz.Document, Dict[int, List[Dict]]]:
    """
    Process the PDF to extract sections (text, images, tables) in reading order.
    
    Args:
        pdf_path (str): Path to the PDF file.
        show_text (bool): Whether to draw borders around text blocks.
        show_images (bool): Whether to draw borders around images.
    
    Returns:
        Tuple[fitz.Document, Dict[int, List[Dict]]]: The processed PDF document and a dictionary of page sections.
    """
    try:
        doc = fitz.open(pdf_path)
        page_sections: Dict[int, List[Dict]] = {}
        
        for page_num in range(len(doc)):
            page = doc.load_page(page_num)
            sections = []
            
            # Extract different types of content
            text_sections = _extract_text_blocks(page, show_text)
            image_sections = _extract_images(page, show_images, process_with_ai=False)
            
            # Combine all sections
            all_sections = text_sections + image_sections
            
            # Sort sections by reading order (top to bottom, left to right)
            all_sections.sort(key=lambda s: (s.bbox[1], s.bbox[0]))
            
            # Convert to dictionary format
            page_sections[page_num] = [section.to_dict() for section in all_sections]
            
            logger.info(f"Page {page_num + 1}: Found {len(text_sections)} text, "
                       f"{len(image_sections)} image sections")
        
        return doc, page_sections
    
    except Exception as e:
        logger.error(f"Error processing PDF: {e}")
        raise

def get_page_images(doc: fitz.Document, zoom_level: float = 1.5) -> Dict[int, str]:
    """
    Render each page of the PDF as a base64-encoded PNG image.
    
    Args:
        doc (fitz.Document): The PDF document.
        zoom_level (float): The zoom level for rendering (affects image size).
    
    Returns:
        Dict[int, str]: A dictionary of page images as base64 strings.
    """
    page_images = {}
    
    try:
        for page_num in range(len(doc)):
            page = doc.load_page(page_num)
            
            # Create transformation matrix for zoom
            matrix = fitz.Matrix(zoom_level, zoom_level)
            pix = page.get_pixmap(matrix=matrix)
            
            # Convert to PNG bytes and encode to base64
            img_bytes = pix.tobytes("png")
            img_base64 = base64.b64encode(img_bytes).decode()
            page_images[page_num] = img_base64
            
            logger.debug(f"Generated preview image for page {page_num + 1}")
            
    except Exception as e:
        logger.error(f"Error generating page images: {e}")
        
    return page_images

def save_and_encode_pdf(doc: fitz.Document) -> Tuple[bytes, str]:
    """
    Save the modified PDF to a temporary file and encode it to base64.
    
    Args:
        doc (fitz.Document): The PDF document to save and encode.
    
    Returns:
        Tuple[bytes, str]: The raw PDF bytes and the base64-encoded string.
    """
    try:
        # Create temporary file
        with tempfile.NamedTemporaryFile(delete=False, suffix=".pdf") as temp_file:
            modified_pdf_path = temp_file.name
        
        # Save document
        doc.save(modified_pdf_path)
        doc.close()
        
        # Read and encode
        with open(modified_pdf_path, "rb") as f:
            modified_pdf_bytes = f.read()
        
        # Clean up temporary file
        try:
            os.unlink(modified_pdf_path)
        except OSError:
            logger.warning(f"Could not delete temporary file: {modified_pdf_path}")
        
        base64_encoded = base64.b64encode(modified_pdf_bytes).decode("utf-8")
        logger.info("PDF successfully saved and encoded")
        
        return modified_pdf_bytes, base64_encoded
        
    except Exception as e:
        logger.error(f"Error saving and encoding PDF: {e}")
        raise

def format_text_with_llama(text: str, model: str) -> str:
    """
    Use the selected Ollama model to format the text into Markdown.
    
    Args:
        text (str): The text to format.
        model (str): The model to use for formatting.
    
    Returns:
        str: The formatted Markdown text.
    """
    if not text.strip():
        logger.warning("Empty text provided for formatting")
        return text
    
    try:
        logger.info(f"Formatting text with model: {model}")
        response = client.generate(
            model=model,
            prompt=text
        )
        
        if "response" in response:
            return response["response"]
        else:
            logger.error("Invalid response format from Ollama")
            return text
            
    except Exception as e:
        logger.error(f"Error formatting text with Ollama: {e}")
        return text  # Return the original text if formatting fails
    
def initialize_a_logger(logger_name: str = "./logs/ChatBot.log"):
    logger = logging.getLogger('example_logger')

    # Set the logging level (DEBUG, INFO, WARNING, ERROR, CRITICAL)
    logger.setLevel(logging.DEBUG)

    log_path = Path(logger_name).expanduser().resolve()
    log_path.parent.mkdir(parents=True, exist_ok=True)
    # Use the same resolved path as mkdir — relative paths + cwd quirks broke FileHandler in Docker.
    fh = logging.FileHandler(str(log_path), mode="w")
    logger.addHandler(fh)
    return logger

def build_prompt(text: str) -> str:
    

    return f"""
        As an expert auditor, read the text line by line, analyze semantic chunks,
        and infer the compliance rules in the text thoroughly.
        Return only valid JSON. If information is missing, leave fields empty.
        
        Input Text:
        {text}
        
        LLM Tasks:
        - Parse text as semantic chunks, analyze input and the Type descriptions and concisely extract compliance rules based on the format.
        1. Ubiquitous Requirements. Description: No pre-condition, Always active
        Format: <system name> <modal> <system response>
        Example: "The Bank SHALL maintain audit trails for all change requests."
        
        2. Event-driven Requirements. Description: Begin with WHEN, Triggered by specific events
        Format: WHEN <trigger> <system name> <modal> <system response>
        Example: "WHEN acquiring software packages, the Bank SHALL perform a detailed evaluation to ensure user and business requirements are met."
        
        3. Unwanted Behavior Requirements. Description: Begin with IF...THEN, Express undesirable situations to be handled
        Format: IF <condition> THEN <system name> <modal> <system response>
        Example: "IF a supplier's financial stability is in doubt, THEN the Bank SHALL have alternatives to mitigate potential loss of service."
        
        4. State-driven Requirements. Description: Begin with WHILE, Active during specific system states
        Format: WHILE <state> <system name> <modal> <system response>
        Example: "WHILE systems are in production, the Bank SHALL maintain hardware and software requirements at recovery locations."
        
        5. Optional Feature Requirements. Description: Begin with WHERE, Apply when certain features are enabled/present
        Format: WHERE <feature> <system name> <modal> <system response>
        Example: "WHERE third-party service providers are involved, the Bank SHALL implement secure remote access mechanisms."

        
        6. Complex Requirements. Description: Combine multiple condition types, Handle more sophisticated scenarios
        Example: "WHEN changes are implemented to information systems, IF the change is an emergency change, THEN the Bank SHALL follow formal emergency change management procedures, INCLUDING regularization of implemented emergency changes."
        
        
        - Normalize language → turn “should/must/required/have to” into clear compliance rules.

        
        Output JSON:
        {{
            "requirements": ["output 1", "output 2", ...]
        }}
        """

def extract_pdf_text(pdf_path: str):
    """
    Extracts text page-by-page from a PDF, sends each to the Ollama endpoint,
    and returns structured compliance rules in JSON.
    """

    def parse_pdf_to_json(text: str):
        prompt = build_prompt(text)
        url = "http://localhost:4402/test_ollama_1"

        try:
            # Send prompt to your FastAPI Ollama endpoint
            response = requests.post(url, json={"prompt": prompt})
            response.raise_for_status()
            raw = response.json()

            # Try to clean and parse LLM JSON output
            content = raw.get("oss_json", "")
            cleaned = content.replace("```json", "").replace("```", "").strip()

            try:
                parsed = json.loads(cleaned)
            except json.JSONDecodeError:
                parsed = {"requirements": [cleaned]}

            return parsed

        except Exception as e:
            print(f"Error processing page: {e}")
            return {"error": str(e)}

    reader = PdfReader(pdf_path)
    all_requirements = []

    for i, page in enumerate(reader.pages):
        text = page.extract_text()
        if not text:
            continue

        print(f"Processing page {i + 1}...")
        result = parse_pdf_to_json(text)
        all_requirements.append({"page": i + 1, "result": result})

    return all_requirements





