import gradio as gr
from fpdf import FPDF
from PIL import Image
from groq import Groq
from datetime import datetime, timezone
import base64
from io import BytesIO
import time, os

max_texts_rows = 5
max_texts_elements = 10
max_img_desc_rows = 5

init_texts_rows_state = 1
init_img_desc_rows_state = 0

custom_css = """
    html {
        height: 100%;
        margin: 0;
        display: flex;
        justify-content: center;
        align-items: center;
    }
    body {
      min-height: 50vh;
      aspect-ratio: 1 / 1.414; /* A4 aspect ratio */
      width: 100%; /* Full width */
      max-width: 210mm; /* Max width for A4 */
      max-height: 297mm; /* Max height for A4 */
      padding: 10px; /* Optional padding */
      box-sizing: border-box;
      margin: 0 auto; /* Center the block */
      overflow: auto; /* Ensure content is scrollable if it exceeds the block size */
      gap: 10px; /* Add gap between elements */
      border: 1px solid #000;
    }
    .gradio-container{
      background-color: white;
      
    }
    #logo-upload {
      min-height: 100px;
      max-height: 100px;
    }
    #text-empresa {
      min-height: 100px;
      max-height: 100px;
    }
    .text-box {
    }
    .image-upload{
      min-height: 300px;
      max-height: 300px;
    }
    img {
      object-fit: contain;
    }
    .description {
      min-height: 300px;
      max-height: 300px;
    }
    h1, h2, h3 {
      color: rgb(60, 60, 60);
    }
"""

api_key = os.getenv("GROQ_API_KEY")

client = Groq(
    api_key=api_key,
)

