# Copyright (c) 2024-2026 YuyAI / InsAIts Team. All rights reserved.
# Proprietary and confidential. See LICENSE.premium for terms.
"""
InsAIts Premium - Adaptive Dictionary System
=============================================
Domain-specific dictionaries, adaptive jargon learning,
persistence, and auto-expansion via LLM.
"""

import json
import re
import time
import logging
from pathlib import Path
from typing import Any, Dict, List, Optional, Set
from collections import defaultdict, Counter

from ..security import (
    encrypt_data,
    decrypt_data,
    get_or_create_encryption_key,
    set_secure_file_permissions,
    set_secure_dir_permissions,
)
from ..config import INSAITS_CACHE_DIR_NAME, ENCRYPTION_KEY_FILENAME

logger = logging.getLogger(__name__)

# Try to import LLM for auto-expansion
try:
    from ..local_llm import ollama_chat
    LLM_AVAILABLE = True
except Exception:
    LLM_AVAILABLE = False

# Default cache directory for persisted dictionaries
INSAITS_CACHE_DIR = Path.home() / INSAITS_CACHE_DIR_NAME
JARGON_FILE = INSAITS_CACHE_DIR / "jargon.bin"
ENCRYPTION_KEY_PATH = INSAITS_CACHE_DIR / ENCRYPTION_KEY_FILENAME


# ============================================
# Domain-specific dictionaries
# ============================================

