import json
from connector_base import call_google_api

DOCS_GET_DOCUMENT_TOOL = {
    "type": "function",
    "function": {
        "name": "docs_get_document",
        "description": "Retrieve full document content, paragraphs, headings, tables, and document metadata of a Google Doc by ID.",
        "parameters": {
            "type": "object",
            "properties": {
                "document_id": {
                    "type": "string",
                    "description": "The ID of the document."
                }
            },
            "required": ["document_id"]
        }
    }
}

DOCS_CREATE_DOCUMENT_TOOL = {
    "type": "function",
    "function": {
        "name": "docs_create_document",
        "description": "Create a new Google Doc with specified title and optional initial text.",
        "parameters": {
            "type": "object",
            "properties": {
                "title": {
                    "type": "string",
                    "description": "Title of the new Google Doc."
                },
                "initial_text": {
                    "type": "string",
                    "description": "Optional starting text to insert into document."
                }
            },
            "required": ["title"]
        }
    }
}

DOCS_APPEND_TEXT_TOOL = {
    "type": "function",
    "function": {
        "name": "docs_append_text",
        "description": "Append text or paragraphs to the end of an existing Google Doc.",
        "parameters": {
            "type": "object",
            "properties": {
                "document_id": {
                    "type": "string",
                    "description": "The ID of the Google Doc."
                },
                "text": {
                    "type": "string",
                    "description": "Text content to insert at the end of the document."
                }
            },
            "required": ["document_id", "text"]
        }
    }
}

DOCS_BATCH_UPDATE_TOOL = {
    "type": "function",
    "function": {
        "name": "docs_batch_update",
        "description": "Apply structural updates to a Google Doc (insertText, deleteContentRange, replaceAllText, updateTextStyle).",
        "parameters": {
            "type": "object",
            "properties": {
                "document_id": {
                    "type": "string",
                    "description": "The ID of the document."
                },
                "requests": {
                    "type": "array",
                    "items": {"type": "object"},
                    "description": "List of update Request objects following Google Docs API schema."
                }
            },
            "required": ["document_id", "requests"]
        }
    }
}

DOCS_FORMAT_TEXT_TOOL = {
    "type": "function",
    "function": {
        "name": "docs_format_text",
        "description": "Apply bold, italic, underline, or heading style to a specific range of text in a Google Doc.",
        "parameters": {
            "type": "object",
            "properties": {
                "document_id": {
                    "type": "string",
                    "description": "The ID of the document."
                },
                "start_index": {
                    "type": "integer",
                    "description": "The start index of the text range."
                },
                "end_index": {
                    "type": "integer",
                    "description": "The end index of the text range."
                },
                "bold": {
                    "type": "boolean",
                    "description": "Apply bold formatting."
                },
                "italic": {
                    "type": "boolean",
                    "description": "Apply italic formatting."
                },
                "underline": {
                    "type": "boolean",
                    "description": "Apply underline formatting."
                },
                "heading": {
                    "type": "string",
                    "enum": ["1", "2", "3", "4", "5", "6"],
                    "description": "Apply heading style (1 to 6)."
                }
            },
            "required": ["document_id", "start_index", "end_index"]
        }
    }
}

DOCS_CREATE_LIST_TOOL = {
    "type": "function",
    "function": {
        "name": "docs_create_list",
        "description": "Create a bulleted or numbered list within a specific range of a Google Doc.",
        "parameters": {
            "type": "object",
            "properties": {
                "document_id": {
                    "type": "string",
                    "description": "The ID of the document."
                },
                "start_index": {
                    "type": "integer",
                    "description": "The start index of the range."
                },
                "end_index": {
                    "type": "integer",
                    "description": "The end index of the range."
                },
                "list_type": {
                    "type": "string",
                    "enum": ["bullet", "numbered"],
                    "description": "Type of list to create."
                }
            },
            "required": ["document_id", "start_index", "end_index", "list_type"]
        }
    }
}

DOCS_ADD_COMMENT_TOOL = {
    "type": "function",
    "function": {
        "name": "docs_add_comment",
        "description": "Add a comment to a Google Doc using the Drive API.",
        "parameters": {
            "type": "object",
            "properties": {
                "document_id": {
                    "type": "string",
                    "description": "The ID of the document (also known as file ID)."
                },
                "text": {
                    "type": "string",
                    "description": "The comment text content."
                }
            },
            "required": ["document_id", "text"]
        }
    }
}

