import json
import os
import re
import time
import zipfile
import math
from pathlib import Path

from typing import Any, Dict, List, Optional, Union
from collections import Counter
import xml.etree.ElementTree as ET
from pandas import Timestamp
from datetime import datetime
from pandas.api.types import is_datetime64_any_dtype

import pandas as pd
from tabulate import tabulate
from qwen_agent.log import logger
from qwen_agent.settings import DEFAULT_WORKSPACE, DEFAULT_MAX_INPUT_TOKENS
from qwen_agent.tools.base import BaseTool, register_tool
from qwen_agent.tools.storage import KeyNotExistsError, Storage
from file_tools.utils import (get_file_type, hash_sha256, is_http_url, get_basename_from_url, 
                                  sanitize_chrome_file_path, save_url_to_local_work_dir)
from qwen_agent.utils.tokenization_qwen import count_tokens, tokenizer
from file_tools.idp import IDP

# Configuration constants
PARSER_SUPPORTED_FILE_TYPES = ['pdf', 'docx', 'pptx', 'txt', 'html', 'csv', 'tsv', 'xlsx', 'xls', 'doc', 'zip', '.mp4', '.mov', '.mkv', '.webm', '.mp3', '.wav']
def str_to_bool(value):
    """Convert string to boolean, handling common true/false representations"""
    if isinstance(value, bool):
        return value
    return str(value).lower() in ('true', '1', 'yes', 'on')
USE_IDP = str_to_bool(os.getenv("USE_IDP", "True"))
IDP_TIMEOUT = 150000
ENABLE_CSI = False
PARAGRAPH_SPLIT_SYMBOL = '\n'


class CustomJSONEncoder(json.JSONEncoder):
    def default(self, obj):
        if isinstance(obj, (datetime, Timestamp)):
            return obj.isoformat()
        return super().default(obj)


class FileParserError(Exception):
    """Custom exception for document parsing errors"""

    def __init__(self, message: str, code: str = '400', exception: Optional[Exception] = None):
        super().__init__(message)
        self.code = code
        self.exception = exception


def parse_file_by_idp(file_path: str = None, file_url: str = None) -> List[dict]:
    idp = IDP()
    try:
        fid = idp.file_submit_with_url(file_url) if file_url else idp.file_submit_with_path(file_path)
        if not fid:
            return []

        for _ in range(10):  
            result, status = idp.file_parser_query(fid)
            if status == 'success':
                return process_idp_result(result)
            time.sleep(10)

        logger.error("IDP parsing timeout")
        return []
    except Exception as e:
        logger.error(f"IDP processing failed: {str(e)}")
        return []


def process_idp_result(result: dict) -> List[dict]:
    pages = []
    current_page = None

    for layout in result.get('layouts', []):
        page_num = layout.get('pageNum', 0)
        content = layout.get('markdownContent', '')

        if current_page and current_page['page_num'] == page_num:
            current_page['content'].append({'text': content})
        else:
            current_page = {'page_num': page_num, 'content': [{'text': content}]}
            pages.append(current_page)

    return pages


def clean_text(text: str) -> str:
    cleaners = [
        lambda x: re.sub(r'\n+', '\n', x),  
        lambda x: x.replace("Add to Qwen's Reading List", ''),
        lambda x: re.sub(r'-{6,}', '-----', x),  
        lambda x: x.strip()
    ]
    for cleaner in cleaners:
        text = cleaner(text)
    return text


def get_plain_doc(doc: list):
    paras = []
    for page in doc:
        for para in page['content']:
            for k, v in para.items():
                if k in ['text', 'table', 'image']:
                    paras.append(v)
    return PARAGRAPH_SPLIT_SYMBOL.join(paras)


def df_to_markdown(df: pd.DataFrame) -> str:
    df = df.dropna(how='all').fillna('')
    return tabulate(df, headers='keys', tablefmt='pipe', showindex=False)


