# === ANCHOR: IMPORTS_START ===
import sqlite3
import json
import urllib.request
import urllib.error
from typing import List, Dict, Any, Tuple
from datetime import datetime

DB_NAME = "context_memory.db"
OLLAMA_URL = "http://192.168.11.3:11434/api/chat"
MODEL_NAME = "qwen2.5:14b"
# === ANCHOR: IMPORTS_END ===


# === ANCHOR: DB_TOOLS_START ===
def db_get_tasks(status: str = None) -> str:
    """Извлекает список задач из реестра БД."""
    conn = sqlite3.connect(DB_NAME)
    conn.row_factory = sqlite3.Row
    cursor = conn.cursor()
    if status:
        cursor.execute("SELECT * FROM tasks WHERE status = ? ORDER BY task_id", (status.upper(),))
    else:
        cursor.execute("SELECT * FROM tasks ORDER BY task_id")
    rows = [dict(r) for r in cursor.fetchall()]
    conn.close()
    return json.dumps(rows, ensure_ascii=False)


def db_update_task_status(task_id: str, status: str = None, due_date: str = None) -> str:
    """Обновляет статус и/или срок выполнения задачи."""
    conn = sqlite3.connect(DB_NAME)
    cursor = conn.cursor()
    updates = []
    params = []

    if status:
        status_upper = status.upper()
        updates.append("status = ?")
        params.append(status_upper)
        if status_upper == "COMPLETED":
            updates.append("completed_at = CURRENT_TIMESTAMP")
        else:
            updates.append("completed_at = NULL")

    if due_date:
        updates.append("due_date = ?")
        params.append(due_date)

    if not updates:
        conn.close()
        return json.dumps({"error": "Не указаны параметры для обновления"}, ensure_ascii=False)

    params.append(task_id.upper())
    query = f"UPDATE tasks SET {', '.join(updates)} WHERE task_id = ?"
    cursor.execute(query, tuple(params))

    if cursor.rowcount == 0:
        conn.close()
        return json.dumps({"error": f"Задача {task_id} не найдена"}, ensure_ascii=False)

    conn.commit()
    conn.close()
    return json.dumps({"status": "success", "message": f"Задача {task_id} обновлена"}, ensure_ascii=False)


def db_add_task(task_id: str, module: str, title: str, priority: str = "HIGH", due_date: str = None) -> str:
    """Добавляет новую задачу в реестр."""
    conn = sqlite3.connect(DB_NAME)
    cursor = conn.cursor()
    try:
        cursor.execute(
            "INSERT INTO tasks (task_id, module, title, status, priority, due_date) VALUES (?, ?, ?, 'BACKLOG', ?, ?)",
            (task_id.upper(), module, title, priority.upper(), due_date)
        )
        conn.commit()
        res = {"status": "success", "message": f"Задача {task_id} создана"}
    except sqlite3.IntegrityError:
        res = {"error": f"Задача с ID {task_id} уже существует"}
    finally:
        conn.close()
    return json.dumps(res, ensure_ascii=False)


def db_get_rules() -> str:
    """Возвращает действующие архитектурные правила."""
    conn = sqlite3.connect(DB_NAME)
    conn.row_factory = sqlite3.Row
    cursor = conn.cursor()
    cursor.execute("SELECT category, rule_text FROM architecture_memory WHERE is_active = 1")
    rows = [dict(r) for r in cursor.fetchall()]
    conn.close()
    return json.dumps(rows, ensure_ascii=False)
# === ANCHOR: DB_TOOLS_END ===


# === ANCHOR: TOOLS_SCHEMA_START ===
TOOLS_SCHEMA = [
    {
        "type": "function",
        "function": {
            "name": "db_get_tasks",
            "description": "Получить список задач проекта с возможностью фильтрации по статусу (BACKLOG, IN_PROGRESS, COMPLETED).",
            "parameters": {
                "type": "object",
                "properties": {
                    "status": {
                        "type": "string",
                        "description": "Фильтр статуса: BACKLOG, IN_PROGRESS или COMPLETED."
                    }
                }
            }
        }
    },
    {
        "type": "function",
        "function": {
            "name": "db_update_task_status",
            "description": "Изменить статус и/или срок выполнения (due_date) задачи в реестре.",
            "parameters": {
                "type": "object",
                "properties": {
                    "task_id": {"type": "string", "description": "Идентификатор задачи, например TASK-17"},
                    "status": {"type": "string", "description": "Новый статус: COMPLETED, IN_PROGRESS или BACKLOG"},
                    "due_date": {"type": "string", "description": "Срок выполнения задачи, например '2026-08-07 14:00'"}
                },
                "required": ["task_id"]
            }
        }
    },
    {
        "type": "function",
        "function": {
            "name": "db_add_task",
            "description": "Добавить новую задачу в бэклог проекта.",
            "parameters": {
                "type": "object",
                "properties": {
                    "task_id": {"type": "string", "description": "Уникальный ID задачи, например TASK-18"},
                    "module": {"type": "string", "description": "Модуль проекта, например services/scud_export"},
                    "title": {"type": "string", "description": "Краткое описание задачи"},
                    "priority": {"type": "string", "description": "Приоритет: HIGH, MEDIUM, LOW"},
                    "due_date": {"type": "string", "description": "Срок выполнения задачи, например '2026-08-07 14:00'"}
                },
                "required": ["task_id", "module", "title"]
            }
        }
    },
    {
        "type": "function",
        "function": {
            "name": "db_get_rules",
            "description": "Получить список архитектурных правил и ограничений проекта.",
            "parameters": {"type": "object", "properties": {}}
        }
    }
]
# === ANCHOR: TOOLS_SCHEMA_END ===