DOCS_SHARE_DOCUMENT_TOOL = {
    "type": "function",
    "function": {
        "name": "docs_share_document",
        "description": "Share a Google Doc with a specific user email and role using the Drive API.",
        "parameters": {
            "type": "object",
            "properties": {
                "document_id": {
                    "type": "string",
                    "description": "The ID of the document."
                },
                "email": {
                    "type": "string",
                    "description": "The email address of the user to share with."
                },
                "role": {
                    "type": "string",
                    "enum": ["reader", "writer", "commenter"],
                    "description": "The permission role."
                }
            },
            "required": ["document_id", "email"]
        }
    }
}

DOCS_GET_DOCUMENT_STRUCTURE_TOOL = {
    "type": "function",
    "function": {
        "name": "docs_get_document_structure",
        "description": "Parse and return a structured summary of headings and tables from a Google Doc for easier AI comprehension.",
        "parameters": {
            "type": "object",
            "properties": {
                "document_id": {
                    "type": "string",
                    "description": "The ID of the document."
                }
            },
            "required": ["document_id"]
        }
    }
}

ALL_DOCS_MCP_TOOLS = [
    DOCS_GET_DOCUMENT_TOOL,
    DOCS_CREATE_DOCUMENT_TOOL,
    DOCS_APPEND_TEXT_TOOL,
    DOCS_BATCH_UPDATE_TOOL,
    DOCS_FORMAT_TEXT_TOOL,
    DOCS_CREATE_LIST_TOOL,
    DOCS_ADD_COMMENT_TOOL,
    DOCS_SHARE_DOCUMENT_TOOL,
    DOCS_GET_DOCUMENT_STRUCTURE_TOOL
]

def get_docs_metadata(func_name: str, args: dict):
    if func_name == "docs_get_document":
        return "docs", "Getting document content"
    elif func_name == "docs_create_document":
        title = args.get("title", "")
        return "docs", f"Creating document: {title}" if title else "Creating document"
    elif func_name == "docs_append_text":
        return "docs", "Appending text to document"
    elif func_name == "docs_batch_update":
        return "docs", "Batch updating document"
    elif func_name == "docs_format_text":
        return "docs", "Formatting text in document"
    elif func_name == "docs_create_list":
        return "docs", "Creating list in document"
    elif func_name == "docs_add_comment":
        return "docs", "Adding comment to document"
    elif func_name == "docs_share_document":
        return "docs", "Sharing document"
    elif func_name == "docs_get_document_structure":
        return "docs", "Getting document structure"
    return "docs", func_name