def parse_word(docx_path: str, extract_image: bool = False):
    if extract_image:
        raise ValueError('Currently, extracting images is not supported!')

    from docx import Document
    doc = Document(docx_path)

    content = []
    for para in doc.paragraphs:
        content.append({'text': para.text})
    for table in doc.tables:
        tbl = []
        for row in table.rows:
            tbl.append('|' + '|'.join([cell.text for cell in row.cells]) + '|')
        tbl = '\n'.join(tbl)
        content.append({'table': tbl})
    return [{'page_num': 1, 'content': content}]


def parse_ppt(path: str, extract_image: bool = False):
    if extract_image:
        raise ValueError('Currently, extracting images is not supported!')

    from pptx import Presentation
    from pptx.exc import PackageNotFoundError
    try:
        ppt = Presentation(path)
    except PackageNotFoundError as ex:
        logger.warning(ex)
        return []
    doc = []
    for slide_number, slide in enumerate(ppt.slides):
        page = {'page_num': slide_number + 1, 'content': []}

        for shape in slide.shapes:
            if not shape.has_text_frame and not shape.has_table:
                pass

            if shape.has_text_frame:
                for paragraph in shape.text_frame.paragraphs:
                    paragraph_text = ''.join(run.text for run in paragraph.runs)
                    paragraph_text = clean_text(paragraph_text)
                    if paragraph_text.strip():
                        page['content'].append({'text': paragraph_text})

            if shape.has_table:
                tbl = []
                for row_number, row in enumerate(shape.table.rows):
                    tbl.append('|' + '|'.join([cell.text for cell in row.cells]) + '|')
                tbl = '\n'.join(tbl)
                page['content'].append({'table': tbl})
        doc.append(page)
    return doc


def parse_pdf(pdf_path: str, extract_image: bool = False) -> List[dict]:
    # Todo: header and footer
    from pdfminer.high_level import extract_pages
    from pdfminer.layout import LTImage, LTRect, LTTextContainer

    doc = []
    import pdfplumber
    pdf = pdfplumber.open(pdf_path)
    for i, page_layout in enumerate(extract_pages(pdf_path)):
        page = {'page_num': page_layout.pageid, 'content': []}

        elements = []
        for element in page_layout:
            elements.append(element)

        # Init params for table
        table_num = 0
        tables = []

        for element in elements:
            if isinstance(element, LTRect):
                if not tables:
                    tables = extract_tables(pdf, i)
                if table_num < len(tables):
                    table_string = table_converter(tables[table_num])
                    table_num += 1
                    if table_string:
                        page['content'].append({'table': table_string, 'obj': element})
            elif isinstance(element, LTTextContainer):
                # Delete line breaks in the same paragraph
                text = element.get_text()
                # Todo: Further analysis using font
                font = get_font(element)
                if text.strip():
                    new_content_item = {'text': text, 'obj': element}
                    if font:
                        new_content_item['font-size'] = round(font[1])
                        # new_content_item['font-name'] = font[0]
                    page['content'].append(new_content_item)
            elif extract_image and isinstance(element, LTImage):
                # Todo: ocr
                raise ValueError('Currently, extracting images is not supported!')
            else:
                pass

        # merge elements
        page['content'] = postprocess_page_content(page['content'])
        doc.append(page)

    return doc


def parse_txt(path: str):
    with open(path, 'r', encoding='utf-8') as f:  
        text = f.read()
    paras = text.split(PARAGRAPH_SPLIT_SYMBOL)
    content = []
    for p in paras:
        content.append({'text': p})
    return [{'page_num': 1, 'content': content}]


def get_font(element):
    from pdfminer.layout import LTChar, LTTextContainer

    fonts_list = []
    for text_line in element:
        if isinstance(text_line, LTTextContainer):
            for character in text_line:
                if isinstance(character, LTChar):
                    fonts_list.append((character.fontname, character.size))

    fonts_list = list(set(fonts_list))
    if fonts_list:
        counter = Counter(fonts_list)
        most_common_fonts = counter.most_common(1)[0][0]
        return most_common_fonts
    else:
        return []


def extract_tables(pdf, page_num):
    table_page = pdf.pages[page_num]
    tables = table_page.extract_tables()
    return tables


