from concurrent.futures import ThreadPoolExecutor, as_completed
from typing import Annotated, List, Tuple, Literal

from pydantic import BaseModel
from tqdm import tqdm
from PIL import Image

from marker.processors.llm import BaseLLMComplexBlockProcessor
from marker.schema import BlockTypes
from marker.schema.blocks import Block, TableCell
from marker.schema.document import Document


class LLMTableMergeProcessor(BaseLLMComplexBlockProcessor):
    block_types: Annotated[
        Tuple[BlockTypes],
        "The block types to process.",
    ] = (BlockTypes.Table, BlockTypes.TableOfContents)
    table_height_threshold: Annotated[
        float,
        "The minimum height ratio relative to the page for the first table in a pair to be considered for merging.",
    ] = 0.6
    table_start_threshold: Annotated[
        float,
        "The maximum percentage down the page the second table can start to be considered for merging."
    ] = 0.2
    vertical_table_height_threshold: Annotated[
        float,
        "The height tolerance for 2 adjacent tables to be merged into one."
    ] = 0.25
    vertical_table_distance_threshold: Annotated[
        int,
        "The maximum distance between table edges for adjacency."
    ] = 20
    horizontal_table_width_threshold: Annotated[
        float,
        "The width tolerance for 2 adjacent tables to be merged into one."
    ] = 0.25
    horizontal_table_distance_threshold: Annotated[
        int,
        "The maximum distance between table edges for adjacency."
    ] = 10
    column_gap_threshold: Annotated[
        int,
        "The maximum gap between columns to merge tables"
    ] = 50
    disable_tqdm: Annotated[
        bool,
        "Whether to disable the tqdm progress bar.",
    ] = False
    table_merge_prompt: Annotated[
        str,
        "The prompt to use for rewriting text.",
        "Default is a string containing the Gemini rewriting prompt."
    ] = """You're a text correction expert specializing in accurately reproducing tables from PDFs.
You'll receive two images of tables from successive pages of a PDF.  Table 1 is from the first page, and Table 2 is from the second page.  Both tables may actually be part of the same larger table. Your job is to decide if Table 2 should be merged with Table 1, and how they should be joined.  The should only be merged if they're part of the same larger table, and Table 2 cannot be interpreted without merging.

You'll specify your judgement in json format - first whether Table 2 should be merged with Table 1, then the direction of the merge, either `bottom` or `right`.  A bottom merge means that the rows of Table 2 are joined to the rows of Table 1. A right merge means that the columns of Table 2 are joined to the columns of Table 1.  (bottom merge is equal to np.vstack, right merge is equal to np.hstack)

Table 2 should be merged at the bottom of Table 1 if Table 2 has no headers, and the rows have similar values, meaning that Table 2 continues Table 1. Table 2 should be merged to the right of Table 1 if each row in Table 2 matches a row in Table 1, meaning that Table 2 contains additional columns that augment Table 1.

Only merge Table 1 and Table 2 if Table 2 cannot be interpreted without merging.  Only merge Table 1 and Table 2 if you can read both images properly.

**Instructions:**
1. Carefully examine the provided table images.  Table 1 is the first image, and Table 2 is the second image.
2. Examine the provided html representations of Table 1 and Table 2.
3. Write a description of Table 1.
4. Write a description of Table 2.
5. Analyze whether Table 2 should be merged into Table 1, and write an explanation.
6. Output your decision on whether they should be merged, and merge direction.
**Example:**
Input:
Table 1
```html
<table>
    <tr>
        <th>Name</th>
        <th>Age</th>
        <th>City</th>
        <th>State</th>
    </tr>
    <tr>
        <td>John</td>
        <td>25</td>
        <td>Chicago</td>
        <td>IL</td>
    </tr>
```
Table 2
```html
<table>
    <tr>
        <td>Jane</td>
        <td>30</td>
        <td>Los Angeles</td>
        <td>CA</td>
    </tr>
```
Output:
```json
{
    "table1_description": "Table 1 has 4 headers, and 1 row.  The headers are Name, Age, City, and State.",
    "table2_description": "Table 2 has no headers, but the values appear to represent a person's name, age, city, and state.",
    "explanation": "The values in Table 2 match the headers in Table 1, and Table 2 has no headers. Table 2 should be merged to the bottom of Table 1.",
    "merge": "true",
    "direction": "bottom"
}
```
**Input:**
Table 1
```html
{{table1}}
Table 2
```html
{{table2}}
```
"""

    @staticmethod
    def get_row_count(cells: List[TableCell]):
        if not cells:
            return 0

        max_rows = None
        for col_id in set([cell.col_id for cell in cells]):
            col_cells = [cell for cell in cells if cell.col_id == col_id]
            rows = 0
            for cell in col_cells:
                rows += cell.rowspan
            if max_rows is None or rows > max_rows:
                max_rows = rows
        return max_rows

    @staticmethod
    def get_column_count(cells: List[TableCell]):
        if not cells:
            return 0

        max_cols = None
        for row_id in set([cell.row_id for cell in cells]):
            row_cells = [cell for cell in cells if cell.row_id == row_id]
            cols = 0
            for cell in row_cells:
                cols += cell.colspan
            if max_cols is None or cols > max_cols:
                max_cols = cols
        return max_cols

    def rewrite_blocks(self, document: Document):
        pbar = tqdm(desc=f"{self.__class__.__name__} running", disable=self.disable_tqdm)
        table_runs = []
        table_run = []
        prev_block = None
        prev_page_block_count = None
        for page in document.pages:
            page_blocks = page.contained_blocks(document, self.block_types)
            for block in page_blocks:
                merge_condition = False
                if prev_block is not None:
                    prev_cells = prev_block.contained_blocks(document, (BlockTypes.TableCell,))
                    curr_cells = block.contained_blocks(document, (BlockTypes.TableCell,))
                    row_match = abs(self.get_row_count(prev_cells) - self.get_row_count(curr_cells)) < 5, # Similar number of rows
                    col_match = abs(self.get_column_count(prev_cells) - self.get_column_count(curr_cells)) < 2

                    subsequent_page_table = all([
                        prev_block.page_id == block.page_id - 1, # Subsequent pages
                        max(prev_block.polygon.height / page.polygon.height,
                            block.polygon.height / page.polygon.height) > self.table_height_threshold, # Take up most of the page height
                            (len(page_blocks) == 1 or prev_page_block_count == 1), # Only table on the page
                            (row_match or col_match)
                        ])

                    same_page_vertical_table = all([
                        prev_block.page_id == block.page_id, # On the same page
                        (1 - self.vertical_table_height_threshold) < prev_block.polygon.height / block.polygon.height < (1 + self.vertical_table_height_threshold), # Similar height
                        abs(block.polygon.x_start - prev_block.polygon.x_end) < self.vertical_table_distance_threshold, # Close together in x
                        abs(block.polygon.y_start - prev_block.polygon.y_start) < self.vertical_table_distance_threshold, # Close together in y
                        row_match
                    ])

                    same_page_horizontal_table = all([
                        prev_block.page_id == block.page_id, # On the same page
                        (1 - self.horizontal_table_width_threshold) < prev_block.polygon.width / block.polygon.width < (1 + self.horizontal_table_width_threshold), # Similar width
                        abs(block.polygon.y_start - prev_block.polygon.y_end) < self.horizontal_table_distance_threshold, # Close together in y
                        abs(block.polygon.x_start - prev_block.polygon.x_start) < self.horizontal_table_distance_threshold, # Close together in x
                        col_match
                    ])

                    same_page_new_column = all([
                        prev_block.page_id == block.page_id, # On the same page
                        abs(block.polygon.x_start - prev_block.polygon.x_end) < self.column_gap_threshold,
                        block.polygon.y_start < prev_block.polygon.y_end,
                        block.polygon.width * (1 - self.vertical_table_height_threshold) < prev_block.polygon.width  < block.polygon.width * (1 + self.vertical_table_height_threshold), # Similar width
                        col_match
                    ])
                    merge_condition = any([subsequent_page_table, same_page_vertical_table, same_page_new_column, same_page_horizontal_table])

                if prev_block is not None and merge_condition:
                    if prev_block not in table_run:
                        table_run.append(prev_block)
                    table_run.append(block)
                else:
                    if table_run:
                        table_runs.append(table_run)
                    table_run = []
                prev_block = block
            prev_page_block_count = len(page_blocks)

        if table_run:
            table_runs.append(table_run)

        with ThreadPoolExecutor(max_workers=self.max_concurrency) as executor:
            for future in as_completed([
                executor.submit(self.process_rewriting, document, blocks)
                for blocks in table_runs
            ]):
                future.result()  # Raise exceptions if any occurred
                pbar.update(1)

        pbar.close()

    def process_rewriting(self, document: Document, blocks: List[Block]):
        if len(blocks) < 2:
            # Can't merge single tables
            return

        start_block = blocks[0]
        for i in range(1, len(blocks)):
            curr_block = blocks[i]
            children = start_block.contained_blocks(document, (BlockTypes.TableCell,))
            children_curr = curr_block.contained_blocks(document, (BlockTypes.TableCell,))
            if not children or not children_curr:
                # Happens if table/form processors didn't run
                break

            start_image = start_block.get_image(document, highres=False)
            curr_image = curr_block.get_image(document, highres=False)
            start_html = start_block.render(document).html
            curr_html = curr_block.render(document).html

            prompt = self.table_merge_prompt.replace("{{table1}}", start_html).replace("{{table2}}", curr_html)

            response = self.llm_service(
                prompt,
                [start_image, curr_image],
                curr_block,
                MergeSchema,
            )

            if not response or ("direction" not in response or "merge" not in response):
                curr_block.update_metadata(llm_error_count=1)
                break

            merge = response["merge"]

            # The original table is okay
            if "true" not in merge:
                start_block = curr_block
                continue

            # Merge the cells and images of the tables
            direction = response["direction"]
            if not self.validate_merge(children, children_curr, direction):
                start_block = curr_block
                continue

            merged_image = self.join_images(start_image, curr_image, direction)
            merged_cells = self.join_cells(children, children_curr, direction)
            curr_block.structure = []
            start_block.structure = [b.id for b in merged_cells]
            start_block.lowres_image = merged_image

    def validate_merge(self, cells1: List[TableCell], cells2: List[TableCell], direction: Literal['right', 'bottom'] = 'right'):
        if direction == "right":
            # Check if the number of rows is the same
            cells1_row_count = self.get_row_count(cells1)
            cells2_row_count = self.get_row_count(cells2)
            return abs(cells1_row_count - cells2_row_count) < 5
        elif direction == "bottom":
            # Check if the number of columns is the same
            cells1_col_count = self.get_column_count(cells1)
            cells2_col_count = self.get_column_count(cells2)
            return abs(cells1_col_count - cells2_col_count) < 2


    def join_cells(self, cells1: List[TableCell], cells2: List[TableCell], direction: Literal['right', 'bottom'] = 'right') -> List[TableCell]:
        if direction == 'right':
            # Shift columns right
            col_count = self.get_column_count(cells1)
            for cell in cells2:
                cell.col_id += col_count
            new_cells = cells1 + cells2
        else:
            # Shift rows up
            row_count = self.get_row_count(cells1)
            for cell in cells2:
                cell.row_id += row_count
            new_cells = cells1 + cells2
        return new_cells

    @staticmethod
    def join_images(image1: Image.Image, image2: Image.Image, direction: Literal['right', 'bottom'] = 'right') -> Image.Image:
        # Get dimensions
        w1, h1 = image1.size
        w2, h2 = image2.size

        if direction == 'right':
            new_height = max(h1, h2)
            new_width = w1 + w2
            new_img = Image.new('RGB', (new_width, new_height), 'white')
            new_img.paste(image1, (0, 0))
            new_img.paste(image2, (w1, 0))
        else:
            new_width = max(w1, w2)
            new_height = h1 + h2
            new_img = Image.new('RGB', (new_width, new_height), 'white')
            new_img.paste(image1, (0, 0))
            new_img.paste(image2, (0, h1))
        return new_img


class MergeSchema(BaseModel):
    table1_description: str
    table2_description: str
    explanation: str
    merge: Literal["true", "false"]
    direction: Literal["bottom", "right"]