import os
from dotenv import load_dotenv
load_dotenv()

import json
import pandas as pd

import zipfile
import xml.etree.ElementTree as ET
from io import BytesIO
import openpyxl

from openai import OpenAI

import re

import logging


HF_API_KEY = os.getenv("HF_API_KEY")

# Configure logging to write to 'zaoju_logs.log' without using pickle
logging.basicConfig(
    filename='extract_po_logs.log',
    level=logging.INFO, 
    format='%(asctime)s - %(levelname)s - %(message)s',
    encoding='utf-8'
)

# Default Word XML namespace
DEFAULT_NS = {'w': 'http://schemas.openxmlformats.org/wordprocessingml/2006/main'}
NS = None  # Global variable to store the namespace

def get_namespace(root):
    """Extracts the primary namespace from the XML root element while keeping the default."""
    global NS

    ns = root.tag.split('}')[0].strip('{')
    NS = {'w': ns} if ns else DEFAULT_NS

    return NS

# --- Helper Functions for DOCX Processing ---

def extract_text_from_cell(cell):
    """Extracts text from a Word table cell, preserving line breaks and reconstructing split words."""
    paragraphs = cell.findall('.//w:p', NS)
    lines = []
    
    for paragraph in paragraphs:
        # Get all text runs and concatenate their contents
        text_runs = [t.text for t in paragraph.findall('.//w:t', NS) if t.text]
        line = ''.join(text_runs).strip()  # Merge split words properly
        
        if line:  # Add only non-empty lines
            lines.append(line)
    
    return lines  # Return list of lines to preserve line breaks

def clean_spaces(text):
    """
    Removes excessive spaces between Chinese characters while preserving spaces in English words.
    """
    # Remove spaces **between** Chinese characters but keep English spaces
    text = re.sub(r'([\u4e00-\u9fff])\s+([\u4e00-\u9fff])', r'\1\2', text)
    return text.strip()

def extract_key_value_pairs(text, target_dict=None):
    """
    Extracts multiple key-value pairs from a given text.
    - First, split by more than 3 spaces (`\s{3,}`) **only if the next segment contains a `:`.**
    - Then, process each segment by splitting at `:` to correctly assign keys and values.
    """
    if target_dict is None:
        target_dict = {}

    text = text.replace("：", ":")  # Normalize Chinese colons to English colons

    # Step 1: Check if splitting by more than 3 spaces is necessary
    segments = re.split(r'(\s{3,})', text)  # Use raw string to prevent invalid escape sequence

    # Step 2: Process each segment, ensuring we only split if the next part has a `:`
    merged_segments = []
    temp_segment = ""

    for segment in segments:
        if ":" in segment:  # If segment contains `:`, it's a valid split point
            if temp_segment:
                merged_segments.append(temp_segment.strip())
                temp_segment = ""
            merged_segments.append(segment.strip())
        else:
            temp_segment += " " + segment.strip()

    if temp_segment:
        merged_segments.append(temp_segment.strip())

    # Step 3: Extract key-value pairs correctly
    for segment in merged_segments:
        if ':' in segment:
            key, value = segment.split(':', 1)  # Only split at the first colon
            key, value = key.strip(), value.strip()  # Clean spaces

            if key in target_dict:
                target_dict[key] += "\n" + value  # Append if key already exists
            else:
                target_dict[key] = value

    return target_dict

# --- Table Processing Functions ---

def process_single_column_table(rows):
    """Processes a single-column table and returns the extracted lines as a list."""
    single_column_data = []

    for row in rows:
        cells = row.findall('.//w:tc', NS)
        if len(cells) == 1:
            cell_lines = extract_text_from_cell(cells[0])  # Extract all lines from the cell
            
            # Append each line directly to the list without splitting
            single_column_data.extend(cell_lines)

    return single_column_data  # Return the list of extracted lines