def table_converter(table):
    table_string = ''
    for row_num in range(len(table)):
        row = table[row_num]
        cleaned_row = [
            item.replace('\n', ' ') if item is not None and '\n' in item else 'None' if item is None else item
            for item in row
        ]
        table_string += ('|' + '|'.join(cleaned_row) + '|' + '\n')
    table_string = table_string[:-1]
    return table_string


def postprocess_page_content(page_content: list) -> list:
    # rm repetitive identification for table and text
    # Some documents may repeatedly recognize LTRect and LTTextContainer
    table_obj = [p['obj'] for p in page_content if 'table' in p]
    tmp = []
    for p in page_content:
        repetitive = False
        if 'text' in p:
            for t in table_obj:
                if t.bbox[0] <= p['obj'].bbox[0] and p['obj'].bbox[1] <= t.bbox[1] and t.bbox[2] <= p['obj'].bbox[
                    2] and p['obj'].bbox[3] <= t.bbox[3]:
                    repetitive = True
                    break

        if not repetitive:
            tmp.append(p)
    page_content = tmp

    # merge paragraphs that have been separated by mistake
    new_page_content = []
    for p in page_content:
        if new_page_content and 'text' in new_page_content[-1] and 'text' in p and abs(
                p.get('font-size', 12) -
                new_page_content[-1].get('font-size', 12)) < 2 and p['obj'].height < p.get('font-size', 12) + 1:
            # Merge those lines belonging to a paragraph
            new_page_content[-1]['text'] += f' {p["text"]}'
            # new_page_content[-1]['font-name'] = p.get('font-name', '')
            new_page_content[-1]['font-size'] = p.get('font-size', 12)
        else:
            p.pop('obj')
            new_page_content.append(p)
    for i in range(len(new_page_content)):
        if 'text' in new_page_content[i]:
            new_page_content[i]['text'] = clean_text(new_page_content[i]['text'])
    return new_page_content


def extract_xls_schema(file_path: str) -> Dict[str, Any]:
    xls = pd.ExcelFile(file_path)
    schema = {
        "sheets": [],
        "n_sheets": len(xls.sheet_names)
    }

    for sheet_name in xls.sheet_names:
        df = xls.parse(sheet_name, nrows=3)  # 读取前3行

        dtype_mapping = {
            'object': 'string',
            'datetime64[ns]': 'datetime',
            'timedelta64[ns]': 'timedelta'
        }
        dtypes = df.dtypes.astype(str).replace(dtype_mapping).to_dict()

        sample_df = df.head(3).copy()
        for col in sample_df.columns:
            if is_datetime64_any_dtype(sample_df[col]):
                sample_df[col] = sample_df[col].dt.strftime('%Y-%m-%dT%H:%M:%S')

        sheet_info = {
            "name": sheet_name,
            "columns": df.columns.tolist(),
            "dtypes": dtypes,  
            "sample_data": sample_df.to_dict(orient='list') 
        }
        schema["sheets"].append(sheet_info)

    return schema


def extract_csv_schema(file_path: str) -> Dict[str, Any]:
    df_dtype = pd.read_csv(file_path, nrows=100)  
    df_sample = pd.read_csv(file_path, nrows=3) 

    return {
        "columns": df_dtype.columns.tolist(),
        "dtypes": df_dtype.dtypes.astype(str).to_dict(),
        "sample_data": df_sample.to_dict(orient='list'),
        "estimated_total_rows": _estimate_total_rows(file_path)
    }


def _estimate_total_rows(file_path) -> int:
    with open(file_path, 'rb') as f:
        line_count = 0
        chunk_size = 1024 * 1024  
        while chunk := f.read(chunk_size):
            line_count += chunk.count(b'\n')
    return line_count - 1  


