# backend/utils.py
import os
import time
import pytesseract
from PIL import Image
from pdf2image import convert_from_path
import shutil
import asyncio
import uuid
import subprocess
import base64
import re
from models import get_prompt

from reportlab.lib.pagesizes import A4
from reportlab.platypus import SimpleDocTemplate, Paragraph, Spacer, PageBreak
from reportlab.lib.styles import getSampleStyleSheet, ParagraphStyle
from reportlab.lib.enums import TA_JUSTIFY, TA_CENTER

# Se agrega BeautifulSoup para procesar HTML
from bs4 import BeautifulSoup

UPLOAD_FOLDER = "uploads"
OUTPUT_FOLDER = "outputs"

def encode_image(file_path):
    with open(file_path, "rb") as f:
        return base64.b64encode(f.read()).decode("utf-8")

def run_tesseract(file_path):
    """
    Si el archivo es un PDF, se realiza OCR página a página añadiendo al principio de cada una
    la cabecera con el formato [Page 0001], [Page 0002], etc.
    Para otros formatos se realiza OCR normal.
    """
    extracted_text = ""
    if file_path.lower().endswith(".pdf"):
        try:
            pages = convert_from_path(file_path)
            for i, page in enumerate(pages, start=1):
                page_text = pytesseract.image_to_string(page, lang='eng')
                extracted_text += f"[Page {i:04d}]\n{page_text}\n\n"
        except Exception as e:
            extracted_text = f"❌ Error processing PDF: {str(e)}"
    else:
        image = Image.open(file_path)
        extracted_text = pytesseract.image_to_string(image, lang='eng')
    return extracted_text

def call_api_correction(api, model, text, prompt_key="ocr_correction"):
    prompt = get_prompt(prompt_key) + text
    try:
        from google import genai
    except ImportError:
        raise ImportError("Please install 'google-genai' to use Gemini.")
    try:
        loop = asyncio.get_event_loop()
    except RuntimeError:
        loop = asyncio.new_event_loop()
        asyncio.set_event_loop(loop)
    client = genai.Client(api_key=os.environ.get("GEMINI_API_KEY"))
    response = client.models.generate_content(model=model, contents=[prompt])
    return response.text

def call_api_ocr(api, model, file_path, prompt_key="ocr"):
    try:
        from google import genai
    except ImportError:
        raise ImportError("Please install 'google-genai' to use Gemini.")
    try:
        loop = asyncio.get_event_loop()
    except RuntimeError:
        loop = asyncio.new_event_loop()
        asyncio.set_event_loop(loop)
    client = genai.Client(api_key=os.environ.get("GEMINI_API_KEY"))
    file_ref = client.files.upload(file=file_path)
    prompt = get_prompt(prompt_key)
    response = client.models.generate_content(model=model, contents=[file_ref, prompt])
    return response.text

def call_api_translation(api, model, text, target_language, prompt_key="translation"):
    prompt_template = get_prompt(prompt_key)
    prompt = prompt_template.format(target_language=target_language) + text
    try:
        from google import genai
    except ImportError:
        raise ImportError("Please install 'google-genai' to use Gemini.")
    try:
        loop = asyncio.get_event_loop()
    except RuntimeError:
        loop = asyncio.new_event_loop()
        asyncio.set_event_loop(loop)
    client = genai.Client(api_key=os.environ.get("GEMINI_API_KEY"))
    response = client.models.generate_content(model=model, contents=[prompt])
    return response.text

def embed_ocr_in_pdf(input_pdf, output_pdf):
    try:
        subprocess.run(["ocrmypdf", input_pdf, output_pdf], check=True)
        return True
    except Exception as e:
        return False

def ocr_file_by_pages(file_path, api, model, prompt_key, update_progress, is_cancelled):
    final_text = ""
    if file_path.lower().endswith(".pdf"):
        try:
            pages = convert_from_path(file_path)
        except Exception as e:
            return f"❌ Error processing PDF: {str(e)}"
        total = len(pages)
        for i, page in enumerate(pages, start=1):
            if is_cancelled():
                update_progress(0, "⏹️ Cancelled")
                return "Process cancelled."
            temp_filename = os.path.join(OUTPUT_FOLDER, f"temp_page_{uuid.uuid4().hex}.png")
            page.save(temp_filename, "PNG")
            page_text = call_api_ocr(api, model, temp_filename, prompt_key)
            final_text += f"[Page {i:04d}]\n{page_text}\n\n"
            os.remove(temp_filename)
            progress = int((i / total) * 100)
            update_progress(progress, f"📄 Processed page {i} of {total}.")
            time.sleep(1)
        return final_text
    else:
        return call_api_ocr(api, model, file_path, prompt_key)