def process_buyer_seller_table(rows):
    """Processes a two-column buyer-seller table into a structured dictionary using the first row as keys."""
    headers = [extract_text_from_cell(cell) for cell in rows[0].findall('.//w:tc', NS)]
    if len(headers) != 2:
        return None  # Not a buyer-seller table

    # determine role based on header text
    def get_role(header_text, default_role):
        header_text = header_text.lower()  # Convert to lowercase
        if '买方' in header_text or 'buyer' in header_text or '甲方' in header_text:
            return 'buyer_info'
        elif '卖方' in header_text or 'seller' in header_text or '乙方' in header_text:
            return 'seller_info'
        else:
            return default_role  # Default if no keyword is found

    # Determine the keys for buyer and seller columns
    buyer_key = get_role(headers[0][0], 'buyer_info')
    seller_key = get_role(headers[1][0], 'seller_info')

    # Initialize the dictionary using the determined keys
    buyer_seller_data = {
        buyer_key: {},
        seller_key: {}
    }

    for row in rows:
        cells = row.findall('.//w:tc', NS)
        if len(cells) == 2:
            buyer_lines = extract_text_from_cell(cells[0])
            seller_lines = extract_text_from_cell(cells[1])

        for line in buyer_lines:
            extract_key_value_pairs(line, buyer_seller_data[buyer_key])

        for line in seller_lines:
            extract_key_value_pairs(line, buyer_seller_data[seller_key])
            
    return buyer_seller_data

def process_summary_table(rows):
    """Processes a two-column summary table where keys are extracted as dictionary keys."""
    extracted_data = []
    
    for row in rows:
        cells = row.findall('.//w:tc', NS)
        if len(cells) == 2:
            key = " ".join(extract_text_from_cell(cells[0]))
            value = " ".join(extract_text_from_cell(cells[1]))
            extracted_data.append({key: value})
    
    return extracted_data

def extract_headers(first_row_cells):
    """Extracts unique column headers from the first row of a table."""
    headers = []
    header_count = {}
    for cell in first_row_cells:
        cell_text = " ".join(extract_text_from_cell(cell))
        grid_span = cell.find('.//w:gridSpan', NS)
        col_span = int(grid_span.attrib.get(f'{{{NS["w"]}}}val', '1')) if grid_span is not None else 1
        for _ in range(col_span):
            # Ensure header uniqueness by appending an index if repeated
            if cell_text in header_count:
                header_count[cell_text] += 1
                unique_header = f"{cell_text}_{header_count[cell_text]}"
            else:
                header_count[cell_text] = 1
                unique_header = cell_text
            headers.append(unique_header if unique_header else f"Column_{len(headers) + 1}")
    return headers

def process_long_table(rows):
    """Processes a standard table and correctly handles horizontally merged cells."""
    if not rows:
        return []  # Avoid IndexError
    
    headers = extract_headers(rows[0].findall('.//w:tc', NS))
    table_data = []
    vertical_merge_tracker = {}

    for row in rows[1:]:
        row_data = {}
        cells = row.findall('.//w:tc', NS)
        running_index = 0
        
        for cell in cells:
            cell_text = " ".join(extract_text_from_cell(cell))
            
            # Consistent Namespace Handling for Horizontal Merge
            grid_span = cell.find('.//w:gridSpan', NS)
            grid_span_val = grid_span.attrib.get(f'{{{NS["w"]}}}val') if grid_span is not None else '1'
            col_span = int(grid_span_val)
            
            # Handle vertical merge
            v_merge = cell.find('.//w:vMerge', NS)
            if v_merge is not None:
                v_merge_val = v_merge.attrib.get(f'{{{NS["w"]}}}val')
                if v_merge_val == 'restart':
                    vertical_merge_tracker[running_index] = cell_text
                else:
                    # Repeat the value from the previous row's merged cell
                    cell_text = vertical_merge_tracker.get(running_index, "")
            
            # Repeat the value for horizontally merged cells
            start_col = running_index
            end_col = running_index + col_span
            
            # Repeat the value for each spanned column
            for col in range(start_col, end_col):
                key = headers[col] if col < len(headers) else f"Column_{col+1}"
                row_data[key] = cell_text
            
            # Update the running index to the end of the merged cell
            running_index = end_col
        
        # Fill remaining columns with empty strings to maintain alignment
        while running_index < len(headers):
            row_data[headers[running_index]] = ""
            running_index += 1
        
        table_data.append(row_data)
    
    return table_data

