"""
AI Chat Service for NOMAZRA — Cambodian AI Legal Assistant.
Uses gpt-4o-mini (cheapest capable model: $0.15/1M input, $0.60/1M output).
All files are pre-converted to text (via file_converter) before being sent here.
"""
import os
import logging
import tiktoken

logger = logging.getLogger(__name__)

MODEL = 'gpt-4o-mini'
MAX_CONTEXT_TOKENS = 12000   # Keep context tight to save costs
MAX_RESPONSE_TOKENS = 2000
ENCODING = 'cl100k_base'  # tokenizer for gpt-4o-mini

SYSTEM_PROMPT = """You are NOMAZRA, an expert AI legal assistant specializing in Cambodian law.

Your expertise covers:
- Cambodian Civil Code, Criminal Code, Labor Law, Land Law, Family Law
- Constitutional law and human rights protections
- Business/commercial law (Law on Commercial Enterprises)
- Immigration and visa regulations
- Contract law and dispute resolution
- Court procedures and legal processes

Guidelines:
1. Always cite specific Cambodian laws, articles, or prakas (sub-decrees) when applicable.
2. Clearly state when information may be outdated or when professional legal counsel is recommended.
3. Provide practical, actionable advice while noting legal limitations.
4. Respond in the language the user writes in (Khmer or English).
5. For complex matters, break down the legal analysis step by step.
6. Never fabricate law numbers or articles — if unsure, say so.
7. Keep responses concise but thorough.
"""

TOOL_PROMPTS = {
    'law_search': (
        "You are searching Cambodian legal databases. "
        "Provide relevant laws, articles, and explanations for the user's query. "
        "Format results clearly with law name, article number, and summary."
    ),
    'document_review': (
        "You are reviewing a legal document. Analyze the document text provided "
        "and identify: key clauses, potential issues, missing provisions, "
        "and compliance with Cambodian law. Provide a structured review."
    ),
    'rights_check': (
        "Explain the user's legal rights under Cambodian law for this situation. "
        "Cite specific laws and articles. Include practical steps they can take."
    ),
    'contract_draft': (
        "Help draft or review contract clauses under Cambodian law. "
        "Ensure compliance with the Cambodian Civil Code."
    ),
    'glossary': (
        "Explain this legal term in the context of Cambodian law. "
        "Provide: definition, relevant laws, practical examples."
    ),
}


def _get_client():
    """Lazy-init OpenAI client."""
    import openai
    api_key = os.getenv('OPENAI_API_KEY')
    if not api_key:
        raise ValueError("OPENAI_API_KEY not configured")
    return openai.OpenAI(api_key=api_key)


def count_tokens(text):
    """Count tokens for cost estimation."""
    try:
        enc = tiktoken.get_encoding(ENCODING)
        return len(enc.encode(text))
    except Exception:
        # Rough fallback: ~4 chars per token
        return len(text) // 4


def _truncate_to_tokens(text, max_tokens):
    """Truncate text to fit within token budget."""
    try:
        enc = tiktoken.get_encoding(ENCODING)
        tokens = enc.encode(text)
        if len(tokens) <= max_tokens:
            return text
        return enc.decode(tokens[:max_tokens]) + "\n\n[...truncated for length]"
    except Exception:
        # Fallback: rough char estimate
        max_chars = max_tokens * 4
        if len(text) <= max_chars:
            return text
        return text[:max_chars] + "\n\n[...truncated for length]"


def _build_messages(conversation_history, user_message, tool=None, document_text=None):
    """
    Build the messages array for the API call.
    Trims conversation history to fit within token budget.
    """
    system = SYSTEM_PROMPT
    if tool and tool in TOOL_PROMPTS:
        system += f"\n\nCurrent task: {TOOL_PROMPTS[tool]}"

    messages = [{'role': 'system', 'content': system}]

    # Budget: MAX_CONTEXT - system - response - current message
    system_tokens = count_tokens(system)
    budget = MAX_CONTEXT_TOKENS - system_tokens - MAX_RESPONSE_TOKENS

    # If document text provided, allocate half the budget to it
    if document_text:
        doc_budget = budget // 2
        truncated_doc = _truncate_to_tokens(document_text, doc_budget)
        doc_msg = {
            'role': 'user',
            'content': f"[Document content for analysis]:\n\n{truncated_doc}"
        }
        messages.append(doc_msg)
        budget -= count_tokens(truncated_doc)

    # User's current message
    user_tokens = count_tokens(user_message)
    budget -= user_tokens

    # Fill remaining budget with conversation history (most recent first)
    if conversation_history:
        history_messages = []
        used = 0
        for msg in reversed(conversation_history):
            msg_tokens = count_tokens(msg.get('content', ''))
            if used + msg_tokens > budget:
                break
            history_messages.insert(0, msg)
            used += msg_tokens
        messages.extend(history_messages)

    messages.append({'role': 'user', 'content': user_message})
    return messages


def chat(user_message, conversation_history=None, tool=None, document_text=None):
    """
    Send a chat message and get AI response.

    Args:
        user_message: The user's current message text
        conversation_history: List of previous {'role': ..., 'content': ...} dicts
        tool: Optional tool context (law_search, document_review, etc.)
        document_text: Optional pre-extracted document text

    Returns:
        dict with 'response', 'input_tokens', 'output_tokens'
    """
    if not user_message or not user_message.strip():
        return {
            'response': 'Please provide a message or question.',
            'input_tokens': 0,
            'output_tokens': 0,
        }

    try:
        client = _get_client()
        messages = _build_messages(
            conversation_history or [],
            user_message,
            tool=tool,
            document_text=document_text,
        )

        response = client.chat.completions.create(
            model=MODEL,
            messages=messages,
            max_tokens=MAX_RESPONSE_TOKENS,
            temperature=0.3,  # Lower temp for legal accuracy
        )

        choice = response.choices[0]
        return {
            'response': choice.message.content.strip(),
            'input_tokens': response.usage.prompt_tokens,
            'output_tokens': response.usage.completion_tokens,
        }
    except Exception as e:
        logger.error("AI chat error: %s", e)
        return {
            'response': (
                "I'm sorry, I'm having trouble processing your request right now. "
                "Please try again in a moment."
            ),
            'input_tokens': 0,
            'output_tokens': 0,
            'error': str(e),
        }


def summarize_document(document_text, document_name=''):
    """Generate a concise summary of a legal document."""
    prompt = (
        f"Summarize this document{f' ({document_name})' if document_name else ''} "
        "in the context of Cambodian law. Provide:\n"
        "1. Document type and purpose\n"
        "2. Key provisions/clauses\n"
        "3. Notable legal implications\n"
        "4. Any potential issues or missing elements"
    )
    return chat(prompt, document_text=document_text, tool='document_review')


def explain_rights(situation):
    """Explain legal rights for a given situation."""
    return chat(situation, tool='rights_check')


def search_law(query):
    """Search for relevant Cambodian laws."""
    return chat(query, tool='law_search')
