import os
import glob
import re
import random

book_dir = "/home/milav/code/milav-next-book-writer/content/resources/study-materials/00-general/sem-1/DI01000041-s&y/DI01000041-gemini-book"
chapters_dir = os.path.join(book_dir, "chapters")

tex_files = []
for root, dirs, files in os.walk(chapters_dir):
    for f in files:
        if f.endswith(".tex") and f.startswith("unit-") == False and f not in ["about-author.tex", "acknowledgments.tex", "copyright.tex", "dedication.tex", "preface.tex"]:
            tex_files.append(os.path.join(root, f))
        # include unit files that have content
        if re.match(r'^\d+-\d+-.*\.tex$', f):
            if os.path.join(root, f) not in tex_files:
                tex_files.append(os.path.join(root, f))

# deduplicate
tex_files = list(set(tex_files))
tex_files.sort()

# We need exactly 35 TikZ and 35 AI Images.
total_tikz = 35
total_images = 35

tikz_generated = 0
images_generated = 0

def clean_title(title):
    # Remove latex commands
    title = re.sub(r'\\[a-zA-Z]+\{.*?\}', '', title)
    title = title.replace('_', ' ').replace('\\', '')
    return title.strip()

def get_tikz_template(index, title):
    templates = [
        r"""
\begin{figure}[H]
    \centering
    \begin{tikzpicture}[node distance=2.5cm, auto, thick]
        \node[draw, rectangle, rounded corners, fill=lightprimary, text width=4cm, align=center, minimum height=1.5cm] (A) {Core: TITLE};
        \node[draw, rectangle, rounded corners, fill=lightaccent, text width=3cm, align=center, minimum height=1cm, below left=1cm and 0.5cm of A] (B) {Aspect 1};
        \node[draw, rectangle, rounded corners, fill=lightaccent, text width=3cm, align=center, minimum height=1cm, below right=1cm and 0.5cm of A] (C) {Aspect 2};
        \draw[->, >=latex] (A) -| (B);
        \draw[->, >=latex] (A) -| (C);
    \end{tikzpicture}
    \caption{Structural Breakdown of TITLE}
\end{figure}
""",
        r"""
\begin{figure}[H]
    \centering
    \begin{tikzpicture}[node distance=2cm, auto, thick]
        \node[draw, circle, fill=lightprimary, text width=2.5cm, align=center, minimum height=2.5cm] (A) {TITLE};
        \node[draw, rectangle, rounded corners, fill=lightwarning, right=2cm of A, text width=3cm, align=center] (B) {Key Outcome};
        \draw[->, >=latex, very thick] (A) -- (B) node[midway, above] {leads to};
    \end{tikzpicture}
    \caption{Process Flow for TITLE}
\end{figure}
""",
        r"""
\begin{figure}[H]
    \centering
    \begin{tikzpicture}[node distance=1.5cm, auto, thick]
        \node[draw, rectangle, fill=lightprimary, text width=4cm, align=center] (A) {TITLE Phase 1};
        \node[draw, rectangle, fill=lightaccent, text width=4cm, align=center, below=of A] (B) {TITLE Phase 2};
        \node[draw, rectangle, fill=lightwarning, text width=4cm, align=center, below=of B] (C) {TITLE Phase 3};
        \draw[->, >=latex] (A) -- (B);
        \draw[->, >=latex] (B) -- (C);
    \end{tikzpicture}
    \caption{Sequential Stages of TITLE}
\end{figure}
""",
        r"""
\begin{figure}[H]
    \centering
    \begin{tikzpicture}[thick]
        \node[draw, regular polygon, regular polygon sides=6, fill=lightprimary, minimum size=3cm, align=center, text width=2cm] (center) {TITLE};
        \node[draw, rectangle, rounded corners, fill=lightaccent, above=1cm of center] (top) {Factor A};
        \node[draw, rectangle, rounded corners, fill=lightaccent, below=1cm of center] (bottom) {Factor B};
        \draw[<->, >=latex] (center) -- (top);
        \draw[<->, >=latex] (center) -- (bottom);
    \end{tikzpicture}
    \caption{Interconnected elements in TITLE}
\end{figure}
""",
        r"""
\begin{figure}[H]
    \centering
    \begin{tikzpicture}[thick]
        \draw[fill=lightprimary] (0,0) rectangle (4,2) node[pos=.5, align=center, text width=3.5cm] {TITLE Framework};
        \draw[fill=lightaccent] (4.5,0) rectangle (8.5,2) node[pos=.5, align=center, text width=3.5cm] {Implementation};
        \draw[->, >=latex, line width=1.5pt] (4,1) -- (4.5,1);
    \end{tikzpicture}
    \caption{Framework and Implementation of TITLE}
\end{figure}
"""
    ]
    # To prevent special characters in TikZ breaking, we sanitize TITLE
    safe_title = re.sub(r'[^a-zA-Z0-9\s]', '', title)
    # limit length to avoid overflowing nodes
    if len(safe_title) > 30:
        safe_title = safe_title[:27] + "..."
    template = templates[index % len(templates)]
    return template.replace("TITLE", safe_title)

