"""Output exporters — save summaries as DOCX, JSON, or CSV."""

from __future__ import annotations

import json
import os
from pathlib import Path
from typing import Union

import pandas as pd
from docx import Document

from res_sum.schemas import StructuredSummary


def save_as_docx(
    summary: Union[str, StructuredSummary],
    output_path: str | Path,
    title: str = "Summary",
) -> None:
    """Save a summary as a DOCX file.

    Args:
        summary: Summary text or StructuredSummary object.
        output_path: Path for the output file.
        title: Document title.
    """
    doc = Document()
    doc.add_heading(title, level=0)

    if isinstance(summary, StructuredSummary):
        if summary.research_question:
            doc.add_heading("Research Question", level=1)
            doc.add_paragraph(summary.research_question)
        if summary.methods:
            doc.add_heading("Methods", level=1)
            doc.add_paragraph(summary.methods)
        if summary.key_findings:
            doc.add_heading("Key Findings", level=1)
            for finding in summary.key_findings:
                doc.add_paragraph(finding, style="List Bullet")
        if summary.ecological_implications:
            doc.add_heading("Ecological Implications", level=1)
            doc.add_paragraph(summary.ecological_implications)
        if summary.limitations:
            doc.add_heading("Limitations", level=1)
            doc.add_paragraph(summary.limitations)
        if summary.raw_summary:
            doc.add_heading("Full Summary", level=1)
            doc.add_paragraph(summary.raw_summary)
    else:
        doc.add_paragraph(str(summary))

    Path(output_path).parent.mkdir(parents=True, exist_ok=True)
    doc.save(str(output_path))


def save_as_json(
    summary: Union[str, StructuredSummary],
    output_path: str | Path,
) -> None:
    """Save a summary as a JSON file."""
    if isinstance(summary, StructuredSummary):
        data = summary.model_dump()
    else:
        data = {"summary": str(summary)}

    Path(output_path).parent.mkdir(parents=True, exist_ok=True)
    with open(output_path, "w") as f:
        json.dump(data, f, indent=2)


def save_as_csv(
    summary: Union[str, StructuredSummary],
    output_path: str | Path,
) -> None:
    """Save a summary as a CSV file."""
    if isinstance(summary, StructuredSummary):
        data = summary.model_dump()
        # Flatten list fields to comma-separated strings
        for key, val in data.items():
            if isinstance(val, list):
                data[key] = "; ".join(str(v) for v in val)
    else:
        data = {"summary": str(summary)}

    Path(output_path).parent.mkdir(parents=True, exist_ok=True)
    df = pd.DataFrame([data])
    df.to_csv(str(output_path), index=False)
