"""Documentation loader for enhanced WhatsApp Flow generation.

This module provides functionality to load and combine documentation
with the base system prompt for improved LLM output quality.
"""

from __future__ import annotations

from functools import lru_cache
from pathlib import Path
from typing import Optional

# Bundled WhatsApp Flows documentation
# This contains the complete reference for WhatsApp Flows JSON schema
BUNDLED_DOCS = """
# WhatsApp Flows JSON Schema Documentation

## Overview

WhatsApp Flows allow businesses to create interactive, multi-screen forms and workflows within WhatsApp. This document covers schema versions **6.3**, **7.0**, **7.1**, and **7.2**.

---

## Schema Structure

```json
{
  "version": "7.0",
  "routing_model": {},  // Optional - defines screen navigation
  "screens": []
}
```

---

## Supported Versions

| Version | Features |
|---------|----------|
| 6.3 | Basic forms, Dropdown, Footer |
| 7.0 | Full component set, routing_model, Switch, If conditionals |
| 7.1 | Enhanced inputs, success screens |
| 7.2 | PhotoPicker improvements, DocumentPicker |

---

## Screen Structure

```json
{
  "id": "SCREEN_ID",
  "title": "Screen Title",
  "terminal": true,      // true if final screen
  "success": true,       // Optional - marks success state
  "data": {},            // Optional - screen data
  "layout": {
    "type": "SingleColumnLayout",
    "children": []
  }
}
```

### Screen Properties

| Property | Type | Description |
|----------|------|-------------|
| `id` | string | Unique screen identifier (UPPERCASE_SNAKE_CASE) |
| `title` | string | Display title |
| `terminal` | boolean | Marks the final screen in flow |
| `success` | boolean | Indicates successful completion |
| `data` | object | Screen-specific data |
| `layout` | object | UI layout configuration |

---

## Routing Model (v7.0+)

Defines navigation paths between screens:

```json
{
  "routing_model": {
    "SCREEN_ONE": ["SCREEN_TWO"],
    "SCREEN_TWO": ["SCREEN_THREE"],
    "SCREEN_THREE": []  // Empty = terminal/no next
  }
}
```

---

## Components Reference

### 1. Text Components

#### TextHeading
```json
{
  "type": "TextHeading",
  "text": "Welcome to Registration"
}
```

#### TextSubheading
```json
{
  "type": "TextSubheading",
  "text": "Please fill in your details"
}
```

#### TextBody
```json
{
  "type": "TextBody",
  "text": "Enter your information below",
  "font-weight": "bold"  // Optional
}
```

#### TextCaption
```json
{
  "type": "TextCaption",
  "text": "* Required fields"
}
```

---

### 2. Input Components

#### TextInput
```json
{
  "type": "TextInput",
  "label": "Full Name",
  "name": "user_name",
  "input-type": "text",
  "required": true,
  "helper-text": "Enter your full name",
  "min-chars": 2,
  "max-chars": 50
}
```

**Input Types:** `text`, `email`, `phone`, `number`, `password`

#### TextArea
```json
{
  "type": "TextArea",
  "name": "address",
  "label": "Residential Address",
  "required": true,
  "helper-text": "Enter complete address"
}
```

---

### 3. Selection Components

#### Dropdown
```json
{
  "type": "Dropdown",
  "label": "Select Country",
  "name": "country",
  "required": true,
  "data-source": [
    { "id": "us", "title": "United States" },
    { "id": "uk", "title": "United Kingdom" },
    { "id": "in", "title": "India" }
  ]
}
```

#### ChipsSelector
```json
{
  "type": "ChipsSelector",
  "name": "gender",
  "label": "Gender",
  "description": "Select your gender",
  "max-selected-items": 1,
  "required": true,
  "data-source": [
    { "id": "male", "title": "Male" },
    { "id": "female", "title": "Female" },
    { "id": "other", "title": "Other" }
  ]
}
```

#### RadioButtonsGroup
```json
{
  "type": "RadioButtonsGroup",
  "name": "category",
  "label": "Race Category",
  "required": true,
  "description": "Choose your race category",
  "data-source": [
    { "id": "sprint", "title": "Sprint" },
    { "id": "half", "title": "Half Marathon" },
    { "id": "olympic", "title": "Olympic" }
  ]
}
```

---

### 4. Date Picker

#### CalendarPicker
```json
{
  "type": "CalendarPicker",
  "name": "dob",
  "label": "Date of Birth",
  "helper-text": "Select your date of birth",
  "required": true,
  "mode": "single",
  "min-date": "1950-01-01",
  "max-date": "2009-12-31",
  "unavailable-dates": ["2024-12-25", "2024-01-01"]
}
```

---

### 5. Media Components

#### Image
```json
{
  "type": "Image",
  "src": "data:image/png;base64,iVBORw0KGgo...",
  "width": 200,
  "height": 200
}
```

#### PhotoPicker (v7.0+)
```json
{
  "type": "PhotoPicker",
  "name": "passport_photo",
  "label": "Passport Size Photo",
  "description": "Upload your passport size photo",
  "photo-source": "camera_gallery",
  "min-uploaded-photos": 1,
  "max-uploaded-photos": 1,
  "max-file-size-kb": 10240
}
```

#### DocumentPicker (v7.0+)
```json
{
  "type": "DocumentPicker",
  "name": "kyc_document",
  "label": "Identity Document",
  "description": "Upload Aadhar/Passport/ID Card",
  "min-uploaded-documents": 1,
  "max-uploaded-documents": 2,
  "max-file-size-kb": 5000,
  "allowed-mime-types": ["application/pdf", "image/jpeg", "image/png"]
}
```

---

### 6. Navigation Components

#### Footer
```json
{
  "type": "Footer",
  "label": "Next",
  "enabled": true,
  "on-click-action": {
    "name": "navigate",
    "next": {
      "type": "screen",
      "name": "NEXT_SCREEN_ID"
    },
    "payload": {
      "user_name": "${form.user_name}",
      "email": "${form.email}"
    }
  }
}
```

#### Footer with Complete Action (Terminal Screen)
```json
{
  "type": "Footer",
  "label": "Submit",
  "on-click-action": {
    "name": "complete",
    "payload": {
      "user_name": "${screen.SCREEN_ONE.form.user_name}",
      "email": "${form.email}",
      "submitted": "true"
    }
  }
}
```

#### Button
```json
{
  "type": "Button",
  "label": "Submit",
  "on-click-action": {
    "name": "complete",
    "payload": {
      "action": "submit"
    }
  }
}
```

#### EmbeddedLink
```json
{
  "type": "EmbeddedLink",
  "text": "View Terms & Conditions",
  "on-click-action": {
    "name": "open_url",
    "url": "https://example.com/terms"
  }
}
```

---

### 7. Form Container

```json
{
  "type": "Form",
  "name": "registration_form",
  "children": [
    // Form components here
  ]
}
```

---

### 8. Conditional Components

#### If Conditional
```json
{
  "type": "If",
  "condition": "${form.age} < 18",
  "then": [
    {
      "type": "TextBody",
      "text": "You must be 18+ to proceed"
    }
  ],
  "else": [
    {
      "type": "Footer",
      "label": "Continue",
      "on-click-action": {
        "name": "navigate",
        "next": { "type": "screen", "name": "NEXT_SCREEN" }
      }
    }
  ]
}
```

#### Switch
```json
{
  "type": "Switch",
  "value": "${form.category}",
  "cases": {
    "sprint": [
      { "type": "TextBody", "text": "Sprint: 750m swim, 5km run" }
    ],
    "half": [
      { "type": "TextBody", "text": "Half: 1900m swim, 21km run" }
    ],
    "default": [
      { "type": "TextBody", "text": "Select a category to continue" }
    ]
  }
}
```

---

## Data Reference Syntax

### Current Form
```
${form.field_name}
```

### Previous Screen
```
${screen.SCREEN_ID.form.field_name}
```

### Examples
```json
{
  "payload": {
    "name": "${form.user_name}",
    "dob": "${screen.SCREEN_ONE.form.dob}",
    "photo": "${screen.PHOTO_SCREEN.form.passport_photo}"
  }
}
```

---

## Best Practices

1. **Screen IDs**: Use UPPERCASE_SNAKE_CASE (e.g., `WELCOME_SCREEN`)
2. **Form Names**: Use lowercase_snake_case (e.g., `registration_form`)
3. **Field Names**: Use lowercase_snake_case (e.g., `user_name`)
4. **Terminal Screens**: Always set `terminal: true` on final screen
5. **Payload References**: Use `${form.field}` for current screen, `${screen.ID.form.field}` for previous screens
6. **Validation**: Mark fields as `required: true` for mandatory inputs
7. **Helper Text**: Provide `helper-text` for better UX
8. **File Sizes**: Keep `max-file-size-kb` reasonable (5-10MB)

---

## Action Types

| Action | Description | Usage |
|--------|-------------|-------|
| `navigate` | Move to next screen | Multi-screen flows |
| `complete` | End flow and submit | Terminal screens |
| `open_url` | Open external URL | EmbeddedLink |

---

## Screen Navigation: With vs Without Payload

### Navigation WITHOUT Payload

Use when the next screen doesn't need data from current screen, OR when you'll reference data directly from previous screens on the terminal screen.

```json
{
  "type": "Footer",
  "label": "Next",
  "on-click-action": {
    "name": "navigate",
    "next": {
      "type": "screen",
      "name": "NEXT_SCREEN"
    }
  }
}
```

### Navigation WITH Payload

Use when you need to explicitly pass data to the next screen for validation, conditional logic, or immediate use.

```json
{
  "type": "Footer",
  "label": "Next",
  "on-click-action": {
    "name": "navigate",
    "next": {
      "type": "screen",
      "name": "NEXT_SCREEN"
    },
    "payload": {
      "user_name": "${form.user_name}",
      "email": "${form.email}"
    }
  }
}
```

### Key Rules

1. **Terminal screens** must use `"name": "complete"` (not `navigate`)
2. **Payload is optional** for `navigate` but **required** for `complete`
3. **Reference syntax:**
   - Current form: `${form.field_name}`
   - Previous screen: `${screen.SCREEN_ID.form.field_name}`
4. **All fields in `complete` payload** are sent to your backend webhook

---

## Component Availability by Version

| Component | 6.3 | 7.0 | 7.1 | 7.2 |
|-----------|-----|-----|-----|-----|
| TextHeading | ✓ | ✓ | ✓ | ✓ |
| TextBody | ✓ | ✓ | ✓ | ✓ |
| TextInput | ✓ | ✓ | ✓ | ✓ |
| Dropdown | ✓ | ✓ | ✓ | ✓ |
| Footer | ✓ | ✓ | ✓ | ✓ |
| ChipsSelector | - | ✓ | ✓ | ✓ |
| CalendarPicker | - | ✓ | ✓ | ✓ |
| RadioButtonsGroup | - | ✓ | ✓ | ✓ |
| PhotoPicker | - | ✓ | ✓ | ✓ |
| DocumentPicker | - | ✓ | ✓ | ✓ |
| Switch | - | ✓ | ✓ | ✓ |
| If | - | ✓ | ✓ | ✓ |
| Image | - | ✓ | ✓ | ✓ |
| EmbeddedLink | - | ✓ | ✓ | ✓ |
| TextArea | - | ✓ | ✓ | ✓ |
| routing_model | - | ✓ | ✓ | ✓ |
"""


