#!/usr/bin/env python3
"""
Feishu Document Markdown Uploader
Usage: python feishu_md_uploader.py --app-id=<id> --app-secret=<secret> --file=<md_file> [--doc-id=<doc_id>] --title=<title>
"""

import json
import re
import requests
import sys
import os
import time
import argparse

MAX_TABLE_ROWS = 8  # Feishu limits tables to ~8 rows

TOKEN_URL = "https://open.feishu.cn/open-apis/auth/v3/tenant_access_token/internal"
DOC_BLOCKS_URL = "https://open.feishu.cn/open-apis/docx/v1/documents/{doc_token}/blocks"
DOC_BLOCK_CHILDREN_URL = "https://open.feishu.cn/open-apis/docx/v1/documents/{doc_token}/blocks/{block_id}/children"


def get_token(app_id, app_secret):
    resp = requests.post(TOKEN_URL, json={"app_id": app_id, "app_secret": app_secret})
    data = resp.json()
    if data.get("code") != 0:
        raise Exception(f"Token error: {data}")
    return data["tenant_access_token"]


def get_doc_root_block(doc_token, token):
    resp = requests.get(
        DOC_BLOCKS_URL.format(doc_token=doc_token),
        headers={"Authorization": f"Bearer {token}"}
    )
    blocks = resp.json()["data"]["items"]
    return blocks[0]["block_id"]


def write_block(doc_token, block_id, block_data, token):
    url = DOC_BLOCK_CHILDREN_URL.format(doc_token=doc_token, block_id=block_id)
    resp = requests.post(
        url,
        json=block_data,
        headers={"Authorization": f"Bearer {token}", "Content-Type": "application/json"}
    )
    result = resp.json()
    if result.get("code") != 0:
        print(f"Warning: {result}")
        return result
    return result


def write_heading(doc_token, content, token, block_id, level=1):
    block_type = 3 + level - 1
    block_data = {
        "children": [{
            "block_type": block_type,
            f"heading{level}": {
                "elements": [{"type": 2, "text_run": {"content": content}}]
            }
        }]
    }
    return write_block(doc_token, block_id, block_data, token)


def parse_inline_formatting(text):
    """
    Parse markdown inline formatting and return Feishu elements
    Supports: **bold**, *italic*, `code`, ~~strikethrough~~
    """
    elements = []
    
    # Pattern to match **bold**, *italic*, `code`, ~~strikethrough~~
    pattern = r'(\*\*|__)(.*?)\1|(\*|_)(.*?)\3|(`)(.*?)`|~~(.*?)~~'
    
    last_end = 0
    for match in re.finditer(pattern, text):
        # Add text before match
        if match.start() > last_end:
            elements.append({
                "type": 2,
                "text_run": {"content": text[last_end:match.start()]}
            })
        
        matched = match.group(0)
        style = {}
        content = ""
        
        if matched.startswith('**') or matched.startswith('__'):
            content = match.group(2)
            style = {"bold": True}
        elif matched.startswith('*') or matched.startswith('_'):
            content = match.group(4)
            style = {"italic": True}
        elif matched.startswith('`'):
            content = match.group(6)
            style = {"inline_code": True}
        elif matched.startswith('~~'):
            content = match.group(7)
            style = {"strikethrough": True}
        
        element = {"type": 2, "text_run": {"content": content}}
        if style:
            element["text_run"]["text_element_style"] = style
        
        elements.append(element)
        last_end = match.end()
    
    # Add remaining text
    if last_end < len(text):
        elements.append({
            "type": 2,
            "text_run": {"content": text[last_end:]}
        })
    
    # If no formatting found, return simple element
    if not elements:
        elements.append({
            "type": 2,
            "text_run": {"content": text}
        })
    
    return elements


def write_paragraph(doc_token, content, token, block_id):
    elements = parse_inline_formatting(content)
    
    block_data = {
        "children": [{
            "block_type": 2,
            "text": {
                "elements": elements
            }
        }]
    }
    return write_block(doc_token, block_id, block_data, token)