def parse_tabular_file(file_path: str, **kwargs) -> List[dict]:
    try:
        df = pd.read_excel(file_path) if file_path.endswith(('.xlsx', '.xls')) else \
            pd.read_csv(file_path, **kwargs)
        if count_tokens(df_to_markdown(df)) > DEFAULT_MAX_INPUT_TOKENS:
            schema = extract_xls_schema(file_path) if file_path.endswith(('.xlsx', '.xls')) else \
                extract_csv_schema(file_path)
            return [{'page_num': 1, 'content': [{'schema': schema}]}]
        else:
            return [{'page_num': 1, 'content': [{'table': df_to_markdown(df)}]}]
    except Exception as e:
        logger.error(f"Table parsing failed: {str(e)}")
        return []  


def parse_zip(file_path: str, extract_dir: str) -> List[dict]:
    with zipfile.ZipFile(file_path, 'r') as zip_ref:
        zip_ref.extractall(extract_dir)
        return [os.path.join(extract_dir, f) for f in zip_ref.namelist()]


def parse_html(file_path: str) -> List[dict]:
    from bs4 import BeautifulSoup  

    with open(file_path, 'r', encoding='utf-8') as f:
        soup = BeautifulSoup(f, 'lxml')

    content = [{'text': clean_text(p.get_text())}
               for p in soup.find_all(['p', 'div']) if p.get_text().strip()]

    return [{
        'page_num': 1,
        'content': content,
        'title': soup.title.string if soup.title else ''
    }]


def extract_xml_skeleton_markdown(xml_file):
    tree = ET.parse(xml_file)
    root = tree.getroot()
    markdown_lines = []

    def process_element(element, level=0, parent_path="", is_last=True, prefix=""):
        if level > 0:
            connector = "└── " if is_last else "├── "
            markdown_lines.append(f"{prefix}{connector}**{element.tag}**")
        else:
            markdown_lines.append(f"## Root: {element.tag}")

        if element.attrib:
            attrs = [f"`{k}`" for k in element.attrib.keys()]
            attr_line = f"{prefix}{'    ' if level > 0 else ''}*Attributes:* {', '.join(attrs)}"
            markdown_lines.append(attr_line)

        if element.text and element.text.strip():
            text_line = f"{prefix}{'    ' if level > 0 else ''}*Has text content*"
            markdown_lines.append(text_line)
        seen_tags = set()
        unique_children = []
        for child in element:
            if child.tag not in seen_tags:
                seen_tags.add(child.tag)
                unique_children.append(child)

        for i, child in enumerate(unique_children):
            is_last_child = (i == len(unique_children) - 1)
            child_prefix = prefix + ("    " if is_last else "│   ")
            process_element(child, level + 1,
                            f"{parent_path}/{element.tag}" if parent_path else element.tag,
                            is_last_child, child_prefix)

    process_element(root)
    markdown_content = "\n".join(markdown_lines)
    return markdown_content


def parse_xml(file_path: str) -> List[dict]:
    with open(file_path, 'r', encoding='utf-8') as f: 
        text = f.read()
    if count_tokens(text) > DEFAULT_MAX_INPUT_TOKENS:
        schema = extract_xml_skeleton_markdown(file_path)
        content = [{'schema': schema}]
    else:
        content = [{'text': text}]
    return [{'page_num': 1, 'content': content}]


def compress(results: list) -> list[str]:
    compress_results = []
    max_token = math.floor(DEFAULT_MAX_INPUT_TOKENS / len(results))
    for result in results:
        token_list = tokenizer.tokenize(result)
        token_list = token_list[:min(len(token_list), max_token)]
        compress_results.append(tokenizer.convert_tokens_to_string(token_list))
    return compress_results


