"""
LaTeX document builder for cheat sheet generation.

Creates structured LaTeX documents from CheatSheetContent
and compiles them to PDF.
"""

import os
import re
import subprocess
import tempfile
from dataclasses import dataclass, field
from pathlib import Path
from typing import Any, Dict, List, Optional, Union

from jinja2 import Environment, BaseLoader

from ..infra.logging import get_logger
from ..infra.tracing import trace_function

logger = get_logger(__name__)


# Default LaTeX template for cheat sheets
DEFAULT_TEMPLATE = r"""
\documentclass[{{ font_size }}pt,{{ paper_size }}]{article}

% Packages
\usepackage[utf8]{inputenc}
\usepackage[T1]{fontenc}
\usepackage{geometry}
\usepackage{multicol}
\usepackage{amsmath,amssymb,amsthm}
\usepackage{algorithm}
\usepackage{algpseudocode}
\usepackage{listings}
\usepackage{xcolor}
\usepackage{hyperref}
\usepackage{titlesec}
\usepackage{enumitem}
\usepackage{fancyhdr}
\usepackage{parskip}

% Page geometry
\geometry{
    top={{ margins.top }},
    bottom={{ margins.bottom }},
    left={{ margins.left }},
    right={{ margins.right }}
}

% Colors
\definecolor{sectioncolor}{RGB}{0,51,102}
\definecolor{codebackground}{RGB}{248,248,248}
\definecolor{codecomment}{RGB}{106,153,85}
\definecolor{codekeyword}{RGB}{0,0,255}

% Section formatting
\titleformat{\section}
    {\normalfont\large\bfseries\color{sectioncolor}}
    {\thesection}{1em}{}
\titleformat{\subsection}
    {\normalfont\normalsize\bfseries\color{sectioncolor}}
    {\thesubsection}{1em}{}
\titlespacing*{\section}{0pt}{2ex plus 1ex minus .2ex}{1ex plus .2ex}
\titlespacing*{\subsection}{0pt}{1.5ex plus 1ex minus .2ex}{0.5ex plus .2ex}

% List formatting
\setlist[itemize]{noitemsep, topsep=0pt, leftmargin=*}
\setlist[enumerate]{noitemsep, topsep=0pt, leftmargin=*}

% Code listing style
\lstset{
    basicstyle=\ttfamily\small,
    backgroundcolor=\color{codebackground},
    commentstyle=\color{codecomment},
    keywordstyle=\color{codekeyword},
    breaklines=true,
    frame=single,
    framesep=2pt,
    xleftmargin=5pt,
    xrightmargin=5pt
}

% Header/Footer
\pagestyle{fancy}
\fancyhf{}
\fancyhead[L]{\small {{ title }}}
\fancyhead[R]{\small \thepage}
\renewcommand{\headrulewidth}{0.4pt}

% Theorem environments
\newtheorem{theorem}{Theorem}
\newtheorem{lemma}{Lemma}
\newtheorem{definition}{Definition}

% Custom commands
\newcommand{\important}[1]{\textbf{#1}}
\newcommand{\concept}[1]{\textit{#1}}

\begin{document}

{% if include_toc %}
\tableofcontents
\newpage
{% endif %}

% Title
\begin{center}
    {\LARGE\bfseries {{ title }}}\\[0.5em]
    {% if authors %}
    {\normalsize {{ authors | join(', ') }}}\\[0.3em]
    {% endif %}
    {\small\today}
\end{center}

\vspace{1em}

{% if columns > 1 %}
\begin{multicols}{ {{ columns }} }
{% endif %}

% Overview
{% if overview %}
\section*{Overview}
{{ overview }}
{% endif %}

% Core Concepts
{% if concepts %}
\section{Core Concepts}
{% for concept in concepts %}
\textbf{ {{ concept.name }} }{% if concept.definition %}: {{ concept.definition }}{% endif %}
{% if concept.intuition %}
\textit{Intuition:} {{ concept.intuition }}
{% endif %}
{% endfor %}
{% endif %}

% Key Formulas
{% if formulas %}
\section{Key Formulas}
{% for group_name, group_formulas in formulas_grouped.items() %}
\subsection{ {{ group_name }} }
{% for formula in group_formulas %}
{% if formula.name %}\textbf{ {{ formula.name }} }:{% endif %}
{% if formula.latex %}
{{ formula.latex }}
{% endif %}
{% if formula.context %}
\textit{ {{ formula.context }} }
{% endif %}

{% endfor %}
{% endfor %}
{% endif %}

% Models
{% if models %}
\section{Models \& Theory}
{% for model in models %}
{% if model.summary %}
{{ model.summary }}
{% else %}
\subsection{ {{ model.name }} }
{% if model.assumptions %}
\textbf{Assumptions:}
\begin{itemize}
{% for assumption in model.assumptions %}
\item {{ assumption }}
{% endfor %}
\end{itemize}
{% endif %}
{% if model.objective %}
\textbf{Objective:} {{ model.objective }}
{% endif %}
{% if model.formulation %}
\textbf{Formulation:} {{ model.formulation }}
{% endif %}
{% endif %}
{% endfor %}
{% endif %}

% Algorithms
{% if algorithms %}
\section{Algorithms}
{% for algo in algorithms %}
\subsection{ {{ algo.name }} }
{% if algo.purpose %}
\textbf{Purpose:} {{ algo.purpose }}
{% endif %}

{% if algo.input %}
\textbf{Input:} {{ algo.input }}
{% endif %}

{% if algo.output %}
\textbf{Output:} {{ algo.output }}
{% endif %}

{% if algo.pseudocode %}
\begin{lstlisting}
{{ algo.pseudocode }}
\end{lstlisting}
{% endif %}

{% if algo.time_complexity or algo.space_complexity %}
\textbf{Complexity:} {% if algo.time_complexity %}Time: {{ algo.time_complexity }}{% endif %}{% if algo.space_complexity %}, Space: {{ algo.space_complexity }}{% endif %}
{% endif %}

{% if algo.key_insight %}
\textbf{Key Insight:} {{ algo.key_insight }}
{% endif %}

{% endfor %}
{% endif %}

% Connections
{% if connections %}
\section{Related Topics}
\begin{itemize}
{% for conn in connections %}
\item \textbf{ {{ conn.topic }} }: {{ conn.description }}
{% endfor %}
\end{itemize}
{% endif %}

% Recap
{% if recap %}
\section{Quick Reference}
{{ recap }}
{% endif %}

{% if columns > 1 %}
\end{multicols}
{% endif %}

% References
{% if include_references and references %}
\section*{References}
\begin{enumerate}
{% for ref in references %}
\item {{ ref }}
{% endfor %}
\end{enumerate}
{% endif %}

\end{document}
"""