def extract_tables(root):
    """Extracts tables from the DOCX document and returns structured data."""
    tables = root.findall('.//w:tbl', NS)
    table_data = {}
    table_paragraphs = set()
    
    for table_index, table in enumerate(tables, start=1):
        rows = table.findall('.//w:tr', NS)
        if not rows:
            continue  # Skip empty tables
        
        for paragraph in table.findall('.//w:p', NS):
            table_paragraphs.add(paragraph)
        
        first_row_cells = rows[0].findall('.//w:tc', NS)
        num_columns = len(first_row_cells)

        if num_columns == 1:
            single_column_data = process_single_column_table(rows)
            if single_column_data:
                table_data[f"table_{table_index}_single_column"] = single_column_data
            continue  # Skip further processing for this table

        summary_start_index = None
        for i, row in enumerate(rows):
            if len(row.findall('.//w:tc', NS)) == 2:
                summary_start_index = i
                break

        long_table_data = []
        summary_data = []

        if summary_start_index is not None and summary_start_index > 0:
            long_table_data = process_long_table(rows[:summary_start_index])
        elif summary_start_index is None:
            long_table_data = process_long_table(rows)

        if summary_start_index is not None:
            is_buyer_seller_table = all(len(row.findall('.//w:tc', NS)) == 2 for row in rows)
            if is_buyer_seller_table:
                buyer_seller_data = process_buyer_seller_table(rows)
                if buyer_seller_data:
                    table_data[f"table_{table_index}_buyer_seller"] = buyer_seller_data
            else:
                summary_data = process_summary_table(rows[summary_start_index:])
    
        if long_table_data:
            table_data[f"long_table_{table_index}"] = long_table_data
        if summary_data:
            table_data[f"long_table_{table_index}_summary"] = summary_data
    
    return table_data, table_paragraphs

# --- Non-Table Processing Functions ---

def extract_text_outside_tables(root, table_paragraphs):
    """Extracts text from paragraphs outside tables in the document."""
    extracted_text = []
    
    for paragraph in root.findall('.//w:p', NS):
        if paragraph in table_paragraphs:
            continue  # Skip paragraphs inside tables
        
        texts = [t.text.strip() for t in paragraph.findall('.//w:t', NS) if t.text]
        line = clean_spaces(' '.join(texts).replace('：',':'))  # Clean colons and spaces
        
        if ':' in line:
            extracted_text.append(line)
    
    return extracted_text

# --- Main Extraction Functions ---

def extract_docx_as_xml(file_bytes, save_xml=False, xml_filename="document.xml"):

    # Ensure file_bytes is at the start position
    file_bytes.seek(0)

    with zipfile.ZipFile(file_bytes, 'r') as docx:
        with docx.open('word/document.xml') as xml_file:
            xml_content = xml_file.read().decode('utf-8')
            if save_xml:
                with open(xml_filename, "w", encoding="utf-8") as f:
                    f.write(xml_content)
            return xml_content


def xml_to_json(xml_content, save_json=False, json_filename="extracted_data.json"):
    
    tree = ET.ElementTree(ET.fromstring(xml_content))
    root = tree.getroot()
    
    table_data, table_paragraphs = extract_tables(root)
    extracted_data = table_data
    extracted_data["non_table_data"] = extract_text_outside_tables(root, table_paragraphs)
    
    if save_json:
        with open(json_filename, "w", encoding="utf-8") as f:
            json.dump(extracted_data, f, ensure_ascii=False, indent=4)
    
    return json.dumps(extracted_data, ensure_ascii=False, indent=4)