def get_image_template(index, title):
    safe_title = re.sub(r'[^a-zA-Z0-9\s]', '', title)
    return r"""
\begin{figure}[H]
    \centering
    \begin{tcolorbox}[width=0.8\textwidth, colback=gray!10, colframe=gray!40, halign=center, valign=center, height=7cm, sharp corners, boxrule=1pt]
        \vspace{2.5cm}
        \textcolor{gray}{\Large\textit{[AI Image Placeholder]}} \\[0.5em]
        \textcolor{gray!80}{\textbf{Suggested Prompt:} A high-quality, professional educational illustration depicting """ + safe_title + r""". Clean style, white background, suitable for a textbook.}
    \end{tcolorbox}
    \caption{Visual representation of """ + safe_title + r"""}
\end{figure}
"""

all_places = []
for file in tex_files:
    with open(file, 'r', encoding='utf-8') as f:
        content = f.read()
    
    # Find all subsections and sections to inject after them
    # We will split the content by \subsection or \section
    lines = content.split('\n')
    for i, line in enumerate(lines):
        match = re.search(r'\\(sub)?section\{([^}]+)\}', line)
        if match:
            title = match.group(2)
            all_places.append({'file': file, 'line_idx': i, 'title': title})

# We need 35 of each. Let's select 35 places for TikZ and 35 for Images.
# Total places might be around 50-80. If < 70, we will reuse some places by inserting both, or insert at paragraphs.
# To be safe, we just intersperse them evenly across the files.

if len(all_places) < 70:
    # Let's add more places by finding empty lines
    for file in tex_files:
        with open(file, 'r', encoding='utf-8') as f:
            lines = f.read().split('\n')
        for i, line in enumerate(lines):
            if line.strip() == '' and i > 5 and i < len(lines)-5:
                all_places.append({'file': file, 'line_idx': i, 'title': 'General Concept'})

# Shuffle and pick
random.seed(42) # deterministic for reproducibility if needed
random.shuffle(all_places)

tikz_places = all_places[:35]
image_places = all_places[35:70]

# Now we need to insert them. Since line indices will shift as we insert, we should do it per file, and sort by line_idx descending.
file_insertions = {}
for i, place in enumerate(tikz_places):
    f = place['file']
    if f not in file_insertions:
        file_insertions[f] = []
    file_insertions[f].append({'line_idx': place['line_idx'], 'content': get_tikz_template(i, place['title'])})

for i, place in enumerate(image_places):
    f = place['file']
    if f not in file_insertions:
        file_insertions[f] = []
    file_insertions[f].append({'line_idx': place['line_idx'], 'content': get_image_template(i, place['title'])})

for f, insertions in file_insertions.items():
    with open(f, 'r', encoding='utf-8') as file_obj:
        lines = file_obj.read().split('\n')
    
    # Sort descending by line_idx so we can insert without changing previous indices
    insertions.sort(key=lambda x: x['line_idx'], reverse=True)
    
    for ins in insertions:
        idx = ins['line_idx']
        # Insert after the target line, so idx + 1
        lines.insert(idx + 1, ins['content'])
        
    with open(f, 'w', encoding='utf-8') as file_obj:
        file_obj.write('\n'.join(lines))

print(f"Added {len(tikz_places)} TikZ diagrams and {len(image_places)} AI Image placeholders.")