import sqlite3
import os
from google import genai
import time
from dotenv import load_dotenv

load_dotenv()

api_key = os.environ.get("GEMINI_API_KEY")
if not api_key:
    print("GEMINI_API_KEY not found in .env")
    exit(1)

client = genai.Client(api_key=api_key)

db_path = "smart_learning.db"
resources_dir = "resources"

# Map subject to their main PDF
pdf_map = {
    "biology": "BIOLOGY-NOTES-FORM-1-4-BOOKLET.pdf",
    "mathematics": "Final-O-LEVEL-MATH-1-4-final.pdf",
    "computer_science": "DeSila-ICT Skills Full Book (xoer).pdf",
    "science": "BIOLOGY-NOTES-FORM-1-4-BOOKLET.pdf" # Fallback
}

uploaded_files = {}

def get_or_upload_file(pdf_name):
    if pdf_name in uploaded_files:
        return uploaded_files[pdf_name]
    
    pdf_path = os.path.join(resources_dir, pdf_name)
    if not os.path.exists(pdf_path):
        print(f"File not found: {pdf_path}")
        return None
        
    print(f"Uploading {pdf_name}...")
    uploaded_file = client.files.upload(file=pdf_path)
    
    # Wait for processing
    print(f"Waiting for {pdf_name} to process...")
    while True:
        file_info = client.files.get(name=uploaded_file.name)
        if file_info.state.name == "ACTIVE":
            break
        elif file_info.state.name == "FAILED":
            print(f"Failed to process {pdf_name}")
            return None
        time.sleep(2)
        
    uploaded_files[pdf_name] = uploaded_file
    return uploaded_file

def generate_explanation(topic_title, form, subject, file_obj):
    prompt = f"""
    You are an expert Zambian teacher.
    Based ON THE PROVIDED DOCUMENT ONLY, create a concise, educational explanation for the topic: "{topic_title}".
    This is for a Form {form} student in the subject of {subject}.
    Keep it well-structured with basic HTML tags: <h3> for the main title, and <p> for paragraphs. 
    Do not include any practical activities or questions, just a solid conceptual explanation of about 150-250 words based on what's in the book.
    """
    
    try:
        response = client.models.generate_content(
            model='gemini-2.5-flash',
            contents=[file_obj, prompt],
        )
        return response.text
    except Exception as e:
        print(f"Error generating explanation: {e}")
        return None

def main():
    conn = sqlite3.connect(db_path)
    cursor = conn.cursor()
    
    cursor.execute("SELECT id, subject, form, title FROM topics")
    topics = cursor.fetchall()
    
    for row in topics:
        t_id, subject, form, title = row
        print(f"Processing Topic {t_id}: {title} ({subject})")
        
        pdf_name = pdf_map.get(subject)
        if not pdf_name:
            continue
            
        file_obj = get_or_upload_file(pdf_name)
        if not file_obj:
            continue
            
        explanation = generate_explanation(title, form, subject, file_obj)
        if explanation:
            # Clean up explanation
            if explanation.startswith("```html"):
                explanation = explanation[7:-3].strip()
            
            # Combine the generic instruction at the end
            final_content = f"{explanation}<p><em>(Use the Simplify Explanation button if you need this broken down further)</em></p>"
            
            cursor.execute("UPDATE topics SET content = ? WHERE id = ?", (final_content, t_id))
            conn.commit()
            print(f"Updated Topic {t_id}")
            time.sleep(1) # Rate limiting
            
    conn.close()
    print("Done!")

if __name__ == "__main__":
    main()