"""
Export functionality for AKShare CLI harness.

Supports exporting DataFrames to CSV, JSON, Excel, and Markdown.
"""

import json
import os
from typing import Optional

import pandas as pd


def export_csv(df: pd.DataFrame, filepath: str, index: bool = False) -> str:
    """Export DataFrame to CSV file."""
    os.makedirs(os.path.dirname(os.path.abspath(filepath)), exist_ok=True)
    df.to_csv(filepath, index=index, encoding="utf-8-sig")
    return os.path.abspath(filepath)


def export_json(df: pd.DataFrame, filepath: str, orient: str = "records") -> str:
    """Export DataFrame to JSON file."""
    os.makedirs(os.path.dirname(os.path.abspath(filepath)), exist_ok=True)
    df.to_json(filepath, orient=orient, force_ascii=False, indent=2)
    return os.path.abspath(filepath)


def export_excel(df: pd.DataFrame, filepath: str, sheet_name: str = "Sheet1") -> str:
    """Export DataFrame to Excel file."""
    os.makedirs(os.path.dirname(os.path.abspath(filepath)), exist_ok=True)
    df.to_excel(filepath, index=False, sheet_name=sheet_name)
    return os.path.abspath(filepath)


def export_markdown(df: pd.DataFrame, filepath: str) -> str:
    """Export DataFrame to Markdown table file."""
    os.makedirs(os.path.dirname(os.path.abspath(filepath)), exist_ok=True)
    with open(filepath, "w", encoding="utf-8") as f:
        f.write(df.to_markdown(index=False))
        f.write("\n")
    return os.path.abspath(filepath)


def format_table(
    df: pd.DataFrame,
    max_rows: Optional[int] = None,
    show_index: bool = False,
    max_col_width: int = 30,
) -> str:
    """Format DataFrame as a text table for terminal display."""
    if max_rows is not None and len(df) > max_rows:
        display_df = pd.concat([df.head(max_rows // 2), df.tail(max_rows // 2)])
        truncated = True
        total_rows = len(df)
    else:
        display_df = df
        truncated = False
        total_rows = len(df)

    # Truncate wide columns for display
    display_df = display_df.copy()
    for col in display_df.columns:
        if display_df[col].dtype == object:
            display_df[col] = display_df[col].astype(str).str[:max_col_width]

    table = display_df.to_string(index=show_index)

    if truncated:
        table += f"\n\n... [{total_rows} rows total, showing {max_rows}]"

    return table


def format_json(df: pd.DataFrame, max_rows: Optional[int] = None) -> str:
    """Format DataFrame as JSON string."""
    if max_rows is not None:
        df = df.head(max_rows)

    records = json.loads(df.to_json(orient="records", force_ascii=False))
    result = {
        "total_rows": len(df),
        "columns": list(df.columns),
        "data": records,
    }
    return json.dumps(result, indent=2, ensure_ascii=False, default=str)


def format_csv(df: pd.DataFrame, max_rows: Optional[int] = None) -> str:
    """Format DataFrame as CSV string."""
    if max_rows is not None:
        df = df.head(max_rows)
    return df.to_csv(index=False)


def auto_export(
    df: pd.DataFrame,
    filepath: str,
    **kwargs,
) -> str:
    """Export DataFrame based on file extension."""
    ext = os.path.splitext(filepath)[1].lower()
    exporters = {
        ".csv": export_csv,
        ".json": export_json,
        ".xlsx": export_excel,
        ".xls": export_excel,
        ".md": export_markdown,
    }
    exporter = exporters.get(ext)
    if exporter is None:
        raise ValueError(
            f"Unsupported file format: {ext}. "
            f"Supported: {', '.join(exporters.keys())}"
        )
    return exporter(df, filepath, **kwargs)