async def execute_docs_tool(user_id: str, fn_name: str, args: dict) -> str:
    base_url = "https://docs.googleapis.com/v1/documents"
    drive_url = "https://www.googleapis.com/drive/v3/files"

    if fn_name == "docs_get_document":
        doc_id = args.get("document_id")
        status, body = await call_google_api(user_id, "docs", "GET", f"{base_url}/{doc_id}")
        return body

    elif fn_name == "docs_create_document":
        title = args.get("title", "Untitled Document")
        status, body = await call_google_api(user_id, "docs", "POST", base_url, json_body={"title": title})
        if status not in [200, 201]:
            return body

        initial_text = args.get("initial_text")
        if initial_text:
            try:
                d = json.loads(body)
                created_id = d.get("documentId")
                if created_id:
                    ins_payload = {
                        "requests": [
                            {
                                "insertText": {
                                    "location": {"index": 1},
                                    "text": initial_text
                                }
                            }
                        ]
                    }
                    await call_google_api(user_id, "docs", "POST", f"{base_url}/{created_id}:batchUpdate", json_body=ins_payload)
            except Exception:
                pass

        return body

    elif fn_name == "docs_append_text":
        doc_id = args.get("document_id")
        text = args.get("text", "")
        payload = {
            "requests": [
                {
                    "insertText": {
                        "endOfSegmentLocation": {},
                        "text": text
                    }
                }
            ]
        }
        status, body = await call_google_api(user_id, "docs", "POST", f"{base_url}/{doc_id}:batchUpdate", json_body=payload)
        return body

    elif fn_name == "docs_batch_update":
        doc_id = args.get("document_id")
        payload = {"requests": args.get("requests", [])}
        status, body = await call_google_api(user_id, "docs", "POST", f"{base_url}/{doc_id}:batchUpdate", json_body=payload)
        return body

    elif fn_name == "docs_format_text":
        doc_id = args.get("document_id")
        start_idx = args.get("start_index")
        end_idx = args.get("end_index")
        bold = args.get("bold", False)
        italic = args.get("italic", False)
        underline = args.get("underline", False)
        heading = args.get("heading")
        requests = []
        if heading:
            style_type = f"HEADING_{heading}" if heading.isdigit() and 1 <= int(heading) <= 6 else "NORMAL_TEXT"
            requests.append({
                "updateParagraphStyle": {
                    "range": {"startIndex": start_idx, "endIndex": end_idx},
                    "paragraphStyle": {"namedStyleType": style_type},
                    "fields": "namedStyleType"
                }
            })
        if bold or italic or underline:
            style = {}
            if bold:
                style["bold"] = True
            if italic:
                style["italic"] = True
            if underline:
                style["underline"] = True
            requests.append({
                "updateTextStyle": {
                    "range": {"startIndex": start_idx, "endIndex": end_idx},
                    "textStyle": style,
                    "fields": ",".join(style.keys())
                }
            })
        payload = {"requests": requests}
        status, body = await call_google_api(user_id, "docs", "POST", f"{base_url}/{doc_id}:batchUpdate", json_body=payload)
        return body

    elif fn_name == "docs_create_list":
        doc_id = args.get("document_id")
        start_idx = args.get("start_index")
        end_idx = args.get("end_index")
        list_type = args.get("list_type", "bullet")
        bullet_preset = "NUMBERED" if list_type == "numbered" else "BULLET"
        payload = {
            "requests": [
                {
                    "createParagraphBullets": {
                        "range": {"startIndex": start_idx, "endIndex": end_idx},
                        "bulletPreset": bullet_preset
                    }
                }
            ]
        }
        status, body = await call_google_api(user_id, "docs", "POST", f"{base_url}/{doc_id}:batchUpdate", json_body=payload)
        return body

    elif fn_name == "docs_add_comment":
        doc_id = args.get("document_id")
        text = args.get("text", "")
        comment_payload = {"content": text}
        status, body = await call_google_api(user_id, "drive", "POST", f"{drive_url}/{doc_id}/comments", json_body=comment_payload)
        return body

    elif fn_name == "docs_share_document":
        doc_id = args.get("document_id")
        email = args.get("email")
        role = args.get("role", "reader")
        permission_payload = {
            "type": "user",
            "role": role,
            "emailAddress": email
        }
        status, body = await call_google_api(user_id, "drive", "POST", f"{drive_url}/{doc_id}/permissions", json_body=permission_payload)
        return body

    elif fn_name == "docs_get_document_structure":
        doc_id = args.get("document_id")
        status, body = await call_google_api(user_id, "docs", "GET", f"{base_url}/{doc_id}")
        if status not in [200, 201]:
            return body
        try:
            data = json.loads(body)
            content = data.get("body", {}).get("content", [])
            headings = []
            tables = []
            for item in content:
                if "paragraph" in item:
                    para = item["paragraph"]
                    elements = para.get("elements", [])
                    if elements:
                        text_run = elements[0].get("textRun", {})
                        text = text_run.get("content", "").strip()
                        if text:
                            style = para.get("paragraphStyle", {})
                            named_style = style.get("namedStyleType", "")
                            if named_style.startswith("HEADING_"):
                                level = named_style.replace("HEADING_", "")
                                if level.isdigit():
                                    headings.append({"level": int(level), "text": text})
                elif "table" in item:
                    table_data = item["table"]
                    rows = table_data.get("tableRows", [])
                    tables.append({"row_count": len(rows)})
            structured = {"headings": headings, "tables": tables}
            return json.dumps(structured)
        except Exception:
            return body

    return json.dumps({"error": f"Unsupported Docs tool: {fn_name}"})