DOMAIN_DICTIONARIES: Dict[str, Dict[str, Any]] = {
    "finance": {
        "known": {
            "EBITDA", "P&L", "ROI", "CAGR", "AUM", "NAV", "YTD", "QTD", "MTD",
            "WACC", "DCF", "NPV", "IRR", "EPS", "PE", "PB", "ROIC", "ROCE",
            "CAPEX", "OPEX", "FCF", "LBO", "M&A", "IPO", "SPV", "ABS", "MBS",
            "CDO", "CDS", "OTC", "FX", "FRA", "IRS", "LIBOR", "SOFR", "EURIBOR"
        },
        "expansions": {
            "EBITDA": "Earnings Before Interest, Taxes, Depreciation, and Amortization",
            "WACC": "Weighted Average Cost of Capital",
            "DCF": "Discounted Cash Flow",
            "NPV": "Net Present Value",
            "IRR": "Internal Rate of Return",
            "EPS": "Earnings Per Share",
            "CAPEX": "Capital Expenditure",
            "OPEX": "Operating Expenditure",
            "FCF": "Free Cash Flow",
            "LBO": "Leveraged Buyout",
            "ABS": "Asset-Backed Securities",
            "MBS": "Mortgage-Backed Securities",
            "CDO": "Collateralized Debt Obligation",
            "CDS": "Credit Default Swap"
        }
    },
    "healthcare": {
        "known": {
            "HIPAA", "PHI", "EHR", "EMR", "ICD", "CPT", "DRG", "HMO", "PPO",
            "PCP", "RN", "MD", "DO", "NP", "PA", "ICU", "ER", "OR", "NICU",
            "FDA", "CDC", "WHO", "NIH", "CMS", "HRSA", "ONC", "AHRQ",
            "SNOMED", "LOINC", "HL7", "FHIR", "ADT", "CCR", "CCD", "CCDA"
        },
        "expansions": {
            "HIPAA": "Health Insurance Portability and Accountability Act",
            "PHI": "Protected Health Information",
            "EHR": "Electronic Health Record",
            "EMR": "Electronic Medical Record",
            "ICD": "International Classification of Diseases",
            "CPT": "Current Procedural Terminology",
            "DRG": "Diagnosis-Related Group",
            "FHIR": "Fast Healthcare Interoperability Resources",
            "HL7": "Health Level Seven International",
            "SNOMED": "Systematized Nomenclature of Medicine"
        }
    },
    "kubernetes": {
        "known": {
            "K8S", "POD", "HELM", "KUBECTL", "CRD", "HPA", "VPA", "PVC", "PV",
            "RBAC", "CNI", "CSI", "CRI", "ETCD", "ISTIO", "ENVOY", "NGINX",
            "DAEMONSET", "STATEFULSET", "REPLICASET", "CONFIGMAP", "SECRET",
            "INGRESS", "EGRESS", "SVC", "NS", "SA", "CM", "NODEPORT", "LB"
        },
        "expansions": {
            "K8S": "Kubernetes",
            "HPA": "Horizontal Pod Autoscaler",
            "VPA": "Vertical Pod Autoscaler",
            "PVC": "Persistent Volume Claim",
            "PV": "Persistent Volume",
            "RBAC": "Role-Based Access Control",
            "CNI": "Container Network Interface",
            "CSI": "Container Storage Interface",
            "CRI": "Container Runtime Interface",
            "CRD": "Custom Resource Definition"
        }
    },
    "machine_learning": {
        "known": {
            "BERT", "GPT", "LSTM", "RNN", "CNN", "GAN", "VAE", "AE", "MLP",
            "SOTA", "FLOPS", "TPU", "CUDA", "ONNX", "RLHF", "DPO", "PPO",
            "ADAM", "SGD", "MSE", "MAE", "BCE", "NLL", "KL", "BLEU", "ROUGE",
            "F1", "AUC", "ROC", "PR", "AP", "MAP", "IOU", "YOLO", "RCNN"
        },
        "expansions": {
            "BERT": "Bidirectional Encoder Representations from Transformers",
            "GPT": "Generative Pre-trained Transformer",
            "LSTM": "Long Short-Term Memory",
            "RLHF": "Reinforcement Learning from Human Feedback",
            "DPO": "Direct Preference Optimization",
            "SOTA": "State of the Art",
            "ONNX": "Open Neural Network Exchange",
            "BLEU": "Bilingual Evaluation Understudy",
            "ROUGE": "Recall-Oriented Understudy for Gisting Evaluation"
        }
    },
    "devops": {
        "known": {
            "CI", "CD", "CICD", "IaC", "SRE", "SLO", "SLI", "SLA", "MTTR",
            "MTTF", "MTBF", "RTO", "RPO", "DR", "HA", "FT", "LB", "CDN",
            "APM", "ELK", "SIEM", "WAF", "DDoS", "TLS", "mTLS", "PKI",
            "OIDC", "SAML", "OAuth", "JWT", "JWK", "JWKS"
        },
        "expansions": {
            "CI": "Continuous Integration",
            "CD": "Continuous Deployment/Delivery",
            "IaC": "Infrastructure as Code",
            "SRE": "Site Reliability Engineering",
            "SLO": "Service Level Objective",
            "SLI": "Service Level Indicator",
            "MTTR": "Mean Time To Recovery",
            "MTTF": "Mean Time To Failure",
            "RTO": "Recovery Time Objective",
            "RPO": "Recovery Point Objective"
        }
    },
    "quantum": {
        "known": {
            "QUBIT", "QPU", "QC", "QML", "NISQ", "VQE", "QAOA", "QFT", "QPE",
            "CNOT", "SWAP", "TOFFOLI", "HADAMARD", "PAULI", "BLOCH",
            "IBM", "IBMQ", "CIRQ", "QISKIT", "PENNYLANE", "BRAKET"
        },
        "expansions": {
            "QUBIT": "Quantum Bit",
            "QPU": "Quantum Processing Unit",
            "QML": "Quantum Machine Learning",
            "NISQ": "Noisy Intermediate-Scale Quantum",
            "VQE": "Variational Quantum Eigensolver",
            "QAOA": "Quantum Approximate Optimization Algorithm",
            "QFT": "Quantum Fourier Transform",
            "QPE": "Quantum Phase Estimation"
        }
    },
    "legal": {
        "known": {
            # Privacy & Data Protection
            "GDPR", "CCPA", "CPRA", "FADP", "LGPD", "PIPL", "PIPA", "PECR",
            "DPO", "DPA", "DPIA", "SCC", "BCR", "TIA", "ROPA",
            "EDPB", "EDPS", "CNIL", "ICO", "BfDI",
            # Compliance & Standards
            "SOX", "SOC", "PCI", "DSS", "ISO", "NIST", "FERPA", "COPPA",
            "HIPAA", "AML", "KYC", "CDD", "EDD", "SAR", "CTR", "BSA",
            # Courts & Tribunals
            "SCOTUS", "CJEU", "ECHR", "ICJ", "ICC", "WTO", "ICSID",
            "CAFC", "TTAB", "ITC", "PTAB",
            # Legal Documents & Agreements
            "NDA", "MOU", "LOI", "TOS", "EULA", "SLA", "DPA", "MSA",
            "SOW", "RFP", "RFQ", "IP", "IPR",
            # Intellectual Property
            "USPTO", "EPO", "WIPO", "PCT", "FRAND", "SEP", "DMCA",
            # International Law
            "UNCITRAL", "UNIDROIT", "OECD", "FATF", "FSB", "BIS",
            # EU Specific
            "TFEU", "TEU", "ECFR", "OJ", "CELEX",
            # German Law
            "StGB", "BGB", "HGB", "ZPO", "StPO", "GG", "BRAO", "BDSG",
            "BGH", "BVerfG", "BVerwG", "BAG", "BSG",
            # Swiss Law
            "OR", "ZGB", "DSG", "BGE", "BV",
        },
        "expansions": {
            "GDPR": "General Data Protection Regulation",
            "CCPA": "California Consumer Privacy Act",
            "CPRA": "California Privacy Rights Act",
            "FADP": "Federal Act on Data Protection (Switzerland)",
            "LGPD": "Lei Geral de Protecao de Dados (Brazil)",
            "DPO": "Data Protection Officer",
            "DPA": "Data Processing Agreement / Data Protection Authority",
            "DPIA": "Data Protection Impact Assessment",
            "SCC": "Standard Contractual Clauses",
            "BCR": "Binding Corporate Rules",
            "TIA": "Transfer Impact Assessment",
            "ROPA": "Record of Processing Activities",
            "EDPB": "European Data Protection Board",
            "SOX": "Sarbanes-Oxley Act",
            "AML": "Anti-Money Laundering",
            "KYC": "Know Your Customer",
            "CDD": "Customer Due Diligence",
            "SAR": "Suspicious Activity Report",
            "BSA": "Bank Secrecy Act",
            "SCOTUS": "Supreme Court of the United States",
            "CJEU": "Court of Justice of the European Union",
            "ECHR": "European Court of Human Rights",
            "ICJ": "International Court of Justice",
            "ICC": "International Criminal Court",
            "ICSID": "International Centre for Settlement of Investment Disputes",
            "NDA": "Non-Disclosure Agreement",
            "MOU": "Memorandum of Understanding",
            "EULA": "End-User License Agreement",
            "MSA": "Master Service Agreement",
            "USPTO": "United States Patent and Trademark Office",
            "EPO": "European Patent Office",
            "WIPO": "World Intellectual Property Organization",
            "PCT": "Patent Cooperation Treaty",
            "FRAND": "Fair, Reasonable, and Non-Discriminatory",
            "DMCA": "Digital Millennium Copyright Act",
            "UNCITRAL": "United Nations Commission on International Trade Law",
            "FATF": "Financial Action Task Force",
            "StGB": "Strafgesetzbuch (German Criminal Code)",
            "BGB": "Buergerliches Gesetzbuch (German Civil Code)",
            "BGH": "Bundesgerichtshof (German Federal Court of Justice)",
            "BVerfG": "Bundesverfassungsgericht (German Federal Constitutional Court)",
            "BRAO": "Bundesrechtsanwaltsordnung (German Federal Lawyers Act)",
            "BDSG": "Bundesdatenschutzgesetz (German Federal Data Protection Act)",
            "BGE": "Bundesgerichtsentscheid (Swiss Federal Court Decision)",
            "DSG": "Datenschutzgesetz (Swiss Data Protection Act)",
        }
    },
    "coding": {
        "known": {
            # Languages & Runtimes
            "JS", "TS", "PY", "RB", "GO", "RS", "WASM", "JVM", "CLR", "V8",
            # Frontend
            "HTML", "CSS", "JSX", "TSX", "DOM", "VDOM", "SPA", "PWA", "SSR",
            "SSG", "ISR", "CSR", "SCSS", "SASS", "LESS", "SVG",
            # Backend & APIs
            "REST", "GRPC", "RPC", "SOAP", "WS", "WSS", "AMQP", "MQTT",
            "GRAPHQL", "TRPC", "SSE",
            # Databases
            "SQL", "NOSQL", "ORM", "ODM", "ACID", "CAP", "CRUD", "OLAP",
            "OLTP", "ETL", "ELT", "CDC",
            # Architecture & Patterns
            "DDD", "CQRS", "SOA", "MSA", "EDA", "SAGA", "BFF",
            "MVC", "MVVM", "MVP", "DI", "IOC",
            # Testing
            "TDD", "BDD", "E2E", "UAT", "QA", "SAST", "DAST", "IAST",
            # Auth & Security
            "JWT", "OAUTH", "OIDC", "CORS", "XSS", "CSRF", "SQLI", "CSP",
            "MFA", "SSO", "TOTP", "SAML", "RBAC", "ABAC",
            # DevOps & Cloud
            "AWS", "GCP", "AZURE", "PAAS", "IAAS", "SAAS", "FAAS", "BAAS",
            "CDN", "DNS", "SSL", "TLS", "LB",
            # Package & Build
            "NPM", "YARN", "PNPM", "PIP", "CARGO", "MAVEN", "GRADLE",
            "VITE", "WEBPACK", "ESM", "CJS", "UMD", "IIFE",
            # Data Formats
            "JSON", "XML", "YAML", "TOML", "CSV", "PROTOBUF", "AVRO",
            "REGEX", "UUID", "BASE64", "UTF8",
            # Version Control
            "GIT", "SVN", "PR", "MR", "SHA", "GPG",
        },
        "expansions": {
            "WASM": "WebAssembly",
            "JVM": "Java Virtual Machine",
            "SPA": "Single Page Application",
            "PWA": "Progressive Web Application",
            "SSR": "Server-Side Rendering",
            "SSG": "Static Site Generation",
            "ISR": "Incremental Static Regeneration",
            "CSR": "Client-Side Rendering",
            "GRPC": "Google Remote Procedure Call",
            "GRAPHQL": "Graph Query Language",
            "SSE": "Server-Sent Events",
            "AMQP": "Advanced Message Queuing Protocol",
            "MQTT": "Message Queuing Telemetry Transport",
            "ORM": "Object-Relational Mapping",
            "ACID": "Atomicity, Consistency, Isolation, Durability",
            "CAP": "Consistency, Availability, Partition Tolerance",
            "ETL": "Extract, Transform, Load",
            "CDC": "Change Data Capture",
            "DDD": "Domain-Driven Design",
            "CQRS": "Command Query Responsibility Segregation",
            "SOA": "Service-Oriented Architecture",
            "MSA": "Microservices Architecture",
            "EDA": "Event-Driven Architecture",
            "BFF": "Backend for Frontend",
            "DI": "Dependency Injection",
            "IOC": "Inversion of Control",
            "TDD": "Test-Driven Development",
            "BDD": "Behavior-Driven Development",
            "E2E": "End-to-End Testing",
            "SAST": "Static Application Security Testing",
            "DAST": "Dynamic Application Security Testing",
            "XSS": "Cross-Site Scripting",
            "CSRF": "Cross-Site Request Forgery",
            "SQLI": "SQL Injection",
            "CSP": "Content Security Policy",
            "MFA": "Multi-Factor Authentication",
            "SSO": "Single Sign-On",
            "TOTP": "Time-Based One-Time Password",
            "RBAC": "Role-Based Access Control",
            "ABAC": "Attribute-Based Access Control",
            "PAAS": "Platform as a Service",
            "IAAS": "Infrastructure as a Service",
            "SAAS": "Software as a Service",
            "FAAS": "Function as a Service",
            "BAAS": "Backend as a Service",
            "ESM": "ECMAScript Modules",
            "CJS": "CommonJS",
            "PROTOBUF": "Protocol Buffers",
            "UUID": "Universally Unique Identifier",
        }
    },
    "universal": {
        "known": {
            # Agent & AI Communication
            "RAG", "COT", "TOT", "ICL", "PEFT", "LORA", "QLORA", "SFT",
            "MCTS", "GGUF", "GGML", "GPTQ", "AWQ", "FP16", "BF16", "INT8",
            "INT4", "KV", "MHA", "MQA", "GQA", "RoPE", "ALiBi",
            # Orchestration & Agents
            "DAG", "FSM", "BFS", "DFS", "FIFO", "LIFO", "LISP",
            "LLM", "SLM", "VLM", "MLM", "NER", "POS", "NLU", "NLG",
            # Communication Protocols
            "HTTP", "HTTPS", "WS", "WSS", "TCP", "UDP", "IP", "SMTP",
            "IMAP", "POP3", "FTP", "SFTP", "SSH",
            # Encoding & Serialization
            "JSON", "XML", "YAML", "CSV", "BASE64", "UTF8", "ASCII",
            "GZIP", "ZSTD", "LZ4", "BROTLI",
            # Operating Systems & Compute
            "OS", "CPU", "GPU", "TPU", "NPU", "RAM", "ROM", "SSD", "HDD",
            "VM", "VPS", "NUMA", "SIMD", "AVX",
            # Project & Process
            "AGILE", "SCRUM", "KANBAN", "SPRINT", "STANDUP",
            "PM", "QA", "UAT", "SIT",
            # Common Abbreviations
            "ETA", "FYI", "TBD", "TBA", "WIP", "EOD", "COB", "OOO",
            "ASAP", "IMO", "TLDR", "POC", "MVP", "MLP",
            # Metrics & Observability
            "KPI", "OKR", "SLA", "SLO", "SLI", "MTTR", "MTTF", "MTBF",
            "APM", "RUM", "SLO",
            # Auth & Identity
            "IAM", "SSO", "MFA", "OAUTH", "OIDC", "JWT", "ACL",
            # Data & Storage
            "DB", "SQL", "NOSQL", "BLOB", "CDN", "S3", "GCS",
        },
        "expansions": {
            "RAG": "Retrieval-Augmented Generation",
            "COT": "Chain of Thought",
            "TOT": "Tree of Thought",
            "ICL": "In-Context Learning",
            "PEFT": "Parameter-Efficient Fine-Tuning",
            "LORA": "Low-Rank Adaptation",
            "QLORA": "Quantized Low-Rank Adaptation",
            "SFT": "Supervised Fine-Tuning",
            "MCTS": "Monte Carlo Tree Search",
            "GGUF": "GPT-Generated Unified Format",
            "GPTQ": "GPT Quantization",
            "AWQ": "Activation-Aware Weight Quantization",
            "MHA": "Multi-Head Attention",
            "MQA": "Multi-Query Attention",
            "GQA": "Grouped-Query Attention",
            "RoPE": "Rotary Position Embedding",
            "ALiBi": "Attention with Linear Biases",
            "KV": "Key-Value (cache)",
            "DAG": "Directed Acyclic Graph",
            "FSM": "Finite State Machine",
            "SLM": "Small Language Model",
            "VLM": "Vision Language Model",
            "NER": "Named Entity Recognition",
            "NLU": "Natural Language Understanding",
            "NLG": "Natural Language Generation",
            "ZSTD": "Zstandard Compression",
            "NPU": "Neural Processing Unit",
            "NUMA": "Non-Uniform Memory Access",
            "SIMD": "Single Instruction Multiple Data",
            "OKR": "Objectives and Key Results",
            "APM": "Application Performance Monitoring",
            "RUM": "Real User Monitoring",
            "IAM": "Identity and Access Management",
            "ACL": "Access Control List",
            "TLDR": "Too Long Didn't Read",
        }
    }
}


