"""
Conversation Extractor v2 — извлекает RICH знания из разговоров.

Извлекает:
1. Entities (person, project, technology, company, concept)
2. Facts — короткие утверждения
3. Relations — связи между entities
4. Knowledge — решения, формулы, рецепты, конфиги, команды (с артефактами)

Knowledge — это killer feature. LLM сам определяет тип знания:
  [solution] — решение проблемы (код, конфиг)
  [formula] — формула, уравнение
  [treatment] — лечение, назначение
  [experiment] — результат эксперимента
  [recipe] — рецепт (кулинария, процесс)
  [decision] — принятое решение
  [command] — полезная команда / инструкция
  [reference] — ссылка, источник
  [insight] — наблюдение, инсайт
  [example] — пример, кейс
  ... любой другой тип который подходит по смыслу
"""

import sys
import json
from dataclasses import dataclass, field
from typing import Optional

from engine.extractor.llm_client import LLMClient


EXTRACTION_PROMPT = """You are a knowledge extraction system. Analyze the conversation and extract IMPORTANT, LASTING knowledge.

Return ONLY valid JSON without markdown. Be selective — quality over quantity.

CRITICAL RULES:
- ONLY extract entities that are NAMED and SPECIFIC (real names, real projects, real technologies)
- If user says "I" or "me" — resolve to their ACTUAL NAME if known from context, otherwise use "User"
- DO NOT create entities for: generic concepts (e.g. "OOM error", "debugging"), temporary states, error messages, vague terms
- DO NOT create entities for things that are just part of a fact about another entity
- entity_type: person, project, technology, company, concept (use concept SPARINGLY — only for important recurring themes)
- facts: short, specific, LASTING statements (preferences, roles, decisions — NOT temporary actions)
- relations: how entities are connected
- knowledge: solutions, code, formulas, commands with artifacts

GOOD entities: "Ali Baizhanov", "Mengram", "Redis", "Uzum Bank", "Kubernetes"
BAD entities: "OOM error", "User", "connection pool", "debugging session", "migration"

GOOD facts: "works as backend developer", "prefers dark mode", "lives in Almaty"
BAD facts: "asked about OOM error", "is debugging something", "sent a message"

Response format (strict JSON, no ```):
{{
  "entities": [
    {{
      "name": "Entity Name",
      "type": "person|project|technology|company|concept",
      "facts": ["fact 1", "fact 2"]
    }}
  ],
  "relations": [
    {{
      "from": "Entity 1",
      "to": "Entity 2",
      "type": "works_at|uses|member_of|related_to|depends_on|created_by",
      "description": "relationship description"
    }}
  ],
  "knowledge": [
    {{
      "entity": "Entity this knowledge belongs to",
      "type": "solution|formula|command|insight|...",
      "title": "Short descriptive title",
      "content": "Detailed explanation",
      "artifact": "code/config/formula/command (optional, null if none)"
    }}
  ]
}}

CONVERSATION:
{conversation}

Extract knowledge (return ONLY JSON):"""


@dataclass
class ExtractedEntity:
    """Извлечённая сущность"""
    name: str
    entity_type: str  # person, project, technology, company, concept
    facts: list[str] = field(default_factory=list)

    def __repr__(self):
        return f"Entity({self.entity_type}: {self.name}, facts={len(self.facts)})"


@dataclass
class ExtractedRelation:
    """Извлечённая связь"""
    from_entity: str
    to_entity: str
    relation_type: str
    description: str = ""

    def __repr__(self):
        return f"Relation({self.from_entity} --{self.relation_type}--> {self.to_entity})"


@dataclass
class ExtractedKnowledge:
    """Извлечённое знание — solution, formula, command, etc."""
    entity: str           # к какой entity относится
    knowledge_type: str   # solution, formula, treatment, command, insight, ...
    title: str            # краткий заголовок
    content: str          # подробное описание
    artifact: Optional[str] = None  # код, конфиг, формула, команда

    def __repr__(self):
        has_artifact = "📎" if self.artifact else ""
        return f"Knowledge([{self.knowledge_type}] {self.title} → {self.entity} {has_artifact})"


@dataclass
class ExtractionResult:
    """Результат извлечения знаний из разговора"""
    entities: list[ExtractedEntity] = field(default_factory=list)
    relations: list[ExtractedRelation] = field(default_factory=list)
    knowledge: list[ExtractedKnowledge] = field(default_factory=list)
    raw_response: str = ""

    def __repr__(self):
        return (
            f"ExtractionResult(entities={len(self.entities)}, "
            f"relations={len(self.relations)}, "
            f"knowledge={len(self.knowledge)})"
        )