@dataclass
class CheatSheet:
    """Represents a complete cheat sheet document."""
    title: str
    content: Dict[str, Any]
    latex_source: str = ""
    pdf_path: Optional[Path] = None
    
    def to_dict(self) -> Dict[str, Any]:
        return {
            "title": self.title,
            "content": self.content,
            "has_pdf": self.pdf_path is not None
        }


class LaTeXBuilder:
    """
    Builder for creating LaTeX cheat sheets.
    
    Uses Jinja2 templates for flexible document generation.
    """
    
    def __init__(
        self,
        settings=None,
        template: Optional[str] = None
    ):
        """
        Initialize LaTeX builder.
        
        Args:
            settings: LaTeXSettings or dict with settings
            template: Custom Jinja2 template string
        """
        self.settings = settings
        self.template_str = template or DEFAULT_TEMPLATE
        
        # Initialize Jinja2 environment
        self.env = Environment(loader=BaseLoader())
        self.env.filters['escape_latex'] = self._escape_latex
        
        # Compile template
        self.template = self.env.from_string(self.template_str)
    
    @staticmethod
    def _escape_latex(text: str) -> str:
        """Escape special LaTeX characters."""
        if not text:
            return ""
        
        # Characters that need escaping
        chars = {
            '&': r'\&',
            '%': r'\%',
            '#': r'\#',
            '_': r'\_',
            '{': r'\{',
            '}': r'\}',
            '~': r'\textasciitilde{}',
            '^': r'\textasciicircum{}',
        }
        
        # Don't escape if already in math mode or LaTeX command
        if text.startswith('$') or text.startswith('\\'):
            return text
        
        for char, replacement in chars.items():
            text = text.replace(char, replacement)
        
        return text
    
    @trace_function(name="build_latex")
    def build(self, content: "CheatSheetContent") -> str:
        """
        Build LaTeX source from content.
        
        Args:
            content: CheatSheetContent to render
            
        Returns:
            LaTeX source string
        """
        # Prepare template context
        context = self._prepare_context(content)
        
        # Render template
        latex_source = self.template.render(**context)
        
        # Post-process
        latex_source = self._post_process(latex_source)
        
        logger.info(f"Built LaTeX document: {len(latex_source)} characters")
        return latex_source
    
    def _prepare_context(self, content: "CheatSheetContent") -> Dict[str, Any]:
        """Prepare template context from content and settings."""
        # Get settings or defaults
        if self.settings:
            font_size = self.settings.font_size
            paper_size = self.settings.paper_size
            margins = {
                "top": self.settings.margins.top,
                "bottom": self.settings.margins.bottom,
                "left": self.settings.margins.left,
                "right": self.settings.margins.right,
            }
            columns = self.settings.columns
            include_toc = self.settings.include_toc
            include_references = self.settings.include_references
        else:
            font_size = 10
            paper_size = "a4paper"
            margins = {"top": "1.5cm", "bottom": "1.5cm", "left": "1.5cm", "right": "1.5cm"}
            columns = 2
            include_toc = False
            include_references = True
        
        # Group formulas by topic
        formulas_grouped = self._group_formulas(content.formulas)
        
        return {
            "title": content.title,
            "authors": content.references[:1] if content.references else [],
            "overview": content.overview,
            "concepts": content.concepts,
            "formulas": content.formulas,
            "formulas_grouped": formulas_grouped,
            "models": content.models,
            "algorithms": content.algorithms,
            "connections": content.connections,
            "recap": content.recap,
            "references": content.references,
            "font_size": font_size,
            "paper_size": paper_size,
            "margins": margins,
            "columns": columns,
            "include_toc": include_toc,
            "include_references": include_references,
        }
    
    def _group_formulas(self, formulas: List[Dict[str, Any]]) -> Dict[str, List[Dict[str, Any]]]:
        """Group formulas by topic."""
        groups: Dict[str, List[Dict[str, Any]]] = {}
        
        for formula in formulas:
            topic = formula.get("topic", "General")
            if not topic:
                topic = "General"
            
            if topic not in groups:
                groups[topic] = []
            groups[topic].append(formula)
        
        return groups if groups else {"General": formulas}
    
    def _post_process(self, latex: str) -> str:
        """Post-process LaTeX source."""
        # Remove excessive blank lines
        latex = re.sub(r'\n{3,}', '\n\n', latex)
        
        # Fix common issues
        latex = re.sub(r'\\textbf\{\s*\}', '', latex)  # Remove empty bold
        latex = re.sub(r'\\textit\{\s*\}', '', latex)  # Remove empty italic
        
        return latex.strip()
    
    @trace_function(name="compile_pdf")
    def compile_pdf(
        self,
        latex_source: str,
        output_path: Union[str, Path],
        compiler: Optional[str] = None
    ) -> Path:
        """
        Compile LaTeX source to PDF.
        
        Args:
            latex_source: LaTeX source string
            output_path: Path for output PDF
            compiler: LaTeX compiler to use (pdflatex, xelatex, lualatex, tectonic)
            
        Returns:
            Path to generated PDF
        """
        output_path = Path(output_path)
        
        # Get compiler from settings or default
        if compiler is None:
            compiler = self.settings.compiler if self.settings else "pdflatex"
        
        # Create temporary directory for compilation
        with tempfile.TemporaryDirectory() as tmpdir:
            tmpdir = Path(tmpdir)
            tex_file = tmpdir / "cheatsheet.tex"
            
            # Write LaTeX source
            tex_file.write_text(latex_source, encoding='utf-8')
            
            # Compile
            try:
                if compiler == "tectonic":
                    self._compile_with_tectonic(tex_file, tmpdir)
                else:
                    self._compile_with_latex(tex_file, tmpdir, compiler)
                
                # Copy PDF to output location
                pdf_file = tmpdir / "cheatsheet.pdf"
                if pdf_file.exists():
                    output_path.parent.mkdir(parents=True, exist_ok=True)
                    
                    import shutil
                    shutil.copy(pdf_file, output_path)
                    
                    logger.info(f"Generated PDF: {output_path}")
                    return output_path
                else:
                    raise RuntimeError("PDF not generated")
                    
            except subprocess.CalledProcessError as e:
                # Try to get error from log file
                log_file = tmpdir / "cheatsheet.log"
                error_msg = ""
                if log_file.exists():
                    log_content = log_file.read_text(encoding='utf-8', errors='ignore')
                    # Extract error lines
                    for line in log_content.split('\n'):
                        if '!' in line or 'Error' in line:
                            error_msg += line + '\n'
                
                logger.error(f"LaTeX compilation failed: {error_msg or str(e)}")
                raise RuntimeError(f"LaTeX compilation failed: {error_msg or str(e)}")
    
    def _compile_with_latex(
        self,
        tex_file: Path,
        tmpdir: Path,
        compiler: str
    ) -> None:
        """Compile with standard LaTeX compiler."""
        # Run twice for references
        for _ in range(2):
            result = subprocess.run(
                [
                    compiler,
                    "-interaction=nonstopmode",
                    "-halt-on-error",
                    tex_file.name
                ],
                cwd=tmpdir,
                capture_output=True,
                text=True,
                timeout=120
            )
            
            if result.returncode != 0:
                # Check if it's a real error or just warnings
                log_file = tmpdir / "cheatsheet.log"
                if log_file.exists():
                    log_content = log_file.read_text(encoding='utf-8', errors='ignore')
                    if "Fatal error" in log_content or "! " in log_content:
                        raise subprocess.CalledProcessError(
                            result.returncode, compiler, result.stdout, result.stderr
                        )
    
    def _compile_with_tectonic(self, tex_file: Path, tmpdir: Path) -> None:
        """Compile with tectonic."""
        result = subprocess.run(
            ["tectonic", tex_file.name],
            cwd=tmpdir,
            capture_output=True,
            text=True,
            timeout=120
        )
        
        if result.returncode != 0:
            raise subprocess.CalledProcessError(
                result.returncode, "tectonic", result.stdout, result.stderr
            )
    
    def save_latex(self, latex_source: str, output_path: Union[str, Path]) -> Path:
        """
        Save LaTeX source to file.
        
        Args:
            latex_source: LaTeX source string
            output_path: Output file path
            
        Returns:
            Path to saved file
        """
        output_path = Path(output_path)
        output_path.parent.mkdir(parents=True, exist_ok=True)
        
        output_path.write_text(latex_source, encoding='utf-8')
        logger.info(f"Saved LaTeX source: {output_path}")
        
        return output_path
    
    def build_and_compile(
        self,
        content: "CheatSheetContent",
        output_path: Union[str, Path]
    ) -> CheatSheet:
        """
        Build LaTeX and compile to PDF.
        
        Args:
            content: CheatSheetContent
            output_path: Output PDF path
            
        Returns:
            CheatSheet with source and PDF path
        """
        # Build LaTeX
        latex_source = self.build(content)
        
        # Compile to PDF
        pdf_path = self.compile_pdf(latex_source, output_path)
        
        return CheatSheet(
            title=content.title,
            content=content.to_dict(),
            latex_source=latex_source,
            pdf_path=pdf_path
        )


def check_latex_installation() -> Dict[str, bool]:
    """
    Check which LaTeX compilers are available.
    
    Returns:
        Dictionary of compiler name to availability
    """
    compilers = ["pdflatex", "xelatex", "lualatex", "tectonic"]
    available = {}
    
    for compiler in compilers:
        try:
            result = subprocess.run(
                [compiler, "--version"],
                capture_output=True,
                timeout=5
            )
            available[compiler] = result.returncode == 0
        except (subprocess.TimeoutExpired, FileNotFoundError):
            available[compiler] = False
    
    return available