def translate_file_by_pages(file_path, api, model, target_language, prompt_key, update_progress, is_cancelled):
    final_translation = ""
    if file_path.lower().endswith(".pdf"):
        try:
            pages = convert_from_path(file_path)
        except Exception as e:
            return f"❌ Error processing PDF: {str(e)}"
        total = len(pages)
        for i, page in enumerate(pages, start=1):
            if is_cancelled():
                update_progress(0, "⏹️ Cancelled")
                return "Process cancelled."
            temp_filename = os.path.join(OUTPUT_FOLDER, f"temp_page_{uuid.uuid4().hex}.png")
            page.save(temp_filename, "PNG")
            page_text = pytesseract.image_to_string(page, lang='eng')
            translated_page = call_api_translation(api, model, page_text, target_language, prompt_key)
            final_translation += f"[Page {i:04d}]\n{translated_page}\n\n"
            os.remove(temp_filename)
            progress = int((i / total) * 100)
            update_progress(progress, f"📄 Processed page {i} of {total}.")
            time.sleep(1)
        return final_translation
    elif file_path.lower().endswith(".txt"):
        with open(file_path, "r", encoding="utf-8") as f:
            text = f.read()
        translated = call_api_translation(api, model, text, target_language, prompt_key)
        update_progress(100, "🎉 Process completed")
        return f"[Page {1:04d}]\n{translated}"
    else:
        return "Unsupported file type for translation."

def process_file(file_path, api, model, mode, prompt_key, update_progress, is_cancelled):
    """
    NOTA: Ahora solo usamos [Page XXXX] al inicio de cada página (por Tesseract).
    """
    if is_cancelled():
        update_progress(0, "⏹️ Cancelled")
        return "Process cancelled."
    update_progress(25, "📤 File uploaded.")
    base_name = os.path.splitext(os.path.basename(file_path))[0]

    if mode == "OCR":
        processed_text = run_tesseract(file_path)
        if is_cancelled():
            update_progress(25, "⏹️ Cancelled")
            return "Process cancelled."
        update_progress(50, "✅ Tesseract OCR completed.")
        pdf_output = os.path.join(OUTPUT_FOLDER, base_name + "_ocr.pdf")
        if file_path.lower().endswith(".pdf"):
            if embed_ocr_in_pdf(file_path, pdf_output):
                update_progress(95, "📄 OCR embedded into PDF.")
            else:
                shutil.copy(file_path, pdf_output)
                update_progress(95, "⚠️ Failed to embed OCR; original PDF copied.")
        else:
            shutil.copy(file_path, pdf_output)

    elif mode == "OCR + AI":
        processed_text = run_tesseract(file_path)
        if is_cancelled():
            update_progress(25, "⏹️ Cancelled")
            return "Process cancelled."
        update_progress(50, "✅ Tesseract OCR completed.")
        processed_text = call_api_correction(api, model, processed_text, prompt_key)
        pdf_output = os.path.join(OUTPUT_FOLDER, base_name + "_ocr.pdf")
        if file_path.lower().endswith(".pdf"):
            if embed_ocr_in_pdf(file_path, pdf_output):
                update_progress(95, "📄 OCR embedded into PDF.")
            else:
                shutil.copy(file_path, pdf_output)
                update_progress(95, "⚠️ Failed to embed OCR; original PDF copied.")
        else:
            shutil.copy(file_path, pdf_output)

    elif mode == "AI":
        update_progress(25, "📂 File ready for full AI processing.")
        if file_path.lower().endswith(".pdf"):
            processed_text = ocr_file_by_pages(file_path, api, model, prompt_key, update_progress, is_cancelled)
        else:
            processed_text = call_api_ocr(api, model, file_path, prompt_key)

    else:
        processed_text = "Unrecognized processing mode."
        update_progress(25, "❌ Error: Unrecognized mode.")

    if is_cancelled():
        update_progress(75, "⏹️ Cancelled")
        return "Process cancelled."

    update_progress(75, "🤖 API processing completed.")

    # Guardar el TXT final
    txt_file = os.path.join(OUTPUT_FOLDER, base_name + ".txt")
    with open(txt_file, "w", encoding="utf-8") as f:
        f.write(processed_text)

    update_progress(100, "🎉 Process completed")
    return processed_text