# === ANCHOR: CHAT_AGENT_START ===
def process_chat_message(user_message: str, chat_history: List[Dict[str, Any]] = None) -> Tuple[str, List[Dict[str, Any]]]:
    if chat_history is None:
        chat_history = []

    current_now = datetime.now().strftime("%Y-%m-%d %H:%M")

    system_prompt = {
        "role": "system",
        "content": (
            f"Текущая дата и время сервера: {current_now}.\n"
            f"Текущая дата и время сервера: {current_now}.\n"
            "Ты — ИИ-ассистент проекта SCUD Orion AI. "
            "У тебя есть доступ к функциям для просмотра, добавления и изменения статусов и сроков задач (due_date) в БД. "
            "Когда пользователь просит показать задачи, добавить новую или изменить статус — обязательно вызывай нужную функцию. "
            "После вызова функции давай пользователю понятное, лаконичное резюме результата на русском языке. "
            "При запросе списка всех задач не расписывай каждую подробно, а давай краткую сводку по количествам, "
            "так как детали отображаются в выезжающем реестре задач."
        )
    }

    messages = [system_prompt] + chat_history + [{"role": "user", "content": user_message}]

    payload = {
    "model": MODEL_NAME,
    "messages": messages,
    "tools": TOOLS_SCHEMA,
    "stream": False,
    "options": {
        "num_predict": 2048,   # Достаточно для любых длинных списков и кода
        "num_ctx": 8192,       # Расширенная память диалога без потери скорости
        "temperature": 0.2,    # Низкая температура: ответы точные и без галлюцинаций
        "top_p": 0.9,          # Оптимальная выборка токенов
        "num_gpu": 99          # Заставляет Ollama загрузить 100% слоев модели в VRAM
    }
}

    req = urllib.request.Request(
        OLLAMA_URL,
        data=json.dumps(payload).encode("utf-8"),
        headers={"Content-Type": "application/json"}
    )

    try:
        with urllib.request.urlopen(req) as response:
            res_data = json.loads(response.read().decode("utf-8"))
            msg = res_data.get("message", {})

            tool_calls = msg.get("tool_calls", [])
            
            # Резервная проверка, если Ollama прислала вызов функции текстом
            content_str = msg.get("content", "")
            if not tool_calls and "name" in content_str and "db_" in content_str:
                try:
                    parsed = json.loads(content_str)
                    if "name" in parsed:
                        tool_calls = [{"function": parsed}]
                except Exception:
                    pass

            if tool_calls:
                messages.append(msg)
                for tool in tool_calls:
                    fn_name = tool["function"]["name"]
                    fn_args = tool["function"].get("arguments", {})

                    if fn_name == "db_get_tasks":
                        result = db_get_tasks(status=fn_args.get("status"))
                    elif fn_name == "db_update_task_status":
                        result = db_update_task_status(
                            task_id=fn_args.get("task_id"),
                            status=fn_args.get("status"),
                            due_date=fn_args.get("due_date")
                        )
                    elif fn_name == "db_add_task":
                        result = db_add_task(
                            task_id=fn_args.get("task_id"),
                            module=fn_args.get("module"),
                            title=fn_args.get("title"),
                            priority=fn_args.get("priority", "HIGH"),
                            due_date=fn_args.get("due_date")
                        )
                    elif fn_name == "db_get_rules":
                        result = db_get_rules()
                    else:
                        result = json.dumps({"error": "Неизвестная функция"})

                    messages.append({
                        "role": "tool",
                        "content": result
                    })

                # Вторым запросом передаем результат обратно модели для текстового ответа
                payload = {
                "model": MODEL_NAME,
                "messages": messages,
                "tools": TOOLS_SCHEMA,
                "stream": False,
                "options": {
                    "num_predict": 2048,   # Достаточно для любых длинных списков и кода
                    "num_ctx": 8192,       # Расширенная память диалога без потери скорости
                    "temperature": 0.2,    # Низкая температура: ответы точные и без галлюцинаций
                    "top_p": 0.9,          # Оптимальная выборка токенов
                    "num_gpu": 99          # Заставляет Ollama загрузить 100% слоев модели в VRAM
                }
            }
                req_second = urllib.request.Request(
                    OLLAMA_URL,
                    data=json.dumps(payload_second).encode("utf-8"),
                    headers={"Content-Type": "application/json"}
                )
                with urllib.request.urlopen(req_second) as resp_second:
                    final_res = json.loads(resp_second.read().decode("utf-8"))
                    final_msg = final_res.get("message", {}).get("content", "")
                    
                    updated_history = chat_history + [
                        {"role": "user", "content": user_message},
                        {"role": "assistant", "content": final_msg}
                    ]
                    return final_msg, updated_history

            else:
                final_msg = msg.get("content", "")
                updated_history = chat_history + [
                    {"role": "user", "content": user_message},
                    {"role": "assistant", "content": final_msg}
                ]
                return final_msg, updated_history

    except urllib.error.URLError as e:
        return f"Ошибка связи с Ollama ({OLLAMA_URL}): {e}", chat_history
# === ANCHOR: CHAT_AGENT_END ===