"""Conversion from Model classes to LaTeX tables."""

import logging
import re
from collections import defaultdict
from pathlib import Path

from astropy import units as u
from astropy.table import Table, join
from pydantic import BaseModel
from pydantic_core import PydanticUndefined

from ._core import type_to_string
from ._visitor import extract_keyword_mapping, walk_model_to_depth

logger = logging.getLogger(__name__)

__all__ = [
    "model_to_table",
    "model_to_latex_table",
    "keyword_mapping_latex_table",
    "generate_latex_table_includes",
]

MISSING = ""  #: blank string for tables with missing values
OPTIONAL_SYMBOL = r"${^\oslash}$"
OPTIONAL_TEXT = f"Fields marked with {OPTIONAL_SYMBOL} are optional."
NONE_TYPE = "Optional"
MODEL_DUMP_OPTIONS = {"by_alias": True}

TABLE_START_MAP = r"""
\begin{longtblr}{
  colspec={Q Q Q},
  row{odd}={bg=moongray},
  column{1-3} = {font=\ttfamily},
  row{1} = {bg=galaxyblue, font=\normalfont\bfseries, fg=white},
}
"""
TABLE_END_MAP = r"\end{longtblr}"


def _string_to_latex(text: str) -> str:
    """Transform text string to LaTeX."""
    if not isinstance(text, str):
        return ""

    # markdown backtics to texttt:
    text = re.sub(r"`([^`]+)`", r"\\texttt{\1}", text)

    # markdown bold to emph:
    text = re.sub(r"\*([^*]+)\*", r"\\emph{\1}", text)

    # remove paragraph breaks that cause issues in includes:
    text = text.replace("\n\n", " ")

    # escape underscores that are not already escaped
    return re.sub(r"(?<!\\)_", r"\\_", text)


def model_to_table(
    model: type[BaseModel], optional_symbol=OPTIONAL_SYMBOL, type_sep: str = " | "
) -> Table:
    """Turn a model into an Astropy table."""
    data = defaultdict(dict)
    types = defaultdict(list)
    has_optional = False

    # want to stay at depth=2, so we don't recurse
    for curmodel, path in walk_model_to_depth(model, max_depth=2):
        # skip parent:
        if curmodel is model:
            continue

        info = path[-1]
        extra = info.field.json_schema_extra or dict()

        # Name
        data[info.item_name]["Name"] = info.item_name
        if info.is_optional:
            data[info.item_name]["Name"] += optional_symbol
            has_optional = True

        # Types (which may be a list of several)
        types[info.item_name].append(info.item_type)

        # Unit in LaTex format
        if "unit" in extra:
            data[info.item_name]["Unit"] = u.Unit(extra["unit"]).to_string("latex")
        else:
            data[info.item_name]["Unit"] = ""

        # Description starting with the base description
        description = [info.field.description if info.field.description else ""]

        #   - Add the default value
        if info.field.get_default() not in (PydanticUndefined, None):
            description.append(f" (default {info.field.get_default()})")

        #   - add the examples
        if info.field.examples:
            description.append("Ex.: \\newline")
            examples = [rf" \tabitem{{`'{ex}'`}}" for ex in info.field.examples]
            description.append("\\newline".join(examples))

        #   - add the UCD
        if "ucd" in extra:
            description.append(rf" \newline(\textbf{{UCD}}: `{extra['ucd']}`)")

        data[info.item_name]["Description"] = _string_to_latex(" ".join(description))

    # compute Type column
    for name in data:
        data[name]["Type"] = type_sep.join(type_to_string(t) for t in types[name])

    return Table(rows=list(data.values()), meta=dict(has_optional=has_optional))


def _merge_type_and_unit_cols(table: Table) -> Table:
    """Turn unit and type columns into a single simplified column to save space."""
    if len(table) == 0:
        return table
    table["Type"] = [
        tp if un == "" else f"{tp} ({un})"
        for tp, un in zip(table["Type"], table["Unit"])
    ]
    del table["Unit"]
    table = table[["Name", "Type", "Description"]]  # col order
    table.columns["Type"].name = "Type (Unit)"
    return table


def _table_to_latex_table_rows(table: Table, **kwargs) -> str:
    """Return a LaTeX string for just the rows of the input Table.

    Values in the table with underscores are escaped.
    """
    import io

    table = table.copy()

    for col in table.colnames:
        if hasattr(col, "mask"):
            table[col] = col.filled("")

        table[col] = [_string_to_latex(x) for x in table[col]]

    out = io.StringIO()
    table.write(out, format="ascii.latex", **kwargs)
    return "\n".join([x for x in out.getvalue().split("\n") if not x.startswith("\\")])