class ConversationExtractor:
    """Извлекает структурированные знания из разговоров"""

    def __init__(self, llm_client: LLMClient):
        self.llm = llm_client

    def extract(self, conversation: list[dict]) -> ExtractionResult:
        conv_text = self._format_conversation(conversation)
        prompt = EXTRACTION_PROMPT.format(conversation=conv_text)
        raw_response = self.llm.complete(prompt)
        return self._parse_response(raw_response)

    def extract_from_text(self, text: str) -> ExtractionResult:
        conversation = [{"role": "user", "content": text}]
        return self.extract(conversation)

    def _format_conversation(self, conversation: list[dict]) -> str:
        lines = []
        for msg in conversation:
            role = "User" if msg["role"] == "user" else "Assistant"
            lines.append(f"{role}: {msg['content']}")
        return "\n\n".join(lines)

    def _parse_response(self, raw: str) -> ExtractionResult:
        result = ExtractionResult(raw_response=raw)

        # Clean markdown
        clean = raw.strip()
        if clean.startswith("```"):
            lines = clean.split("\n")
            clean = "\n".join(lines[1:-1] if lines[-1].strip() == "```" else lines[1:])

        try:
            data = json.loads(clean)
        except json.JSONDecodeError:
            start = raw.find("{")
            end = raw.rfind("}") + 1
            if start >= 0 and end > start:
                try:
                    data = json.loads(raw[start:end])
                except json.JSONDecodeError:
                    print(f"⚠️  Failed to parse JSON from LLM", file=sys.stderr)
                    return result
            else:
                print(f"⚠️  LLM returned no JSON", file=sys.stderr)
                return result

        # Entities
        for e in data.get("entities", []):
            result.entities.append(ExtractedEntity(
                name=e.get("name", "Unknown"),
                entity_type=e.get("type", "concept"),
                facts=e.get("facts", []),
            ))

        # Relations
        for r in data.get("relations", []):
            result.relations.append(ExtractedRelation(
                from_entity=r.get("from", ""),
                to_entity=r.get("to", ""),
                relation_type=r.get("type", "related_to"),
                description=r.get("description", ""),
            ))

        # Knowledge (NEW)
        for k in data.get("knowledge", []):
            result.knowledge.append(ExtractedKnowledge(
                entity=k.get("entity", ""),
                knowledge_type=k.get("type", "insight"),
                title=k.get("title", ""),
                content=k.get("content", ""),
                artifact=k.get("artifact"),
            ))

        return result


# --- Mock for testing ---

class MockLLMClient(LLMClient):
    """Mock LLM for testing without API"""

    def complete(self, prompt: str, system: str = "") -> str:
        return json.dumps({
            "entities": [
                {
                    "name": "User",
                    "type": "person",
                    "facts": [
                        "Работает backend разработчиком",
                        "Работает в Uzum Bank",
                        "Основной стек: Java, Spring Boot"
                    ]
                },
                {
                    "name": "Uzum Bank",
                    "type": "company",
                    "facts": ["Банк в Узбекистане", "Микросервисная архитектура"]
                },
                {
                    "name": "Проект Alpha",
                    "type": "project",
                    "facts": ["Backend сервис для платежей", "Проблема с connection pool"]
                },
                {
                    "name": "PostgreSQL",
                    "type": "technology",
                    "facts": ["Основная БД", "Версия 15"]
                },
                {
                    "name": "Spring Boot",
                    "type": "technology",
                    "facts": ["Основной фреймворк для микросервисов"]
                }
            ],
            "relations": [
                {"from": "User", "to": "Uzum Bank", "type": "works_at", "description": "Backend разработчик"},
                {"from": "User", "to": "Проект Alpha", "type": "member_of", "description": "Работает над проектом"},
                {"from": "Проект Alpha", "to": "PostgreSQL", "type": "uses", "description": "Основная БД"},
                {"from": "Проект Alpha", "to": "Spring Boot", "type": "uses", "description": "Backend фреймворк"},
                {"from": "Uzum Bank", "to": "Проект Alpha", "type": "related_to", "description": "Проект банка"}
            ],
            "knowledge": [
                {
                    "entity": "PostgreSQL",
                    "type": "solution",
                    "title": "Connection pool exhaustion fix",
                    "content": "OOM при 200+ WebSocket соединениях. Каждый WS держал отдельный connection. Решение: Redis кеш для UserService и BlockedAccountService.",
                    "artifact": "spring.datasource.hikari.maximum-pool-size: 20\nspring.datasource.hikari.idle-timeout: 30000\nspring.datasource.hikari.connection-timeout: 5000"
                },
                {
                    "entity": "PostgreSQL",
                    "type": "command",
                    "title": "Check active connections",
                    "content": "Мониторинг активных соединений PostgreSQL",
                    "artifact": "SELECT count(*), state FROM pg_stat_activity GROUP BY state;"
                }
            ]
        }, ensure_ascii=False)