def deepseek_extract_contract_summary(json_data, save_json=False, json_filename="contract_summary.json"):
    """Sends extracted JSON data to OpenAI and returns formatted structured JSON."""
    
    # Step 1: Convert JSON string to Python dictionary
    contract_data = json.loads(json_data)

    # Step 2: Remove keys that contain "long_table"
    filtered_contract_data = {key: value for key, value in contract_data.items() if "long_table" not in key}

    # Step 3: Convert back to JSON string (if needed)
    json_output = json.dumps(contract_data, ensure_ascii=False, indent=4)
    
    prompt = """You are given a contract in JSON format. Extract the following information:

# Response Format
Return the extracted information as a structured JSON in the exact format shown below (Note: Do not repeat any keys, if unsure leave the value empty):

{
    "合同编号":
    "接收人": （注意：不是买家必须是接收人，不是一个公司而是一个人）
    "Recipient":
    "接收地": （注意：不是交货地点是目的港，只写中文，英文写在 place of receipt）
    "Place of receipt": （只写英文, 如果接收地/目的港/Port of destination 有英文可填在这里）
    "供应商":
    "币种": （主要用的货币，填英文缩写。GNF一般是为了方便而转换出来的, 除非只有GNF，GNF一般不是主要币种。）
    "供货日期": （如果合同里有写才填，不要自己推理出日期，必须是一个日期，而不是天数）
}

Contract data in JSON format:""" + f"""
{json_output}"""

    messages = [
        {
            "role": "user",
            "content": prompt
        }
    ]
    
    # Deepseek R1 Distilled Qwen 2.5 14B --------------------------------
    client = OpenAI(
        base_url="https://router.huggingface.co/novita",
        api_key=HF_API_KEY,
    )
    
    completion = client.chat.completions.create(
        model="deepseek/deepseek-r1-distill-qwen-14b", 
        messages=messages, 
        temperature=0.5,
    )
    
    # Deepseek V3 --------------------------------
    # client = OpenAI(
    #     base_url="https://router.huggingface.co/novita",
    #     api_key=HF_API_KEY,
    # )
    
    # completion = client.chat.completions.create(
    #     model="deepseek/deepseek_v3", 
    #     messages=messages, 
    #     temperature=0.1,
    # )
    
    # Qwen 2.5 7B --------------------------------
    # client = OpenAI(
    #     base_url="https://router.huggingface.co/together",
    #     api_key=HF_API_KEY,
    # )
    
    # completion = client.chat.completions.create(
    #     model="Qwen/Qwen2.5-7B-Instruct-Turbo", 
    #     messages=messages, 
    # )
    
    think_text = re.findall(r"<think>(.*?)</think>", completion.choices[0].message.content, flags=re.DOTALL)
    if think_text:
        print(f"Thought Process: {think_text}")
        logging.info(f"Think text: {think_text}")
        
    contract_summary = re.sub(r"<think>.*?</think>\s*", "", completion.choices[0].message.content, flags=re.DOTALL) # Remove think
    
    contract_summary = re.sub(r"^```json\n|```$", "", contract_summary, flags=re.DOTALL) # Remove ```
    
    if save_json:
        with open(json_filename, "w", encoding="utf-8") as f:
            f.write(contract_summary)
    
    return json.dumps(contract_summary, ensure_ascii=False, indent=4)


def deepseek_extract_price_list(json_data):
    """Sends extracted JSON data to OpenAI and returns formatted structured JSON."""
    
    # Step 1: Convert JSON string to Python dictionary
    contract_data = json.loads(json_data)

    # Step 2: Remove keys that contain "long_table"
    filtered_contract_data = {key: value for key, value in contract_data.items() if "long_table" in key}

    # Step 3: Convert back to JSON string (if needed)
    json_output = json.dumps(filtered_contract_data, ensure_ascii=False, indent=4)
    
    prompt = """You are given a price list in JSON format. Extract the following information in CSV format:

# Response Format
Return the extracted information as a CSV in the exact format shown below:

物料名称, 物料名称(英文), 物料规格, 采购数量, 单位, 单价, 计划号

JSON data:""" + f"""
{json_output}"""

    messages = [
        {
            "role": "user",
            "content": prompt
        }
    ]

    client = OpenAI(
	    base_url="https://router.huggingface.co/novita",
	    api_key=HF_API_KEY,
    )
    
    completion = client.chat.completions.create(
        model="deepseek/deepseek-r1-distill-qwen-14b", 
        messages=messages, 
    )
    
    price_list = re.sub(r"<think>.*?</think>\s*", "", completion.choices[0].message.content, flags=re.DOTALL)
    
    price_list = re.sub(r"^```json\n|```$", "", price_list, flags=re.DOTALL)
    

