#!/usr/bin/env python3
"""
Chat message formatter — converts JSON chat data to Markdown or HTML.

Usage:
    python chat_formatter.py -f html chat1.json chat2.json -o output.html
    python chat_formatter.py -f markdown chat1.json chat2.json -o output.md
"""

import json
import sys
from pathlib import Path

import click
from jinja2 import FileSystemLoader, Environment

template_dir = Path(__file__).parent / "templates"
file_loader = FileSystemLoader(template_dir)
env = Environment(loader=file_loader)
HTML_TEMPLATE = env.get_template("chat.html")


def load_messages(filepaths: list[str]):
    """Load and merge messages from one or more JSON files."""
    all_messages = []
    for fp in filepaths:
        path = Path(fp)
        if not path.exists():
            click.echo(f"Warning: {fp} not found, skipping.", err=True)
            continue
        try:
            with path.open("r", encoding="utf-8") as f:
                data = json.load(f)
        except json.JSONDecodeError as exc:
            click.echo(f"Error parsing {fp}: {exc}", err=True)
            continue
        if not isinstance(data, list):
            click.echo(f"Warning: {fp} is not a JSON array, skipping.", err=True)
            continue
        all_messages.extend(data)
    return all_messages


def normalize_images(images):
    """Ensure each image entry is a dict with a 'url' key."""
    result = []
    if not images:
        return result
    for img in images:
        if isinstance(img, str):
            result.append({"url": img})
        elif isinstance(img, dict):
            url = img.get("url") or img.get("src") or ""
            if url:
                result.append({"url": url})
    return result


def escape_html(text):
    return (
        text.replace("&", "&")
        .replace("<", "&lt;")
        .replace(">", "&gt;")
        .replace('"', "&quot;")
        .replace("'", "&#39;")
    )


def prepare_messages(messages):
    """Pre-process messages for template rendering."""
    prepared = []
    for msg in messages:
        prepared.append(
            {
                "role": msg.get("role", "unknown"),
                "content_esc": escape_html(msg.get("content", "")),
                "images": normalize_images(msg.get("images", [])),
            }
        )
    return prepared


def format_markdown(messages):
    """Render messages as Markdown text."""
    lines = []
    for msg in messages:
        role = msg.get("role", "unknown")
        content = msg.get("content", "")
        images = msg.get("images", [])

        role_label = "👤 User" if role == "user" else "🤖 Assistant"
        lines.append(f"**{role_label}**\n")

        if content:
            lines.append(content)
            lines.append("")

        for i, img in enumerate(images):
            if isinstance(img, dict):
                url = img.get("url") or img.get("src") or ""
            else:
                url = str(img)
            if url:
                lines.append(f"![Image {i + 1}]({url})")
        if images:
            lines.append("")

    return "\n".join(lines)


def format_html(messages, title="Chat History"):
    """Render messages as a standalone HTML document via Jinja2."""
    prepared = prepare_messages(messages)
    return HTML_TEMPLATE.render(title=title, messages=prepared)


@click.command()
@click.argument("files", nargs=-1, type=click.Path())
@click.option(
    "-f",
    "--format",
    type=click.Choice(["html", "markdown"], case_sensitive=False),
    default="markdown",
    show_default=True,
    help="Output format.",
)
@click.option(
    "-o",
    "--output",
    type=click.Path(),
    default=None,
    help="Output file path. Defaults to stdout.",
)
@click.option(
    "--title",
    default="Chat History",
    show_default=True,
    help="Title for the HTML document.",
)
def main(files, format, output, title):
    """Format JSON chat logs as Markdown or HTML."""
    if not files:
        click.echo("Error: no input files provided.", err=True)
        sys.exit(1)

    messages = load_messages(files)
    if not messages:
        click.echo("Error: no valid messages loaded.", err=True)
        sys.exit(1)

    fmt = format.lower()

    if fmt == "html":
        result = format_html(messages, title=title)
    else:
        result = format_markdown(messages)

    if output:
        _ = Path(output).write_text(result, encoding="utf-8")
        click.echo(f"Wrote {len(messages)} messages to {output}")
    else:
        click.echo(result)


if __name__ == "__main__":
    main()