def organize_paragraphs(text):
    """
    Organiza el texto plano en párrafos de forma más flexible.
    Si se detecta doble salto de línea se usa como separador;
    si no, se procesa línea a línea para unirlas en párrafos, creando uno nuevo cuando:
      - Se encuentra una línea vacía, o
      - La línea actual termina en punto.
    Dentro de cada párrafo se unen las líneas; se inserta un <br/> si la línea termina en punto.
    """
    text = text.strip()
    paragraphs = []
    if "\n\n" in text:
        blocks = re.split(r'\n\s*\n', text)
    else:
        lines = text.splitlines()
        blocks = []
        buffer = ""
        for line in lines:
            stripped = line.strip()
            if not stripped:
                if buffer:
                    blocks.append(buffer)
                    buffer = ""
            else:
                if buffer and buffer.endswith("."):
                    blocks.append(buffer)
                    buffer = stripped
                else:
                    if buffer:
                        buffer += " " + stripped
                    else:
                        buffer = stripped
        if buffer:
            blocks.append(buffer)
    for block in blocks:
        line_list = block.splitlines()
        if len(line_list) == 1:
            paragraphs.append(line_list[0].strip())
        else:
            new_block = ""
            for i, line in enumerate(line_list):
                line = line.strip()
                if not line:
                    continue
                if i < len(line_list) - 1:
                    if line.endswith("."):
                        new_block += line + "<br/>"
                    else:
                        new_block += line + " "
                else:
                    new_block += line
            paragraphs.append(new_block.strip())
    return paragraphs

def convert_txt_to_pdf(txt_file_path):
    """
    Convierte un archivo TXT a PDF siguiendo estas reglas:
      - Se detecta el patrón [Page XXXX] para separar páginas (solo al inicio).
      - Este marcador se convierte en un encabezado (h1).
      - Si el contenido proviene de un bloque markdown con ```html se elimina ese marcador y
        se parsea con BeautifulSoup para generar párrafos independientes.
      - Para contenido en texto plano se organiza en párrafos con organize_paragraphs.
      - Se inserta un PageBreak después de cada bloque de página.
    """
    with open(txt_file_path, "r", encoding="utf-8") as f:
        content = f.read()

    base_name = os.path.splitext(os.path.basename(txt_file_path))[0]
    output_pdf = os.path.join(OUTPUT_FOLDER, base_name + "_txt.pdf")
    doc = SimpleDocTemplate(
        output_pdf,
        pagesize=A4,
        rightMargin=40, leftMargin=40,
        topMargin=40, bottomMargin=40
    )
    styles = getSampleStyleSheet()
    header_styles = {
        "h1": ParagraphStyle('Heading1', parent=styles['Heading1'], alignment=TA_CENTER),
        "h2": ParagraphStyle('Heading2', parent=styles['Heading2'], alignment=TA_CENTER),
        "h3": ParagraphStyle('Heading3', parent=styles['Heading3'], alignment=TA_CENTER),
        "h4": ParagraphStyle('Heading4', parent=styles['Heading4'], alignment=TA_CENTER),
        "h5": ParagraphStyle('Heading5', parent=styles['Heading5'], alignment=TA_CENTER),
        "h6": ParagraphStyle('Heading6', parent=styles['Heading6'], alignment=TA_CENTER),
    }
    normal_style = ParagraphStyle(
        'Normal',
        parent=styles['Normal'],
        alignment=TA_JUSTIFY,
        leading=15,
        leftIndent=20
    )

    flowables = []
    # Patrón para detectar el marcador [Page XXXX]
    page_pattern = re.compile(r'\[Page\s+\d{4}\]')
    parts = re.split(r'(\[Page\s+\d{4}\])', content)
    first_page_encountered = False

    for part in parts:
        part = part.strip()
        if not part:
            continue

        if page_pattern.fullmatch(part):
            if first_page_encountered:
                flowables.append(PageBreak())
            else:
                first_page_encountered = True
            header_para = Paragraph(part, styles['Heading1'])
            flowables.append(header_para)
            flowables.append(Spacer(1, 12))
        else:
            if part.startswith("```"):
                lines = part.splitlines()
                if lines and lines[0].startswith("```"):
                    lines = lines[1:]
                if lines and lines[-1].strip() == "```":
                    lines = lines[:-1]
                part = "\n".join(lines)
            # Si parece HTML, se procesa para separar cada etiqueta de interés
            if re.search(r'<\s*html', part, re.IGNORECASE) or re.search(r'<\s*(p|h[1-6])', part, re.IGNORECASE):
                if not part.lower().startswith("<html"):
                    part = "<html>" + part + "</html>"
                soup = BeautifulSoup(part, "html.parser")
                for element in soup.find_all(['h1', 'h2', 'h3', 'h4', 'h5', 'h6', 'p']):
                    if element.name.lower() in ["h1", "h2", "h3", "h4", "h5", "h6"]:
                        style = header_styles.get(element.name.lower(), styles['Heading1'])
                    else:
                        style = normal_style
                    text_content = element.get_text().strip()
                    if text_content:
                        flowables.append(Paragraph(text_content, style))
                        flowables.append(Spacer(1, 12))
            else:
                # Texto plano -> organizar en párrafos
                paragraphs = organize_paragraphs(part)
                for para_text in paragraphs:
                    if para_text:
                        flowables.append(Paragraph(para_text, normal_style))
                        flowables.append(Spacer(1, 12))

    doc.build(flowables)
    return output_pdf