class AdaptiveDictionary:
    """
    Adaptive jargon dictionary with domain-specific knowledge,
    persistence, auto-learning, and LLM-based expansion.

    This is a premium feature that provides:
    - Auto-promotion of frequently seen terms
    - Domain dictionary loading (universal, coding, legal, finance, healthcare, k8s, ML, devops, quantum)
    - Dictionary import/export to JSON
    - LLM-based term expansion via Ollama
    - Persistent storage in ~/.insaits/jargon.json
    """

    CANDIDATE_PROMOTION_THRESHOLD = 5
    MAX_CANDIDATES = 500

    def __init__(self, auto_learn: bool = True, seed_terms: Optional[Set[str]] = None):
        """
        Initialize the adaptive dictionary.

        Args:
            auto_learn: If True, automatically learn new terms from conversations
            seed_terms: Optional set of seed terms to start with
        """
        self.auto_learn = auto_learn

        self.jargon_dict: Dict[str, Any] = {
            "known": seed_terms.copy() if seed_terms else set(),
            "candidate": defaultdict(int),
            "learned": set(),
            "expanded": {}
        }

        self._load_dict()

    def _load_dict(self) -> None:
        """Load persisted jargon dictionary from encrypted disk cache."""
        try:
            if not JARGON_FILE.exists():
                return

            enc_key = get_or_create_encryption_key(ENCRYPTION_KEY_PATH)
            raw = JARGON_FILE.read_bytes()

            plaintext = decrypt_data(raw, enc_key) if enc_key else raw
            if plaintext is None:
                logger.debug("Jargon file decryption failed — ignoring stale cache")
                return

            data = json.loads(plaintext.decode('utf-8'))

            if "learned" in data:
                self.jargon_dict["learned"] = set(data["learned"])
            if "candidate" in data:
                self.jargon_dict["candidate"] = defaultdict(int, data["candidate"])
            if "expanded" in data:
                self.jargon_dict["expanded"] = data["expanded"]

            logger.info(
                f"Loaded jargon dictionary: "
                f"{len(self.jargon_dict['learned'])} learned terms"
            )
        except Exception:
            logger.warning("Could not load jargon dictionary")

    def _save_dict(self) -> None:
        """Persist jargon dictionary to encrypted disk cache."""
        try:
            INSAITS_CACHE_DIR.mkdir(parents=True, exist_ok=True)
            set_secure_dir_permissions(INSAITS_CACHE_DIR)

            data = {
                "learned": list(self.jargon_dict["learned"]),
                "candidate": dict(self.jargon_dict["candidate"]),
                "expanded": self.jargon_dict["expanded"],
                "last_updated": time.time()
            }

            plaintext = json.dumps(data).encode('utf-8')

            enc_key = get_or_create_encryption_key(ENCRYPTION_KEY_PATH)
            payload = encrypt_data(plaintext, enc_key) if enc_key else plaintext

            JARGON_FILE.write_bytes(payload)
            set_secure_file_permissions(JARGON_FILE)

            logger.debug(
                f"Saved jargon dictionary: "
                f"{len(self.jargon_dict['learned'])} learned terms"
            )
        except Exception:
            logger.error("Could not save jargon dictionary")

    def is_known_term(self, term: str) -> bool:
        """Check if a term is known (seed or learned)."""
        upper = term.upper()
        return (
            upper in self.jargon_dict["known"] or
            upper in self.jargon_dict["learned"]
        )

    def track_candidate(self, term: str) -> bool:
        """
        Track a candidate term and promote if threshold reached.
        Returns True if the term was promoted to learned.
        """
        if not self.auto_learn:
            return False

        upper = term.upper()

        if self.is_known_term(upper):
            return False

        self.jargon_dict["candidate"][upper] += 1

        if self.jargon_dict["candidate"][upper] >= self.CANDIDATE_PROMOTION_THRESHOLD:
            self.jargon_dict["learned"].add(upper)
            del self.jargon_dict["candidate"][upper]
            logger.info(f"Auto-learned new term: {upper}")
            self._save_dict()
            return True

        if len(self.jargon_dict["candidate"]) > self.MAX_CANDIDATES:
            sorted_candidates = sorted(
                self.jargon_dict["candidate"].items(),
                key=lambda x: x[1]
            )
            for term_to_remove, _ in sorted_candidates[:100]:
                del self.jargon_dict["candidate"][term_to_remove]

        return False

    def add_learned_term(self, term: str, expanded: Optional[str] = None) -> None:
        """Manually add a term to the learned dictionary."""
        upper = term.upper()
        self.jargon_dict["learned"].add(upper)
        if expanded:
            self.jargon_dict["expanded"][upper] = expanded
        self._save_dict()
        logger.info(f"Manually added term: {upper}")

    def get_jargon_stats(self) -> Dict:
        """Return statistics about the jargon dictionary."""
        return {
            "known_terms": len(self.jargon_dict["known"]),
            "learned_terms": len(self.jargon_dict["learned"]),
            "candidate_terms": len(self.jargon_dict["candidate"]),
            "expanded_terms": len(self.jargon_dict["expanded"]),
            "top_candidates": dict(
                sorted(
                    self.jargon_dict["candidate"].items(),
                    key=lambda x: x[1],
                    reverse=True
                )[:10]
            )
        }

    # ============================================
    # Domain Dictionary Management
    # ============================================

    def load_domain(self, domain: str) -> Dict:
        """Load a domain-specific dictionary to reduce false positives."""
        if domain not in DOMAIN_DICTIONARIES:
            available = list(DOMAIN_DICTIONARIES.keys())
            return {"error": f"Unknown domain: {domain}", "available_domains": available}

        domain_data = DOMAIN_DICTIONARIES[domain]

        terms_before = len(self.jargon_dict["known"])
        self.jargon_dict["known"].update(domain_data["known"])
        terms_added = len(self.jargon_dict["known"]) - terms_before

        expansions_before = len(self.jargon_dict["expanded"])
        self.jargon_dict["expanded"].update(domain_data.get("expansions", {}))
        expansions_added = len(self.jargon_dict["expanded"]) - expansions_before

        self._save_dict()

        logger.info(f"Loaded domain '{domain}': {terms_added} terms, {expansions_added} expansions")

        return {
            "loaded": domain,
            "terms_added": terms_added,
            "expansions_added": expansions_added,
            "total_known": len(self.jargon_dict["known"]),
            "total_expanded": len(self.jargon_dict["expanded"])
        }

    def get_available_domains(self) -> List[str]:
        """Return list of available domain dictionaries."""
        return list(DOMAIN_DICTIONARIES.keys())

    def export_dictionary(self, filepath: str) -> Dict:
        """Export the current dictionary to a JSON file."""
        data = {
            "known": list(self.jargon_dict["known"]),
            "learned": list(self.jargon_dict["learned"]),
            "expanded": self.jargon_dict["expanded"],
            "metadata": {
                "exported_at": time.time(),
                "version": "2.0",
                "known_count": len(self.jargon_dict["known"]),
                "learned_count": len(self.jargon_dict["learned"]),
                "expanded_count": len(self.jargon_dict["expanded"])
            }
        }

        try:
            with open(filepath, 'w', encoding='utf-8') as f:
                json.dump(data, f, indent=2)
            logger.info(f"Exported dictionary to {filepath}")
            return {
                "exported": filepath,
                "total_terms": len(data["known"]) + len(data["learned"]),
                "known_terms": len(data["known"]),
                "learned_terms": len(data["learned"]),
                "expanded_terms": len(data["expanded"])
            }
        except Exception as e:
            logger.error(f"Failed to export dictionary: {e}")
            return {"error": str(e)}

    def import_dictionary(self, filepath: str, merge: bool = True) -> Dict:
        """Import a dictionary from a JSON file."""
        try:
            with open(filepath, 'r', encoding='utf-8') as f:
                data = json.load(f)

            known_before = len(self.jargon_dict["known"])
            learned_before = len(self.jargon_dict["learned"])
            expanded_before = len(self.jargon_dict["expanded"])

            if merge:
                self.jargon_dict["known"].update(set(data.get("known", [])))
                self.jargon_dict["learned"].update(set(data.get("learned", [])))
                self.jargon_dict["expanded"].update(data.get("expanded", {}))
            else:
                imported_known = set(data.get("known", []))
                self.jargon_dict["known"] = self.jargon_dict["known"].union(imported_known)
                self.jargon_dict["learned"] = set(data.get("learned", []))
                self.jargon_dict["expanded"] = data.get("expanded", {})

            self._save_dict()

            return {
                "imported": filepath,
                "mode": "merge" if merge else "replace",
                "known_added": len(self.jargon_dict["known"]) - known_before,
                "learned_added": len(self.jargon_dict["learned"]) - learned_before,
                "expanded_added": len(self.jargon_dict["expanded"]) - expanded_before,
                "total_known": len(self.jargon_dict["known"]),
                "total_learned": len(self.jargon_dict["learned"]),
                "total_expanded": len(self.jargon_dict["expanded"])
            }
        except FileNotFoundError:
            return {"error": f"File not found: {filepath}"}
        except json.JSONDecodeError as e:
            return {"error": f"Invalid JSON: {e}"}
        except Exception as e:
            logger.error(f"Failed to import dictionary: {e}")
            return {"error": str(e)}

    def auto_expand_terms(
        self,
        terms: Optional[List[str]] = None,
        model: str = "phi3"
    ) -> Dict:
        """Use LLM to automatically expand undefined terms."""
        if not LLM_AVAILABLE:
            return {"error": "LLM not available (Ollama required)"}

        if terms is None:
            terms = [
                t for t in self.jargon_dict["learned"]
                if t not in self.jargon_dict["expanded"]
            ]

        if not terms:
            return {"status": "all_terms_expanded", "expanded": {}}

        expansions = {}
        errors = []

        for term in terms[:20]:
            messages = [{
                "role": "user",
                "content": (
                    f"What does the acronym '{term}' stand for? "
                    f"Reply with ONLY the expansion, nothing else. "
                    f"If unknown, reply 'UNKNOWN'."
                )
            }]

            try:
                response = ollama_chat(messages, model=model, temperature=0.1)
                if response and len(response) < 200 and "UNKNOWN" not in response.upper():
                    expansion = response.strip()
                    expansions[term] = expansion
                    self.jargon_dict["expanded"][term] = expansion
            except Exception as e:
                errors.append({"term": term, "error": str(e)})

        if expansions:
            self._save_dict()

        logger.info(f"Auto-expanded {len(expansions)} terms")

        return {
            "expanded": expansions,
            "count": len(expansions),
            "remaining": len(terms) - len(expansions) - len(errors),
            "errors": errors if errors else None
        }

    # ============================================
    # Session Learning
    # ============================================

    def learn_from_session(
        self,
        history: Dict,
        min_occurrences: int = 3,
        auto_save: bool = True
    ) -> Dict:
        """
        Analyze session messages and learn new jargon terms.

        Args:
            history: The conversation history dict from monitor
            min_occurrences: Minimum times a term must appear to be learned
            auto_save: Whether to persist after learning

        Returns:
            Dict with learning statistics
        """
        all_text = []
        for agent_hist in history.values():
            for llm_msgs in agent_hist.values():
                for msg in llm_msgs:
                    all_text.append(msg["text"])

        if not all_text:
            return {
                "status": "no_data",
                "message": "No messages in session to learn from",
                "terms_learned": 0
            }

        full_text = " ".join(all_text)
        acronyms = re.findall(r'\b[A-Z]{2,}\b', full_text)
        term_counts = Counter(acronyms)

        new_terms = []
        skipped_known = []
        skipped_low_count = []

        for term, count in term_counts.items():
            if self.is_known_term(term):
                skipped_known.append(term)
            elif count < min_occurrences:
                skipped_low_count.append((term, count))
            else:
                self.add_learned_term(term)
                new_terms.append((term, count))

        if auto_save and new_terms:
            self._save_dict()

        return {
            "status": "success",
            "terms_learned": len(new_terms),
            "learned_terms": new_terms,
            "already_known": len(skipped_known),
            "below_threshold": len(skipped_low_count),
            "jargon_stats": self.get_jargon_stats()
        }
