#!/usr/bin/env python3
"""
Извлечение тегов по категориям из анализа главы через Ollama (шаг 3).
Вход: merge.json (анализ), вход_главы.json (текст главы), allowed_tags.json (допустимые теги).
Выход: один JSON-файл (tags по категориям, proposed).
"""

import argparse
import json
import re
import sys
import time
import urllib.request
from pathlib import Path

OLLAMA_URL = "http://localhost:11434"
MODEL = "qwen3:14b-q8_0"
DIR = Path(__file__).resolve().parent
DEFAULT_MERGE = DIR.parent / "1_анализ_главы" / "merge.json"
DEFAULT_CHAPTER = DIR.parent / "1_анализ_главы" / "вход_главы.json"
DEFAULT_ALLOWED_TAGS = DIR / "allowed_tags.json"

OLLAMA_OPTIONS = {
    "temperature": 0.3,
    "num_ctx": 9000,
    "repeat_penalty": 1.1,
}

PROMPT_FILE = "extract_tags.txt"


def load_json(path: Path) -> dict:
    """Загружает JSON из файла."""
    with open(path, encoding="utf-8") as f:
        return json.load(f)


def load_prompt(filename: str) -> str:
    """Загружает шаблон промпта из файла."""
    with open(DIR / filename, encoding="utf-8") as f:
        return f.read()


def substitute_prompt(
    prompt: str,
    book_title: str,
    chapter_title: str,
    framework_str: str,
    insights_str: str,
    application_str: str,
    limitations_str: str,
    chapter_text: str,
    allowed_tags_str: str,
) -> str:
    """Подставляет в промпт все поля."""
    return (
        prompt.replace("{book_title}", book_title)
        .replace("{chapter_title}", chapter_title)
        .replace("{framework}", framework_str)
        .replace("{insights}", insights_str)
        .replace("{application}", application_str)
        .replace("{limitations}", limitations_str)
        .replace("{chapter_text}", chapter_text)
        .replace("{allowed_tags_json}", allowed_tags_str)
    )


def extract_json_from_response(text: str) -> dict:
    """Достаёт JSON из ответа модели (может быть обёрнут в ```json ... ```)."""
    text = text.strip()
    match = re.search(r"```(?:json)?\s*([\s\S]*?)\s*```", text)
    if match:
        text = match.group(1).strip()
    return json.loads(text)


def call_ollama(prompt: str) -> str:
    """Вызывает Ollama /api/chat и возвращает content ответа."""
    body = json.dumps(
        {
            "model": MODEL,
            "messages": [{"role": "user", "content": prompt}],
            "stream": False,
            "format": "json",
            "options": OLLAMA_OPTIONS,
            "keep_alive": 0,
        },
        ensure_ascii=False,
    ).encode("utf-8")
    req = urllib.request.Request(
        f"{OLLAMA_URL}/api/chat",
        data=body,
        headers={"Content-Type": "application/json"},
        method="POST",
    )
    try:
        with urllib.request.urlopen(req, timeout=None) as resp:
            data = json.load(resp)
        return data.get("message", {}).get("content", "")
    except urllib.error.HTTPError as e:
        body_b = b""
        if e.fp:
            try:
                body_b = e.fp.read()[:1000]
            except Exception:
                pass
        raise RuntimeError(
            f"Ollama HTTP {e.code}: {e.reason}. Body: {body_b.decode('utf-8', errors='replace')}"
        ) from e


def main() -> int:
    """Загружает данные, вызывает извлечение тегов, пишет результат в JSON."""
    parser = argparse.ArgumentParser(
        description="Извлечение тегов по категориям из анализа главы через Ollama (шаг 3)."
    )
    parser.add_argument(
        "--merge",
        type=Path,
        default=DEFAULT_MERGE,
        help=f"Путь к merge.json (по умолчанию: {DEFAULT_MERGE})",
    )
    parser.add_argument(
        "--chapter",
        type=Path,
        default=DEFAULT_CHAPTER,
        help=f"Путь к вход_главы.json (по умолчанию: {DEFAULT_CHAPTER})",
    )
    parser.add_argument(
        "--allowed-tags",
        type=Path,
        default=DEFAULT_ALLOWED_TAGS,
        help=f"Путь к JSON со списком допустимых тегов по категориям (по умолчанию: allowed_tags.json). Формат: объект с ключами principle, psychology, method, result, context и массивами строк.",
    )
    parser.add_argument(
        "-o",
        "--output",
        type=Path,
        default=DIR / "extracted_tags.json",
        help="Путь к выходному JSON (по умолчанию: extracted_tags.json)",
    )
    args = parser.parse_args()

    if not args.merge.is_file():
        print(f"Файл не найден: {args.merge}", file=sys.stderr)
        return 1
    if not args.chapter.is_file():
        print(f"Файл не найден: {args.chapter}", file=sys.stderr)
        return 1

    print("Загрузка merge.json, вход_главы.json и allowed_tags...")
    merge = load_json(args.merge)
    chapter = load_json(args.chapter)

    if args.allowed_tags.is_file():
        allowed_tags = load_json(args.allowed_tags)
    else:
        allowed_tags = {}
        print(
            f"Файл {args.allowed_tags} не найден; используется пустой список тегов (модель вернёт только proposed).",
            file=sys.stderr,
        )

    book_title = chapter.get("book_title", "")
    chapter_title = chapter.get("chapter_title", "")
    chapter_text = chapter.get("chapter_text", "")

    framework_str = json.dumps(merge.get("framework", {}), ensure_ascii=False, indent=2)
    insights_str = json.dumps(merge.get("insights", []), ensure_ascii=False, indent=2)
    application_str = json.dumps(
        merge.get("application", {}), ensure_ascii=False, indent=2
    )
    limitations_str = json.dumps(
        merge.get("limitations", []), ensure_ascii=False, indent=2
    )
    allowed_tags_str = json.dumps(allowed_tags, ensure_ascii=False, indent=2)

    prompt_tpl = load_prompt(PROMPT_FILE)
    prompt = substitute_prompt(
        prompt_tpl,
        book_title,
        chapter_title,
        framework_str,
        insights_str,
        application_str,
        limitations_str,
        chapter_text,
        allowed_tags_str,
    )

    print(f"Вызов Ollama {MODEL} — извлечение тегов...")
    t0 = time.monotonic()
    try:
        raw = call_ollama(prompt)
    except Exception as e:
        print(f"Ошибка вызова Ollama: {e}", file=sys.stderr)
        return 1
    elapsed = time.monotonic() - t0
    print(f"Ответ получен за {elapsed:.1f} сек ({elapsed / 60:.1f} мин)")

    try:
        result = extract_json_from_response(raw)
    except json.JSONDecodeError as e:
        print(f"Не удалось распарсить JSON ответа: {e}", file=sys.stderr)
        print("Первые 500 символов ответа:", raw[:500], file=sys.stderr)
        return 1

    if not isinstance(result, dict):
        print(
            f"Ожидался объект JSON, получен: {type(result).__name__}",
            file=sys.stderr,
        )
        return 1

    args.output.parent.mkdir(parents=True, exist_ok=True)
    with open(args.output, "w", encoding="utf-8") as f:
        json.dump(result, f, ensure_ascii=False, indent=2)
    print(f"Записано: {args.output}")
    return 0


if __name__ == "__main__":
    sys.exit(main())