# @register_tool('file_parser')
class SingleFileParser(BaseTool):
    name="file_parser"
    description = f"File parsing tool, supports parsing data in  {'/'.join(PARSER_SUPPORTED_FILE_TYPES)} formats, and returns the parsed markdown format data."
    parameters = [{
        'name': 'url',
        'type': 'string',
        'description': 'The full path of the file to be parsed, which can be a local path or a downloadable http(s) link.',
        'required': True
    }]

    def __init__(self, cfg: Optional[Dict] = None):
        super().__init__(cfg)
        self.data_root = self.cfg.get('path', os.path.join(DEFAULT_WORKSPACE, 'tools', self.name))
        self.db = Storage({'storage_root_path': self.data_root})
        self.structured_doc = self.cfg.get('structured_doc', True)

  
        self.parsers = {
            'pdf': parse_pdf,
            'docx': parse_word,
            'doc': parse_word,
            'pptx': parse_ppt,
            'txt': parse_txt,
            'jsonl': parse_txt,
            'jsonld': parse_txt,
            'pdb': parse_txt,
            'py': parse_txt,
            'html': parse_html,
            'xml': parse_xml,
            'csv': lambda p: parse_tabular_file(p, sep=','),
            'tsv': lambda p: parse_tabular_file(p, sep='\t'),
            'xlsx': parse_tabular_file,
            'xls': parse_tabular_file,
            'zip': self.parse_zip
        }

    def call(self, params: Union[str, dict], **kwargs) -> Union[str, list]:
        params = self._verify_json_format_args(params)
        file_path = self._prepare_file(params['url'])
        try:
            cached = self.db.get(f'{hash_sha256(file_path)}_ori')
            return self._flatten_result(json.loads(cached))
        except KeyNotExistsError:
            return self._flatten_result(self._process_new_file(file_path))

    def _prepare_file(self, path: str) -> str:
        if is_http_url(path):
            download_dir = os.path.join(self.data_root, hash_sha256(path))
            os.makedirs(download_dir, exist_ok=True)
            return save_url_to_local_work_dir(path, download_dir)
        return sanitize_chrome_file_path(path)

    def _process_new_file(self, file_path: str) -> Union[str, list]:
        file_type = get_file_type(file_path)
        idp_types = ['pdf', 'docx', 'pptx', 'xlsx', 'jpg', 'png', 'mp3']
        logger.info(f'Start parsing {file_path}...')
        logger.info(f'File type {file_type}...')
        logger.info(f"structured_doc {self.cfg.get('structured_doc')}...")

        if file_type not in idp_types:
            file_type = get_basename_from_url(file_path).split('.')[-1].lower()

        try:
            if USE_IDP and file_type in idp_types:
                try:
                    results = parse_file_by_idp(file_path=file_path)
                except Exception as e:
                    results = self.parsers[file_type](file_path)
            else:
                results = self.parsers[file_type](file_path)
            tokens = 0
            for page in results:
                for para in page['content']:
                    if 'schema' in para:
                        para['token'] = count_tokens(json.dumps(para['schema']))
                    else:
                        para['token'] = count_tokens(para.get('text', para.get('table')))
                    tokens += para['token']

            if not results or not tokens:
                logger.error(f"Parsing failed: No information was parsed")
                raise FileParserError("Document parsing failed")
            else:
                self._cache_result(file_path, results)
                return results
        except Exception as e:
            logger.error(f"Parsing failed: {str(e)}")
            raise FileParserError("Document parsing failed", exception=e)

    def _cache_result(self, file_path: str, result: list):
        cache_key = f'{hash_sha256(file_path)}_ori'
        self.db.put(cache_key, json.dumps(result, ensure_ascii=False))
        logger.info(f'The parsing result of {file_path} has been cached')

    def _flatten_result(self, result: list) -> str:
        return PARAGRAPH_SPLIT_SYMBOL.join(
            para.get('text', para.get('table', ''))
            for page in result for para in page['content']
        )

    def parse_zip(self, file_path: str) -> List[dict]:
        extract_dir = os.path.join(self.data_root, f"zip_{hash_sha256(file_path)}")
        os.makedirs(extract_dir, exist_ok=True)

        results = []
        for extracted_file in parse_zip(file_path, extract_dir):
            if (ft := get_file_type(extracted_file)) in self.parsers:
                try:
                    results.extend(self.parsers[ft](extracted_file))
                except Exception as e:
                    logger.warning(f"Skip files {extracted_file}: {str(e)}")

        if not results:
            raise ValueError("No parseable content found in the ZIP file")
        return results