@lru_cache(maxsize=1)
def get_bundled_docs() -> str:
    """Get the bundled WhatsApp Flows documentation."""
    return BUNDLED_DOCS


def load_docs_from_file(file_path: str) -> str:
    """Load documentation from a file.

    Args:
        file_path: Path to the documentation file.

    Returns:
        The contents of the file, or empty string if file cannot be read.
    """
    try:
        path = Path(file_path)
        return path.read_text(encoding="utf-8")
    except Exception as exc:
        print(f"⚠️  Could not load documentation file: {exc}")
        return ""


def get_enhanced_prompt(
    base_prompt: str,
    *,
    use_bundled_docs: bool = True,
    custom_docs_path: Optional[str] = None,
    custom_docs: Optional[str] = None,
) -> str:
    """Create an enhanced system prompt with documentation.

    Args:
        base_prompt: The base system prompt.
        use_bundled_docs: Whether to include bundled documentation.
        custom_docs_path: Optional path to a custom documentation file.
        custom_docs: Optional custom documentation string.

    Returns:
        Enhanced prompt with documentation appended.
    """
    parts = [base_prompt]

    # Add bundled docs
    if use_bundled_docs:
        bundled = get_bundled_docs()
        if bundled:
            parts.append(bundled)

    # Add custom docs from file
    if custom_docs_path:
        file_docs = load_docs_from_file(custom_docs_path)
        if file_docs:
            parts.append("\n## CUSTOM DOCUMENTATION:\n\n" + file_docs)

    # Add custom docs string
    if custom_docs:
        parts.append("\n## ADDITIONAL CONTEXT:\n\n" + custom_docs)

    if len(parts) > 1:
        return "\n\n".join(parts) + (
            "\n\n---\n\n"
            "IMPORTANT: Use the above documentation as your primary reference. "
            "Follow all specifications, examples, and best practices exactly as documented."
        )

    return base_prompt