def model_to_latex_table(model: type[BaseModel]) -> str:
    r"""Return LaTeX table in format used by Data Model documents.

    To use this in a document, you will need the following in your preamble:

    .. code-block:: latex

        \usepackage{tocbasic}
        \DeclareNewTOC[
          type=classdef,
          types=classdefs,
          name=Class Definition,
          listname={List of Class Definitions},
        ]{cls}

        \NewTblrEnviron{classdef}
        \SetTblrOuter[classdef]{long}
        \SetTblrInner[classdef]{
          colspec={Q[wd=3.3cm] Q[wd=3cm] X[l]},
          row{odd}={bg=moongray},
          column{1-2} = {font=\scriptsize},
          column{1} = {font=\ttfamily\bfseries\tiny},
          column{3} = {font=\scriptsize},
          row{1} = {bg=galaxyblue, font=\normalfont\bfseries, fg=white}
        }
        \NewTblrTheme{classdef}{
          \DeclareTblrTemplate{caption-tag}{default}{\textbf{Class Definition\hspace{0.25em}\thetable}}
          \DeclareTblrTemplate{caption-lot}{default}{%
            \addcontentsline{cls}{classdef}{\protect\numberline{\thetable}\InsertTblrText{entry}}%
          }
          \SetTblrStyle{firsthead}{font=\normalfont\captionsize}
          \SetTblrStyle{middlehead}{font=\normalfont\captionsize}
          \SetTblrStyle{lasthead}{font=\normalfont\captionsize}
        }

    """
    doc = _string_to_latex(model.__doc__) if model.__doc__ else ""
    table = model_to_table(model)
    opt = table.meta.get("has_optional", True) if table.meta else True

    return "\n".join(
        [
            r"\begin{classdef}[",
            r"    theme=classdef,",
            (
                rf"    caption={{\textbf{{\texttt{{{model.__name__}}}}}: "
                rf" {doc.replace('\n',' ')} {OPTIONAL_TEXT if opt else ''}}},"
            ),
            rf"    entry={{{model.__name__}}},",
            rf"    label={{tab:{model.__name__}}}",
            r"]{}",
            _table_to_latex_table_rows(_merge_type_and_unit_cols(table)),
            r"\end{classdef}",
        ]
    )


def keyword_mapping_latex_table(
    model: type[BaseModel], separator: str = ".", return_table: bool = False
):
    """Return LaTeX table with mappings to FITS and IVOA.

    Parameters
    ----------
    model : type[BaseModel]
        Model class to use.
    sep : str
        Separator to use when generating the CTAO key
    return_table : bool
        if True, return the Astropy table instead of the
        LaTeX representation.
    """
    fits_dict = extract_keyword_mapping(model, "fits_keyword", sep=separator)
    ivoa_dict = extract_keyword_mapping(model, "ivoa_keyword", sep=separator)

    fits = Table(
        dict(
            CTAO=list(fits_dict.keys()),
            FITS=list(" | ".join(f) for f in fits_dict.values()),
        )
    )
    ivoa = Table(
        dict(
            CTAO=list(ivoa_dict.keys()),
            IVOA=list(" | ".join(f) for f in ivoa_dict.values()),
        )
    )
    joined = join(ivoa, fits, join_type="outer")
    best = None

    if len(ivoa) > 0 and len(fits) > 0:
        best = joined
    elif len(ivoa) > 0:
        best = ivoa
    elif len(fits) > 0:
        best = fits
    else:
        return "no mappings found"

    if return_table:
        return best

    rows = _table_to_latex_table_rows(best)
    return f"{TABLE_START_MAP}\n{rows}\n{TABLE_END_MAP}"


def generate_latex_table_includes(
    models: list[type[BaseModel]], output_dir: Path | str
) -> None:
    """Write descriptive LaTeX tables for given models to the output_dir.

    These can then be included in latex documents.

    Parameters
    ----------
    models: list[BaseModel]
        models to write
    output_dir: Path|str
        output directory
    """
    output_dir = Path(output_dir)
    output_dir.mkdir(parents=True, exist_ok=True)

    for model in models:
        filename = output_dir / f"table_{model.__name__}.inc.tex"
        logger.debug("Writing: %s", filename)
        filename.write_text(model_to_latex_table(model))
