import os
import sys
import time
from datetime import datetime
from google import genai
from google.genai import types

# -------------------------
# CONFIG
# -------------------------
API_KEY = "AIzaSyBsgUkZFqiwJ-hDlW_bQ-uzTXocXaBT8g4"
INPUT_DIR = "v-md"
OUTPUT_DIR = "v-md-notes"

# -------------------------
# Gemini Setup
# -------------------------
client = genai.Client(api_key=API_KEY)

# -------------------------
# Function: Generate Elaborated Note
# -------------------------
def elaborate_note(note_text):
    prompt = f"""
You are a UPSC academic assistant. Based on the following topic summary, generate detailed Elaborate notes in the following format:

## Elaborate Notes
Write a comprehensive, point-by-point explanation of the topic using academic tone. Each bullet or concept in the summary should be expanded with historical context, examples, and references to scholars and their works along with year, Historical facts or archaeological findings where appropriate.

## Prelims Pointers
Extract factual data suitable for UPSC Prelims. Use bullet points with '-' for unordered lists and '1.' for ordered lists. Avoid analysis.

## Mains Insights
Provide analytical perspectives suitable for GS Paper I–IV. Include cause-effect relationships, debates, and historiographical viewpoints. Use structured lists where appropriate.

Do not use numbered headings for sections. Do not use emojis or casual tone. Maintain academic rigor and Markdown formatting. Don't miss even a minor information. 

Topic Summary:
{note_text}
"""
    try:
        response = client.models.generate_content(
            model="gemini-2.5-pro",
            contents=prompt
        )

        # Debug: token usage
        if hasattr(response, "usage_metadata") and response.usage_metadata:
            input_tokens = response.usage_metadata.prompt_token_count
            print("Input tokens:", input_tokens)
            print("Thoughts tokens:", response.usage_metadata.thoughts_token_count)
            print("Output tokens:", response.usage_metadata.candidates_token_count)

        # If normal text returned
        if response.text:
            return response.text

        # Debug info if no text
        print("[Debug] No text returned by Gemini.")
        if response.candidates:
            cand = response.candidates[0]
            print("Finish reason:", cand.finish_reason)
            if hasattr(cand, "safety_ratings"):
                print("Safety ratings:", cand.safety_ratings)

        return None

    except Exception as e:
        print(f"[Error] Gemini SDK failed: {e}")
        return None

# -------------------------
# Function: Process Directory
# -------------------------
def process_directory():
    if not API_KEY or API_KEY.startswith("your_gemini_api_key"):
        print("[Error] API key not set. Please update the script with a valid key.")
        sys.exit(1)

    total_start = time.time()
    file_count = 0
    fail_count = 0

    for root, dirs, files in os.walk(INPUT_DIR):
        rel_path = os.path.relpath(root, INPUT_DIR)
        output_subdir = os.path.join(OUTPUT_DIR, rel_path)
        os.makedirs(output_subdir, exist_ok=True)

        md_files = [f for f in files if f.endswith(".md")]
        for filename in sorted(md_files):
            input_path = os.path.join(root, filename)
            output_filename = filename.replace(".md", "_notes.md")
            output_path = os.path.join(output_subdir, output_filename)

            if os.path.exists(output_path):
                print(f"[Skipped] Already processed → {output_filename}")
                continue

            with open(input_path, "r", encoding="utf-8") as f:
                note_text = f.read().strip()
                if not note_text:
                    print(f"[Skipped] Empty file → {filename}")
                    continue

            print(f"[Processing] {input_path}")
            start_time = time.time()
            elaborated = elaborate_note(note_text)
            duration = time.time() - start_time

            if elaborated:
                with open(output_path, "w", encoding="utf-8") as out:
                    out.write(elaborated)
                print(f"[Done] {filename} in {duration:.2f} sec")
                file_count += 1
            else:
                print(f"[Error] API did not return a note for {filename}. Skipping.")
                with open("skipped.log", "a", encoding="utf-8") as log:
                    log.write(f"{datetime.now()} | {filename} | No output (see debug above)\n")
                fail_count += 1

                if fail_count >= 2:
                    print(f"[Fatal] 2 files failed. Exiting early.")
                    sys.exit(1)

            # Reverse countdown for 60 seconds
            print("[Sleep] Waiting 60 seconds before next file...")
            for i in range(60, 0, -1):
                print(f"  → {i} sec remaining", end="\r")
                time.sleep(1)
            print(" " * 30, end="\r")  # Clear countdown line

    total_duration = time.time() - total_start
    print(f"\n[Summary] Processed {file_count} new file(s). Failed: {fail_count}")
    print(f"[Complete] Total time: {total_duration:.2f} seconds.")

# -------------------------
# Main
# -------------------------
if __name__ == "__main__":
    process_directory()