#!/usr/bin/env python3
"""
Скрипт загружает .txt файлы из lists/, разбивает по разделителю,
отправляет каждую запись в Ollama, собирает JSON в res/.
"""

import json
import os
import pathlib
import shutil
import sys
import time

from ollama import Client

# --- Конфигурация ---
LISTS_DIR = pathlib.Path(__file__).parent / "lists"
RES_DIR = pathlib.Path(__file__).parent / "res"
DONE_DIR = pathlib.Path(__file__).parent / "done"
OLLAMA_HOST = os.environ.get("OLLAMA_HOST", "http://localhost:11434")
MODEL = os.environ.get("MODEL", "mithunrajm06/gemma-3n-e4b:latest")
NUM_CTX = int(os.environ.get("NUM_CTX", "4096"))
NUM_PREDICT = int(os.environ.get("NUM_PREDICT", "512"))
RETRY_COUNT = 3
RETRY_DELAY = 5
LOG_STREAM = os.environ.get("LOG_STREAM", "").lower() in ("1", "true", "yes")
# показывать вход и выход по каждой записи (0/false — отключить)
SHOW_IO = os.environ.get("SHOW_IO", "1").lower() in ("1", "true", "yes")

# Разделители записей при разбиении файла (\n — перенос строки, \r\n — Windows)
RECORD_DELIMITERS = ("\n", "\r\n", "\r")

EXTRACT_PROMPT = """Extract brand, model, engine_code, date_range from the text.
If 1 record — return JSON object: {"brand":"","model":"","engine_code":"","date_range":""}
If 2+ records on the line — return JSON array: [{"brand":"",...},{"brand":"",...}]
Empty fields as "". No other text."""


def prepare_records(content: str) -> list[str]:
    """1) Подготовка списка записей из файла по разделителям."""
    for sep in RECORD_DELIMITERS:
        content = content.replace(sep, RECORD_DELIMITERS[0])
    return [r.strip() for r in content.split(RECORD_DELIMITERS[0]) if r.strip()]


def _normalize_obj(obj: dict, original: str) -> dict:
    result = {
        "brand": obj.get("brand", ""),
        "model": obj.get("model", ""),
        "engine_code": obj.get("engine_code", ""),
        "date_range": obj.get("date_range", ""),
    }
    if not result["engine_code"].strip():
        result["has_error"] = True
        result["original"] = original.replace("\t", " ")
    return result


def extract_record(client: Client, record: str) -> list[dict]:
    """2) Отправка промпт + 1 запись в нейронку. 3) Получение ответа (1 или более объектов)."""
    for attempt in range(1, RETRY_COUNT + 1):
        try:
            response = client.chat(
                model=MODEL,
                messages=[
                    {"role": "system", "content": EXTRACT_PROMPT},
                    {"role": "user", "content": record},
                ],
                options={"num_ctx": NUM_CTX, "num_predict": NUM_PREDICT},
            )
            result = response.get("message", {}).get("content", "")
            if "```" in result:
                result = result.split("```")[1]
                if result.strip().lower().startswith("json"):
                    result = result.strip()[4:].strip()
            parsed = json.loads(result)
            if isinstance(parsed, list):
                return [_normalize_obj(o, record) for o in parsed]
            return [_normalize_obj(parsed, record)]
        except (json.JSONDecodeError, KeyError):
            if attempt < RETRY_COUNT:
                time.sleep(RETRY_DELAY)
            else:
                return [{"brand": "", "model": "", "engine_code": "", "date_range": record, "has_error": True, "original": record.replace("\t", " ")}]
        except Exception:
            if attempt < RETRY_COUNT:
                time.sleep(RETRY_DELAY)
            else:
                raise
    return [{"brand": "", "model": "", "engine_code": "", "date_range": record, "has_error": True, "original": record.replace("\t", " ")}]


def process_file(client: Client, filepath: pathlib.Path) -> bool:
    """Обработка файла: подготовка → запрос по записи → добавление в JSON."""
    content = filepath.read_text(encoding="utf-8")
    records = prepare_records(content)
    out_path = RES_DIR / filepath.with_suffix(".json").name

    if not records:
        print(f"  {filepath.name} — 0 записей")
        return False

    out_path.write_text("[]", encoding="utf-8")
    total = len(records)
    t0 = time.perf_counter()
    data = []

    for i, record in enumerate(records, 1):
        print(f"  [{i}/{total}] ", end="", flush=True)
        objs = extract_record(client, record)
        data.extend(objs)
        out_path.write_text(json.dumps(data, ensure_ascii=False, indent=2), encoding="utf-8")
        if SHOW_IO:
            inp = record[:60] + "…" if len(record) > 60 else record
            out = json.dumps(objs, ensure_ascii=False)
            print(f"вход: {inp}", flush=True)
            print(f"      выход: {out}", flush=True)
        else:
            print("✓", flush=True)
    elapsed = time.perf_counter() - t0
    print(f"  {filepath.name} — {len(data)} объектов, {elapsed:.1f}s — OK")
    return True


def main() -> None:
    RES_DIR.mkdir(exist_ok=True)
    DONE_DIR.mkdir(exist_ok=True)

    client = Client(host=OLLAMA_HOST)
    txt_files = sorted(LISTS_DIR.glob("*.txt"))

    if not txt_files:
        print(f"Нет .txt файлов в {LISTS_DIR}")
        return

    print(f"Файлов к обработке: {len(txt_files)}")

    for filepath in txt_files:
        if process_file(client, filepath):
            shutil.move(str(filepath), DONE_DIR / filepath.name)


if __name__ == "__main__":
    main()