def json_to_excel(contract_summary, json_data, excel_path):
    """Converts extracted JSON tables to an Excel file."""

    # Correctly parse the JSON string
    contract_summary_json = json.loads(json.loads(contract_summary))
    
    contract_summary_df = pd.DataFrame([contract_summary_json])
    
    # Ensure json_data is a dictionary
    if isinstance(json_data, str):
        json_data = json.loads(json_data)
    
    long_tables = [pd.DataFrame(table) for key, table in json_data.items() if "long_table" in key and "summary" not in key]
    long_table = long_tables[-1] if long_tables else pd.DataFrame()

    with pd.ExcelWriter(excel_path) as writer:
        contract_summary_df.to_excel(writer, sheet_name="Contract Summary", index=False)
        long_table.to_excel(writer, sheet_name="Price List", index=False)

#--- Extract PO ------------------------------

def extract_po(docx_path):
    """Processes a single .docx file, extracts tables, formats with OpenAI, and saves as an Excel file."""
    if not os.path.exists(docx_path) or not docx_path.endswith(".docx"):
        raise ValueError(f"Invalid file: {docx_path}")
    
    # Read the .docx file as bytes
    with open(docx_path, "rb") as f:
        docx_bytes = BytesIO(f.read())
        
    # Step 1: Extract XML content from DOCX
    print("Extracting Docs data to XML...")
    xml_filename = os.path.splitext(os.path.basename(docx_path))[0] + "_document.xml"
    xml_file = extract_docx_as_xml(docx_bytes, save_xml=True, xml_filename=xml_filename)
    
    get_namespace(ET.fromstring(xml_file))
    
    # Step 2: Extract tables from DOCX and save JSON
    print("Extracting XML data to JSON...")
    json_filename = os.path.splitext(os.path.basename(docx_path))[0] + "_extracted_data.json"
    extracted_data = xml_to_json(xml_file, save_json=True, json_filename=json_filename)

    # Step 2: Process JSON with OpenAI to get structured output
    print("Processing JSON data with AI...")
    contract_summary_filename = os.path.splitext(os.path.basename(docx_path))[0] + "_contract_summary.json"
    contract_summary = deepseek_extract_contract_summary(extracted_data, save_json=True, json_filename=contract_summary_filename)

    # Step 3: Save formatted data as Excel
    print("Converting AI Generated JSON to Excel...")
    excel_output_path = os.path.splitext(docx_path)[0] + ".xlsx"
    json_to_excel(contract_summary, extracted_data, excel_output_path)

    print(f"Excel file saved at: {excel_output_path}")
    
    
    # Logging
    log = f"""Results:
    
    Contract Summary: {contract_summary},
    
    RAW Extracted Data: {extracted_data},
    
    XML Preview: {xml_file[:1000]}"""
    
    print(log)
    
    logging.info(f"""{log}""")
    
    
    return excel_output_path

# Example Usage

# extract_po("test-contract-converted.docx")
# extract_po("test-contract.docx")

# Gradio Interface ------------------------------

import gradio as gr
from gradio.themes.base import Base

interface = gr.Interface(
    fn=extract_po,
    title="PO Extractor 买卖合同数据提取", 
    inputs=gr.File(label="买卖合同 （.docx）"),
    outputs=gr.File(label="数据提取结果 （.xlsx）"),
    flagging_mode="never",
    theme=Base()
)

interface.launch()