def write_table(doc_token, table_data, token, block_id):
    """Create and fill a table in Feishu document, splitting if too large"""
    rows = len(table_data)
    cols = len(table_data[0]) if table_data else 0
    
    if rows == 0 or cols == 0:
        return None
    
    # Calculate column widths once for the entire table
    col_widths = _calculate_column_widths(table_data)
    print(f"📐 Table column widths: {col_widths}", flush=True, file=sys.stderr)
    
    # If table is too large, split into multiple tables with header
    if rows > MAX_TABLE_ROWS:
        print(f"⚠️ Table too large ({rows}x{cols}), splitting into chunks...")
        header = table_data[0]  # First row is header
        data_rows = table_data[1:]  # Rest is data
        
        # First chunk: header + first max_rows of data
        first_chunk = [header] + data_rows[:MAX_TABLE_ROWS]
        _write_table_single(doc_token, first_chunk, token, block_id, col_widths)
        time.sleep(0.1)
        
        # Remaining chunks: insert empty text block, then table WITHOUT header
        for i in range(MAX_TABLE_ROWS, len(data_rows), MAX_TABLE_ROWS):
            chunk = data_rows[i:i + MAX_TABLE_ROWS]
            # Insert empty text block to remove dividing line
            write_paragraph(doc_token, "", token, block_id)
            time.sleep(0.1)
            # Insert next chunk WITHOUT header, use same column widths
            _write_table_single(doc_token, chunk, token, block_id, col_widths)
            time.sleep(0.1)
        
        return None
    
    return _write_table_single(doc_token, table_data, token, block_id, col_widths)


