"""
parser.py - LLM-powered JSON parsing with retry and validation

The framework first routes a message to a discovered feature module, then
uses that module's focused prompt and Python validator to parse its action.
Both synchronous and asynchronous entrypoints are kept for reusable modules.
"""

import json
import logging
import os
from pathlib import Path
import re
from datetime import datetime, timezone

from dotenv import load_dotenv
from openai import AsyncOpenAI, OpenAI


logger = logging.getLogger(__name__)

load_dotenv(Path(__file__).resolve().parents[1] / ".env", override=False)

CONFIG_PATH = os.environ.get(
    "AI_CONFIG_PATH", os.path.join(os.path.dirname(__file__), "ai_config.json")
)

with open(CONFIG_PATH, "r", encoding="utf-8") as config_file:
    AI_CONFIG = json.load(config_file)

VALIDATORS = {}
_sync_client = None
_async_client = None


def _get_client(async_client=False):
    global _sync_client, _async_client
    client_class = AsyncOpenAI if async_client else OpenAI
    current = _async_client if async_client else _sync_client
    if current is None:
        current = client_class(
            api_key=os.getenv("OPENROUTER_API_KEY"),
            base_url=os.getenv(
                "OPENROUTER_BASE_URL", "https://openrouter.ai/api/v1"
            ),
        )
        if async_client:
            _async_client = current
        else:
            _sync_client = current
    return current


def _extract_json_from_text(text):
    """Decode the first complete JSON object, including nested objects."""
    if not isinstance(text, str):
        return None

    stripped = text.strip()
    fence = chr(96) * 3
    if stripped.startswith(fence):
        stripped = re.sub(
            rf"^{re.escape(fence)}(?:json)?\s*",
            "",
            stripped,
            flags=re.IGNORECASE,
        )
        stripped = re.sub(rf"\s*{re.escape(fence)}$", "", stripped)

    decoder = json.JSONDecoder()
    for index, character in enumerate(stripped):
        if character not in "[{":
            continue
        try:
            value, _ = decoder.raw_decode(stripped[index:])
            return value
        except json.JSONDecodeError:
            continue
    return None


def _render_template(template, values):
    """Replace named placeholders without treating literal JSON braces as fields."""
    pattern = re.compile(r"\{([A-Za-z_][A-Za-z0-9_]*)\}")

    def replace(match):
        key = match.group(1)
        return str(values[key]) if key in values else match.group(0)

    return pattern.sub(replace, template)


def _response_text(response):
    if not response.choices:
        return None
    message = response.choices[0].message
    if message.content:
        return message.content.strip()
    reasoning = getattr(message, "reasoning", None)
    return reasoning.strip() if reasoning else None


def _request_args(system_prompt, user_prompt):
    args = {
        "model": os.getenv("AI_MODEL", AI_CONFIG["model"]),
        "max_tokens": AI_CONFIG.get("max_tokens", 2048),
        "timeout": AI_CONFIG["validation"].get("timeout_seconds", 15),
        "messages": [
            {"role": "system", "content": system_prompt},
            {"role": "user", "content": user_prompt},
        ],
    }
    if AI_CONFIG.get("json_mode", False):
        args["response_format"] = {"type": "json_object"}
    return args


def _call_llm(system_prompt, user_prompt):
    """Call an OpenAI-compatible API and return response text or None."""
    try:
        response = _get_client().chat.completions.create(
            **_request_args(system_prompt, user_prompt)
        )
        return _response_text(response)
    except Exception as error:
        logger.warning("LLM call failed: %s: %s", type(error).__name__, error)
        return None


async def _call_llm_async(system_prompt, user_prompt):
    """Asynchronously call an OpenAI-compatible API."""
    try:
        response = await _get_client(async_client=True).chat.completions.create(
            **_request_args(system_prompt, user_prompt)
        )
        return _response_text(response)
    except Exception as error:
        logger.warning("LLM call failed: %s: %s", type(error).__name__, error)
        return None


def _history_context(history):
    if not history:
        return "No previous context"

    history_lines = []
    for index, (message, result) in enumerate(history[-3:]):
        history_lines.append(f"{index + 1}. User: {message}")
        history_lines.append(f"   Parsed: {json.dumps(result, default=str)}")
    return "\n".join(history_lines)


def _build_prompt(user_input, prompt_config, history, errors, template_values):
    values = {
        "user_input": user_input,
        "history_context": _history_context(history),
    }
    values.update(template_values or {})
    user_prompt = _render_template(prompt_config["user_template"], values)
    if errors:
        user_prompt += (
            "\n\nThe previous response failed validation:\n- "
            + "\n- ".join(str(error) for error in errors)
            + "\nReturn a corrected JSON object."
        )
    return user_prompt


