import sys
sys.path.append('/app')
from sqlmodel import Session, select
from database import engine
import models
import glob
import os
import fitz
import re

def parse_pdf(filepath):
    doc = fitz.open(filepath)
    text = ""
    for page in doc:
        text += page.get_text() + "\n"
        
    q_block = text.split("Questions\n")[-1].split("Answer Key\n")[0]
    q_pattern = re.compile(r'^(\d+)\.\s+(.*?)(?=^\d+\.\s+|\Z)', re.DOTALL | re.MULTILINE)
    
    questions_dict = {}
    for match in q_pattern.finditer(q_block):
        q_num = int(match.group(1))
        if 1 <= q_num <= 15:
            questions_dict[q_num] = match.group(2).strip()

    answers_dict = {}
    if "Answer Key\n" in text:
        a_block = text.split("Answer Key\n")[-1]
        a_pattern = re.compile(r'^(\d+)\s+(.*?)(?=^\d+\s+|\Z)', re.DOTALL | re.MULTILINE)
        for match in a_pattern.finditer(a_block):
            q_num = int(match.group(1))
            if 1 <= q_num <= 15:
                if q_num not in answers_dict:
                    answers_dict[q_num] = match.group(2).strip()
            
    return questions_dict, answers_dict

def format_text(text):
    if not text: return text
    
    # Original targets
    text = re.sub(r'\bsp3d2\b', r'$sp^3d^2$', text)
    text = re.sub(r'\bsp3d\b', r'$sp^3d$', text)
    text = re.sub(r'\bsp3\b', r'$sp^3$', text)
    text = re.sub(r'\bsp2\b', r'$sp^2$', text)
    
    def repl(m):
        word = m.group(0)
        if word.lower() in ["sp2", "sp3", "sp3d", "sp3d2", "3d", "a1", "b1", "c1"]: 
            return word
        return '$' + re.sub(r'(\d+)', r'_\1', word) + '$'
        
    text = re.sub(r'\b[A-Z][a-z]*[A-Z]*[a-z]*\d+[A-Za-z0-9]*\b', repl, text)

    text = text.replace('z1 = 1 + i', '$z_1 = 1 + i$')
    text = text.replace('z2 = 2 − i', '$z_2 = 2 - i$')
    text = text.replace('z1\nz2', '$\\frac{z_1}{z_2}$')
    text = text.replace('z−1\nz+1', '$\\frac{z-1}{z+1}$')
    text = text.replace('z2 = ¯z', '$z^2 = \\bar{z}$')
    text = text.replace('x2 − 3x + 5 = 0', '$x^2 - 3x + 5 = 0$')
    text = text.replace('x2 + 2x + 5 = 0', '$x^2 + 2x + 5 = 0$')
    text = text.replace('1+i\n1−i', '$\\frac{1+i}{1-i}$')
    text = text.replace('z + 1\nz', '$\\frac{z+1}{z}$')
    text = text.replace('z − 1\nz i', '$\\frac{z-1}{z i}$')
    text = text.replace('30◦', '30°')
    text = text.replace('⃗v = 3', '$\\vec{v} = 3$')
    
    text = text.replace('σ bonds', '$\\sigma$ bonds')
    text = text.replace('π bonds', '$\\pi$ bonds')
    
    return text

def fix_all_from_source():
    pdf_files = glob.glob('/app/pdfs_to_parse/*.pdf')
    with Session(engine) as session:
        for pdf_path in pdf_files:
            filename = os.path.basename(pdf_path)
            subject = "Physics"
            topic = "General"
            if "phy" in filename.lower():
                subject = "Physics"
                topic = "Laws of Motion"
            elif "che" in filename.lower():
                subject = "Chemistry"
                topic = "Chemical Bonding"
            elif "mth" in filename.lower():
                subject = "Mathematics"
                topic = "Calculus"
            elif "bio" in filename.lower():
                subject = "Biology"
                topic = "Genetics"

            q_dict, a_dict = parse_pdf(pdf_path)
            
            # Find the test for this topic
            test = session.exec(select(models.Test).where(models.Test.title == f"DPP: {subject} - {topic}")).first()
            if not test:
                continue
                
            q_ids = test.sections.get("Section 1", {}).get("questions", [])
            for i, q_id in enumerate(q_ids):
                q_num = i + 1
                if q_num in q_dict:
                    raw_text = q_dict[q_num].replace('\x00', '')
                    formatted = format_text(raw_text)
                    
                    q = session.exec(select(models.Question).where(models.Question.id == q_id)).first()
                    if q:
                        q.question_text = formatted
                        session.add(q)
        
        session.commit()
        print("Successfully updated questions from source.")

if __name__ == "__main__":
    fix_all_from_source()