def _calculate_column_widths(table_data, total_width=700, threshold=100, default_width=20):
    """Calculate column widths based on content length (considering Chinese vs English)"""
    if not table_data or not table_data[0]:
        return None
    
    cols = len(table_data[0])
    
    # Calculate pixel width for each column (all rows including header)
    # Start with default width + character-based width
    # Chinese: 16px per char, English/digit: 10px per char
    col_max_px = []
    for col_idx in range(cols):
        max_px = default_width  # Start with default width
        for row in table_data:
            if col_idx < len(row) and row[col_idx]:
                text = str(row[col_idx])
                char_px = sum(16 if '\u4e00' <= c <= '\u9fff' else 10 for c in text)
                max_px = max(max_px, default_width + char_px)
        col_max_px.append(max_px)
    
    total_px = sum(col_max_px)
    
    # If total width < 700px, use actual widths
    if total_px <= total_width:
        return col_max_px
    
    # Total exceeds 700, apply the rules
    # Rule (a): columns with max width < threshold keep their width
    # Rule (b): columns with max width >= threshold share remaining width proportionally
    
    small_cols_px = sum(c for c in col_max_px if c < threshold)
    large_cols = [c for c in col_max_px if c >= threshold]
    large_cols_total = sum(large_cols)
    
    if large_cols_total == 0:
        # All columns are small, use equal distribution
        return [total_width // cols] * cols
    
    remaining_width = total_width - small_cols_px
    
    result = []
    for c in col_max_px:
        if c < threshold:
            result.append(c)
        else:
            # Proportionally distribute remaining width
            width = int(remaining_width * c / large_cols_total)
            result.append(width)
    
    return result


def _write_table_single(doc_token, table_data, token, block_id, col_widths=None):
    """Create and fill a single table (internal)"""
    rows = len(table_data)
    cols = len(table_data[0]) if table_data else 0
    
    if rows == 0 or cols == 0:
        return None
    
    # Use provided column widths or calculate if not provided
    if col_widths is None:
        col_widths = _calculate_column_widths(table_data)
    
    # Create table
    block_data = {
        "children": [{
            "block_type": 31,
            "table": {
                "rows": rows,
                "columns": cols,
                "property": {"row_size": rows, "column_size": cols, "column_width": col_widths}
            }
        }]
    }
    
    result = write_block(doc_token, block_id, block_data, token)
    
    if result.get("code") != 0:
        print(f"Warning: Failed to create table: {result}")
        return None
    
    # Get cell block IDs from children
    table_block = result["data"]["children"][0]
    cell_block_ids = table_block["children"]
    
    # Fill table data
    for i, row in enumerate(table_data):
        for j, val in enumerate(row):
            cell_idx = i * cols + j
            if cell_idx >= len(cell_block_ids):
                break
            cell_block_id = cell_block_ids[cell_idx]
            
            cell_data = {
                "children": [{
                    "block_type": 2,
                    "text": {
                        "elements": [{"type": 2, "text_run": {"content": str(val)}}]
                    }
                }]
            }
            
            cell_url = DOC_BLOCK_CHILDREN_URL.format(doc_token=doc_token, block_id=cell_block_id)
            requests.post(
                cell_url,
                json=cell_data,
                headers={"Authorization": f"Bearer {token}", "Content-Type": "application/json"}
            )
            time.sleep(0.05)
    
    return result


def markdown_to_feishu(doc_token, content, token):
    root_block_id = get_doc_root_block(doc_token, token)
    
    lines = content.split("\n")
    i = 0
    
    while i < len(lines):
        line = lines[i].rstrip()
        
        if not line.strip():
            i += 1
            continue
        
        # Heading
        if line.startswith("# "):
            write_heading(doc_token, line[2:], token, root_block_id, 1)
        elif line.startswith("## "):
            write_heading(doc_token, line[3:], token, root_block_id, 2)
        elif line.startswith("### "):
            write_heading(doc_token, line[4:], token, root_block_id, 3)
        # Divider
        elif re.match(r'^---+$', line) or re.match(r'^\*\*\*+$', line):
            i += 1
            continue
        # Unordered list
        elif line.startswith("- ") or line.startswith("* "):
            write_paragraph(doc_token, "• " + line[2:], token, root_block_id)
        # Ordered list
        elif re.match(r'^\d+\. ', line):
            write_paragraph(doc_token, line, token, root_block_id)
        # Table
        elif line.startswith("|"):
            table_rows, i = parse_table(lines, i)
            if table_rows:
                write_table(doc_token, table_rows, token, root_block_id)
            continue
        # Quote
        elif line.startswith("> "):
            write_paragraph(doc_token, "❝ " + line[2:], token, root_block_id)
        # Code block
        elif line.startswith("```"):
            lang = line[3:].strip() if len(line) > 3 else ""
            code_lines = []
            i += 1
            while i < len(lines) and not lines[i].startswith("```"):
                code_lines.append(lines[i])
                i += 1
            write_paragraph(doc_token, "```" + lang + "\n" + "\n".join(code_lines) + "\n```", token, root_block_id)
        # Regular paragraph
        else:
            write_paragraph(doc_token, line, token, root_block_id)
        
        i += 1
        time.sleep(0.05)


def parse_table(lines, start_idx):
    """Parse markdown table - validates column count"""
    rows = []
    i = start_idx
    expected_cols = None
    
    while i < len(lines):
        line = lines[i].strip()
        if not line.startswith("|"):
            break
        # Skip separator line (e.g., |---|)
        if re.match(r'\|[\s\-:]+\|', line):
            i += 1
            continue
        
        cells = [c.strip() for c in line.split("|")[1:-1]]
        
        # Validate column count
        if expected_cols is None:
            expected_cols = len(cells)
        elif len(cells) != expected_cols:
            raise ValueError(f"Table row {len(rows)+1} has {len(cells)} columns, expected {expected_cols}")
        
        rows.append(cells)
        i += 1
    
    return rows, i


def create_doc(title, token):
    resp = requests.post(
        "https://open.feishu.cn/open-apis/docx/v1/documents",
        headers={"Authorization": f"Bearer {token}", "Content-Type": "application/json"},
        json={"title": title}
    )
    data = resp.json()
    
    if data.get("code") != 0:
        raise Exception(f"创建文档失败: {data}")
    
    return data["data"]["document"]["document_id"]


def main():
    parser = argparse.ArgumentParser(description="Upload Markdown to Feishu Document")
    parser.add_argument("--app-id", required=True, help="Feishu App ID")
    parser.add_argument("--app-secret", required=True, help="Feishu App Secret")
    parser.add_argument("--file", required=True, help="Path to markdown file")
    parser.add_argument("--title", required=True, help="Document title")
    parser.add_argument("--doc-id", help="Existing document ID (optional, creates new if not provided)")
    
    args = parser.parse_args()
    
    # Read file
    with open(args.file, "r", encoding="utf-8") as f:
        content = f.read()
    
    print(f"Read file: {args.file}, length: {len(content)}")
    
    # Get token
    print("Getting token...")
    token = get_token(args.app_id, args.app_secret)
    print(f"Token: {token[:20]}...")
    
    # Determine doc_id
    if args.doc_id:
        doc_id = args.doc_id
        print(f"Using existing doc: {doc_id}")
    else:
        print("Creating new document...")
        doc_id = create_doc(args.title, token)
        print(f"Created doc: {doc_id}")
    
    # Write content
    print("Writing content...")
    markdown_to_feishu(doc_id, content, token)
    print("Content written!")
    
    print(f"\n✅ Success!")
    print(f"📄 Document: https://pingcap.feishu.cn/docx/{doc_id}")


if __name__ == "__main__":
    main()