with gr.Blocks(css=custom_css) as demo:

    gr.Markdown("# Plataforma de Generación de Informes de ViLAB", elem_id="title")

    gr.Markdown("Importante: Plataforma de prueba experimental, se limita solo al uso didáctico.")

    gr.Markdown("### Ingrese el nombre de su empresa y su logo [opcional].", elem_id="subtitle")

    header_items = []

    with gr.Row(equal_height=True):
      with gr.Column(scale=1, min_width=160):
        nombre_empresa = gr.Textbox("Nombre de la Empresa", label=None, visible = True, container = False, elem_id="text-empresa", interactive=True)
        header_items.append(nombre_empresa)
      with gr.Column(scale=1, min_width=160):
        logo = gr.Image(label="Logo", elem_id="logo-upload", visible=True, interactive=True, mirror_webcam=False, sources="upload")
        header_items.append(logo)

    gr.Markdown("---")

    # -----------------------------------------------------------------------------------------------------

    gr.Markdown("## Plantilla de Informes")

    gr.Markdown("### Ingrese los datos para el informe.")
    gr.Markdown("### (Agregue campos con el botón '+', remueva campos con el botón '-')")

    texts_state = gr.State(2)

    textboxes = []

    for i in range(max_texts_rows):
      with gr.Row():
          with gr.Column(scale=1, min_width=160):
            textbox_l = gr.Textbox("Título: Descripción", show_label=False, elem_classes="text-box", visible = (i<init_texts_rows_state), interactive=True )
            textboxes.append(textbox_l)
          with gr.Column(scale=1, min_width=160):
            if i == 0:
              textbox_r = gr.Textbox("Fecha: " + datetime.now().strftime("%d/%m/%Y"), show_label=False, elem_classes="text-box", visible = (i<init_texts_rows_state), interactive=False )
            else:
              textbox_r = gr.Textbox("Título: Descripción", show_label=False, elem_classes="text-box", visible = (i<init_texts_rows_state), interactive=True )
            textboxes.append(textbox_r)

    with gr.Row():
      with gr.Column(scale=1, min_width=160):
        plus_button = gr.Button("+")
      with gr.Column(scale=1, min_width=160):
        minus_button = gr.Button("-")

    def variable_outputs_plus(k):
        k = int(k)
        k = min(max_texts_elements, k + 1)
        return [gr.Textbox(visible=True)]*(k) + [gr.Textbox(visible=False)]*(max_texts_elements-k) + [k]

    def variable_outputs_minus(k):
        k = int(k)
        k = max(1, k - 1)
        return [gr.Textbox(visible=True)]*(k) + [gr.Textbox(visible=False)]*(max_texts_elements-k) + [k]

    plus_button.click(variable_outputs_plus, texts_state, textboxes + [texts_state])
    minus_button.click(variable_outputs_minus, texts_state, textboxes + [texts_state])

    gr.Markdown("---")

    # -----------------------------------------------------------------------------------------------------

    gr.Markdown("## Plantilla de Resultados")

    gr.Markdown("### Ingrese los resultados con imagen y descripción.")
    gr.Markdown("### (Agregue campos con el botón '+', remueva campos con el botón '-')")

    images_description_state = gr.State(0)

    img_desc_boxes = []

    for i in range(max_img_desc_rows):
      with gr.Row():
        with gr.Column(scale=1, min_width=160):
          image = gr.Image(label=f"Imagen {i}", interactive=True, mirror_webcam=False, sources="upload", elem_classes="image-upload", container = True, visible = (i<init_img_desc_rows_state))
          img_desc_boxes.append(image)
        with gr.Column(scale=1, min_width=160):
          description = gr.TextArea("Describa su muestra ... ", label=f"Descripción", visible=(i<init_img_desc_rows_state), elem_classes="description", interactive=True, max_lines=10 )
          img_desc_boxes.append(description)

    with gr.Row():
      with gr.Column(scale=1, min_width=160):
        plus_img_desc_button = gr.Button("+")
      with gr.Column(scale=1, min_width=160):
        minus_img_desc_button = gr.Button("-")


    def variable_outputs_plus_img_desc(k):
        k = int(k)
        k = min(max_img_desc_rows, k + 1)
        return [gr.Image(visible=True), gr.TextArea(visible=True)]*k + [gr.Image(visible=False), gr.TextArea(visible=False)]*(max_img_desc_rows-k) + [k]

    def variable_outputs_minus_img_desc(k):
        k = int(k)
        k = max(0, k - 1)
        return [gr.TextArea(visible=True)]*(2*k) + [gr.TextArea(visible=False)]*(2*(max_texts_rows-k)) + [k]

    plus_img_desc_button.click(variable_outputs_plus_img_desc, images_description_state, img_desc_boxes + [images_description_state])
    minus_img_desc_button.click(variable_outputs_minus_img_desc, images_description_state, img_desc_boxes + [images_description_state])

    gr.Markdown("---")

    # -----------------------------------------------------------------------------------------------------

    state_data = gr.State([])

    with gr.Row():
      submit_button = gr.DownloadButton (label="Generar Reporte")

    file_output = gr.File(interactive=False, visible=False)

    # Función para mantener la relación de aspecto de la imagen
    def get_aspect_ratio_image_size(img, max_width, max_height):
        width, height = img.size
        aspect_ratio = width / height
        if width > max_width or height > max_height:
            if width / max_width > height / max_height:
                width = max_width
                height = width / aspect_ratio
            else:
                height = max_height
                width = height * aspect_ratio
        return width, height

    class PDF(FPDF):

      def __init__(self, datos, *args, **kwargs):
          super().__init__(*args, **kwargs)
          self.datos = datos

      def set_text(self, datos):
          self.datos = datos

      def footer(self):
          self.set_y(-15)
          self.set_font('Helvetica', '', 8)
          self.cell(0, 10, 'Página %s' % self.page_no(), 0, new_x='RIGHT', new_y='TOP', align='C')

      def body(self):

          # Datos de la Empresa
          nombre_empresa = "Nombre de la Empresa: " + self.datos[0][0]
          logo_path = self.datos[0][1]
          if logo_path is not None:
            image = Image.fromarray(logo_path)
            logo_width, logo_height = get_aspect_ratio_image_size(image, 30, 30)
            self.image(image, 10, self.get_y(), logo_width, logo_height)
            self.ln(logo_height + 5)
          self.set_font('Helvetica', '', 12)
          self.cell(0, 10, nombre_empresa, 0, new_x='LMARGIN', new_y='NEXT', align='L')

          # Línea horizontal divisoria
          self.line(10, self.get_y(), self.w - 10, self.get_y())
          self.ln(5)

          # Datos del Informe
          self.set_font('Helvetica', '', 11)
          for i in range(len(self.datos[1])):
            if i%2 != 0:
              if self.datos[1][i] != "Título: Descripción":
                self.cell(self.w / 2 - 10, 10, self.datos[1][i], 0, new_x='LMARGIN', new_y='NEXT', align='L')
            if i%2 == 0:
              if self.datos[1][i] != "Título: Descripción":
                self.cell(self.w / 2 - 10, 10, self.datos[1][i], 0, new_x='RIGHT', new_y='TOP', align='L')

          # Línea horizontal divisoria
          self.ln(10)
          self.line(10, self.get_y(), self.w - 10, self.get_y())
          self.ln(5)

          # Datos del Resultado
          for i in range(len(self.datos[2])):
              if self.datos[2][i] is not None:
                  image = Image.fromarray(self.datos[2][i])
                  img_width, img_height = get_aspect_ratio_image_size(image, self.w / 2 - 20, 100)
                  if self.get_y() + img_height + 10 > self.h - 30:
                      self.add_page()
                  self.image(image, 10, self.get_y(), img_width, img_height)
                  if self.datos[3][i] != "Describa su muestra ... ":
                      self.set_font('Helvetica', '', 10)
                      self.set_xy(self.w / 2 + 10, self.get_y())
                      self.multi_cell(self.w / 2 - 20, None, self.datos[3][i])
                  self.ln(img_height + 5)
              else:
                  if self.datos[3][i] != "Describa su muestra ... ":
                      self.set_font('Helvetica', '', 10)
                      if self.get_y() + 10 > self.h - 30:
                          self.add_page()
                      self.set_xy(self.w / 2 + 10, self.get_y())
                      self.multi_cell(self.w / 2 - 20, None, self.datos[3][i])

          # Resumen del Informe
          self.set_font('Helvetica', '', 10)
          self.ln(10)  # Añadir espacio antes del resumen
          self.multi_cell(0, 6, "Resumen - Historial Clínico: " + self.datos[4], 0, 'L')


    def get_values(*args):
      n_args = len(args)

      data = [
          list(args[0:2]),
          list(args[2:2+max_texts_elements-1]),
          list(args[2+max_texts_elements:n_args:2]),
          list(args[2+max_texts_elements+1:n_args+1:2])
      ]

      nombre_empresa = data[0][0]
      datos_informe = [i for i in data[1] if i != "Título: Descripción"]
      resultados_informe = [i for i in data[3] if i != "Describa su muestra ... "]  

      prompt = f"""Eres un narrador del historial clínico en español basado en el siguiente contexto:
-------------------------------
Nombre de la empresa: {nombre_empresa}
{' - '.join(datos_informe)}
Descripciones según las imágenes digitalizadas con ViLAB desde el microscopio: {' - '.join(resultados_informe)}
-------------------------------
- Nunca generes información extra fuera del contexto.
- Añade informaciones útiles

Solo presenta el resumen del historial clínico en un párrafo.
Resumen:"""

      chat_completion = client.chat.completions.create(
          messages=[
              {
                  "role": "user",
                  "content": prompt,
              }
          ],
          model="llama3-70b-8192",
      )

      data.append(chat_completion.choices[0].message.content)

      # Crear PDF
      pdf = PDF(data)
      pdf.add_page()
      pdf.body()
      pdf_path = "./documento_generado.pdf"
      pdf.output(pdf_path)

      return gr.File(pdf_path, interactive=False, visible=True)

    submit_button.click(get_values, header_items + textboxes + img_desc_boxes, file_output)

if __name__ == "__main__":
    
    demo.launch(share=False)