from flask import request, jsonify
from Services.ai_worker import chatgpt_call
from Services.prompts import Prompts
import PyPDF2
import base64
import imageio
import io

prompts = Prompts()


def parse_medical_reports():
    """
    Accepts multiple files (PDF, images) and extracts medical data using AI
    Returns structured JSON - NO database storage (privacy-first)
    """
    try:
        if 'files' not in request.files:
            return jsonify({'message': 'No files uploaded', 'error': 'true'}), 400
        
        files = request.files.getlist('files')
        
        if len(files) == 0:
            return jsonify({'message': 'No files provided', 'error': 'true'}), 400
        
        extracted_data = []
        
        # Process each file
        for file in files:
            if not file.filename:
                continue
                
            file_type = file.content_type
            
            if 'pdf' in file_type.lower():
                # Extract text from PDF
                try:
                    pdf_reader = PyPDF2.PdfReader(file)
                    text = ""
                    
                    for page in pdf_reader.pages:
                        page_text = page.extract_text()
                        if page_text:
                            text += page_text + "\n"
                    
                    if text.strip():
                        extracted_data.append({
                            'type': 'pdf_text',
                            'content': text.strip(),
                            'filename': file.filename
                        })
                    else:
                        # PDF has no extractable text - might be image-based
                        file.seek(0)
                        image_result = process_image_file(file)
                        extracted_data.append(image_result)
                        
                except Exception as e:
                    print(f"Error processing PDF {file.filename}: {str(e)}")
                    return jsonify({
                        'message': f'Error processing PDF {file.filename}: {str(e)}',
                        'error': 'true'
                    }), 500
            
            elif 'image' in file_type.lower():
                # Process image file
                try:
                    image_result = process_image_file(file)
                    extracted_data.append(image_result)
                except Exception as e:
                    print(f"Error processing image {file.filename}: {str(e)}")
                    return jsonify({
                        'message': f'Error processing image {file.filename}: {str(e)}',
                        'error': 'true'
                    }), 500
            else:
                return jsonify({
                    'message': f'Unsupported file type: {file_type}. Only PDF and images are supported.',
                    'error': 'true'
                }), 400
        
        if not extracted_data:
            return jsonify({
                'message': 'No data could be extracted from the uploaded files',
                'error': 'true'
            }), 400
        
        # Parse all extracted data with AI
        parsed_data = parse_with_ai(extracted_data)
        
        return jsonify({
            'data': parsed_data,
            'files_processed': len(files),
            'error': False
        }), 200
        
    except Exception as e:
        print(f"Error in parse_medical_reports: {str(e)}")
        return jsonify({
            'message': f'Error parsing reports: {str(e)}',
            'error': 'true'
        }), 500


def process_image_file(file):
    """
    Process image file using imageio
    Reads image, converts to JPEG format, and encodes to base64
    """
    try:
        # Read image using imageio
        image = imageio.imread(file)
        
        # Create a BytesIO object to store the image
        img_byte_arr = io.BytesIO()
        
        # Write image to BytesIO in JPEG format
        imageio.imwrite(img_byte_arr, image, format='JPEG')
        
        # Get the byte data
        img_byte_arr.seek(0)
        image_bytes = img_byte_arr.read()
        
        # Encode to base64
        image_base64 = base64.b64encode(image_bytes).decode('utf-8')
        
        # Extract text using AI vision
        extracted_text = parse_image_with_ai_vision(image_base64, file.filename)
        
        return {
            'type': 'image',
            'content': extracted_text,
            'filename': file.filename
        }
        
    except Exception as e:
        raise Exception(f"Error processing image with imageio: {str(e)}")


def parse_image_with_ai_vision(image_base64, filename=""):
    """
    Use AI vision to extract text from medical images
    Works with DeepSeek vision models
    """
    from openai import OpenAI
    
    client = OpenAI(
        api_key="sk-e108ff73c53144379d0b8a4aefbaf525",
        base_url="https://api.deepseek.com",
    )
    
    try:
        response = client.chat.completions.create(
            model="deepseek-chat",
            messages=[
                {
                    "role": "system",
                    "content": "You are a medical document OCR expert. Extract ALL text and data from medical reports accurately, preserving the structure and format."
                },
                {
                    "role": "user",
                    "content": [
                        {
                            "type": "text",
                            "text": """Extract all medical information from this report image. Include:
                            - Test names and their values
                            - Units and reference ranges
                            - Doctor's notes and observations
                            - Patient information (if visible)
                            - Any abnormal findings or flags
                            - Date of report
                            - Laboratory name
                            
                            Format the output as clear, structured text, preserving the original layout."""
                        },
                        {
                            "type": "image_url",
                            "image_url": {
                                "url": f"data:image/jpeg;base64,{image_base64}"
                            }
                        }
                    ]
                }
            ],
            max_tokens=2000
        )
        
        return response.choices[0].message.content
        
    except Exception as e:
        # Fallback if vision processing fails
        print(f"Vision API failed for {filename}: {str(e)}")
        return f"[Image data from {filename} - vision processing failed: {str(e)}]"


def parse_with_ai(extracted_data):
    """
    Send all extracted data to AI for structured parsing
    Combines text from all files and creates structured JSON output
    """
    # Combine all extracted content
    separator = "\n\n" + "="*80 + "\n=== NEW DOCUMENT ===\n" + "="*80 + "\n\n"
    
    combined_content = separator.join([
        f"📄 File: {item['filename']}\n📋 Type: {item['type']}\n\n📝 Content:\n{item['content']}"
        for item in extracted_data
    ])
    
    # Create prompt for structured parsing
    prompt = prompts.parse_medical_reports(combined_content)
    
    # Get structured data from AI
    parsed_data = chatgpt_call(prompt, max_tokens=3000)
    
    return parsed_data