import requests
import json
import time
import re
import sqlite3
import sys
from pathlib import Path


OLLAMA_URL = "http://localhost:11434/api/generate"
MODEL = "mistral"

MAX_CHARS = 4000
RETRIES = 3


with open("tags.txt", "r", encoding="utf8") as file:
    TAG_VOCAB = ", ".join(file.read().strip().split("\n"))


PROMPT_TEMPLATE = """
You are a strict document tagging system.

Your job is to assign tags to the following markdown note.

## RULES

- Choose tags from this allowed list: {allowed}
- Use ONLY tags from the allowed list
- Do NOT invent new tags
- Prefer general topics over niche ones
- If unsure, choose fewer tags, not more

## OUTPUT FORMAT

Return ONLY valid JSON:

{
  "tags": ["tag1", "tag2", ...]
}

No explanation. No extra text.

## NOTE

\"\"\"
{note}
\"\"\"
"""


def call_ollama(prompt):
    response = requests.post(
        OLLAMA_URL,
        json={
            "model": MODEL,
            "prompt": prompt,
            "stream": False,
            "options": {
                "temperature": 0.2
            }
        },
        timeout=120
    )
    response.raise_for_status()
    return response.json()["response"]


def extract_json(text):
    """Extract JSON even if model adds garbage"""
    try:
        return json.loads(text)
    except:
        match = re.search(r"\{.*\}", text, re.DOTALL)
        if match:
            return json.loads(match.group())
        raise ValueError("No valid JSON found")


def tag_note(content: str) -> list[str]:
    content = content[:MAX_CHARS]

    prompt = PROMPT_TEMPLATE.replace("{note}", content).replace("{allowed}", TAG_VOCAB)

    for attempt in range(RETRIES):
        try:
            raw = call_ollama(prompt)
            data = extract_json(raw)

            tags = data.get("tags", [])

            if not isinstance(tags, list):
                raise ValueError("Tags not a list")

            return tags

        except Exception as e:
            print(f"[retry {attempt+1}] error:", e)
            time.sleep(1)

    return []


conn = sqlite3.connect("db.sqlite3")
cursor = conn.cursor()

note_count = next(cursor.execute("SELECT COUNT(id) FROM orgapy_document "))[0]

with open("tagged.tsv", "w", encoding="utf8") as file:
    file.write(f"doc_id\ttitle\ttags\n")

for i, row in enumerate(list(cursor.execute("SELECT id, title, content FROM orgapy_document WHERE type = 'note' "))):
    doc_id, title, content = row
    categories = [
        r[0] for r in cursor.execute("SELECT orgapy_category.name FROM orgapy_document_categories JOIN orgapy_category ON orgapy_category.id = orgapy_document_categories.category_id WHERE orgapy_document_categories.document_id = ?", (doc_id,))
    ]
    categories_string = "\n".join(f"  - {name}" for name in categories)
    full_content = f"""---\ntitle: {title}\ncategories:\n{categories_string}\n---\n\n""" + content
    print(f"\n[{i+1}/{note_count}] {doc_id}. {title}")
    tags = tag_note(full_content)
    print(tags)
    with open("tagged.tsv", "a", encoding="utf8") as file:
        file.write(f"{doc_id}\t{title}\t{';'.join(tags)}\n")