def _get_prompt(interaction_type, prompt_override=None):
    if prompt_override:
        return prompt_override
    return AI_CONFIG.get("prompts", {}).get(interaction_type)


def _validation_errors(parsed, validator):
    if not isinstance(parsed, dict):
        return ["Response must be a JSON object"]
    if validator:
        return list(validator(parsed) or [])
    return []


def parse(
    user_input,
    interaction_type,
    retry_count=0,
    errors=None,
    history=None,
    prompt_override=None,
    validator=None,
    template_values=None,
):
    """Synchronously parse one prompt into a validated dictionary."""
    prompt_config = _get_prompt(interaction_type, prompt_override)
    if not prompt_config:
        return {"error": f"Unknown interaction type: {interaction_type}"}

    validator = validator or VALIDATORS.get(interaction_type)
    max_attempts = AI_CONFIG["validation"].get("max_retries", 3)
    attempt = retry_count
    current_errors = errors

    while attempt < max_attempts:
        user_prompt = _build_prompt(
            user_input, prompt_config, history, current_errors, template_values
        )
        response_text = _call_llm(prompt_config["system"], user_prompt)
        if not response_text:
            return {"error": "AI service unavailable", "user_input": user_input}

        parsed = _extract_json_from_text(response_text)
        current_errors = (
            ["Response was not valid JSON"]
            if parsed is None
            else _validation_errors(parsed, validator)
        )
        if not current_errors:
            return parsed
        attempt += 1

    return {
        "error": f"Failed to parse after {max_attempts} attempts",
        "validation_errors": current_errors or [],
        "user_input": user_input,
    }


async def parse_async(
    user_input,
    interaction_type,
    retry_count=0,
    errors=None,
    history=None,
    prompt_override=None,
    validator=None,
    template_values=None,
):
    """Asynchronously parse one prompt into a validated dictionary."""
    prompt_config = _get_prompt(interaction_type, prompt_override)
    if not prompt_config:
        return {"error": f"Unknown interaction type: {interaction_type}"}

    validator = validator or VALIDATORS.get(interaction_type)
    max_attempts = AI_CONFIG["validation"].get("max_retries", 3)
    attempt = retry_count
    current_errors = errors

    while attempt < max_attempts:
        user_prompt = _build_prompt(
            user_input, prompt_config, history, current_errors, template_values
        )
        response_text = await _call_llm_async(prompt_config["system"], user_prompt)
        if not response_text:
            return {"error": "AI service unavailable", "user_input": user_input}

        parsed = _extract_json_from_text(response_text)
        current_errors = (
            ["Response was not valid JSON"]
            if parsed is None
            else _validation_errors(parsed, validator)
        )
        if not current_errors:
            return parsed
        attempt += 1

    return {
        "error": f"Failed to parse after {max_attempts} attempts",
        "validation_errors": current_errors or [],
        "user_input": user_input,
    }


async def parse_command_async(user_input, module_registry, history=None, timezone_name="UTC"):
    """Route a command, then parse it with the selected feature prompt."""
    now = datetime.now(timezone.utc).isoformat()
    template_values = {
        "module_context": module_registry.router_context(),
        "current_time": now,
        "timezone": timezone_name,
    }
    routed = await parse_async(
        user_input,
        "command_parser",
        history=history,
        template_values=template_values,
    )
    if routed.get("error") or routed.get("needs_clarification"):
        return routed

    threshold = AI_CONFIG["validation"].get("confidence_threshold", 0.8)
    confidence = routed.get("confidence")
    if isinstance(confidence, (int, float)) and confidence < threshold:
        return {
            "needs_clarification": "Could you rephrase that with a little more detail?",
            "confidence": confidence,
        }

    interaction_type = routed.get("interaction_type")
    command = module_registry.get_command(interaction_type)
    if not command:
        return {"error": f"Unknown command type: {interaction_type}"}

    parsed = await parse_async(
        user_input,
        interaction_type,
        history=history,
        prompt_override=command["prompt"],
        validator=command["validator"],
        template_values=template_values,
    )
    if isinstance(parsed, dict) and not parsed.get("error"):
        parsed["interaction_type"] = interaction_type
    return parsed


def register_validator(interaction_type, validator_fn):
    """Keep the original validator registration API for direct parser users."""
    VALIDATORS[interaction_type] = validator_fn