"""Skill execution — LLM calls, JSON extraction, and run logging."""

from __future__ import annotations

import json
import re
import time

import asyncpg
from anthropic import APIStatusError, AsyncAnthropic, AuthenticationError
from jinja2 import Template

from loom.config import SkillsConfig
from loom.exceptions import SkillError
from loom.skills.loader import SkillDefinition


def render_prompt(skill: SkillDefinition, inputs: dict) -> str:
    """Render a skill's Jinja2 prompt template with the given inputs."""
    return Template(skill.prompt_template).render(**inputs)


def _repair_truncated_json(text: str) -> str | None:
    """Try to repair truncated JSON by closing open brackets/braces.

    Strips any trailing partial value (incomplete string, number, etc.)
    then appends the necessary closing characters.
    """
    # Find the first { or [
    for opener, closer in [("{", "}"), ("[", "]")]:
        start = text.find(opener)
        if start == -1:
            continue

        fragment = text[start:]
        # Strip trailing partial tokens (incomplete strings, trailing commas, colons)
        fragment = re.sub(r',\s*$', '', fragment)
        fragment = re.sub(r':\s*$', ': null', fragment)
        # Close any open string literal
        if fragment.count('"') % 2 == 1:
            fragment += '"'

        # Count open brackets/braces and close them
        stack = []
        in_string = False
        escape = False
        for ch in fragment:
            if escape:
                escape = False
                continue
            if ch == '\\' and in_string:
                escape = True
                continue
            if ch == '"':
                in_string = not in_string
                continue
            if in_string:
                continue
            if ch in ('{', '['):
                stack.append('}' if ch == '{' else ']')
            elif ch in ('}', ']') and stack:
                stack.pop()

        if stack:
            # Strip trailing comma before closing
            fragment = re.sub(r',\s*$', '', fragment)
            fragment += ''.join(reversed(stack))
            try:
                return json.loads(fragment)
            except (json.JSONDecodeError, ValueError):
                pass
    return None


def _extract_json(text: str, truncated: bool = False):
    """Extract JSON from LLM response text.

    Tries up to four strategies:
    1. Full text as JSON
    2. ```json fenced code block
    3. First {...} or [...] match (balanced brackets)
    4. Truncated JSON repair (only when truncated=True)
    """
    # Strategy 1: full text
    stripped = text.strip()
    try:
        return json.loads(stripped)
    except (json.JSONDecodeError, ValueError):
        pass

    # Strategy 2: fenced code block
    match = re.search(r"```(?:json)?\s*\n(.*?)\n```", text, re.DOTALL)
    if match:
        try:
            return json.loads(match.group(1).strip())
        except (json.JSONDecodeError, ValueError):
            pass

    # Strategy 3: first { ... } or [ ... ]
    for opener, closer in [("{", "}"), ("[", "]")]:
        start = text.find(opener)
        if start == -1:
            continue
        depth = 0
        for i in range(start, len(text)):
            if text[i] == opener:
                depth += 1
            elif text[i] == closer:
                depth -= 1
            if depth == 0:
                try:
                    return json.loads(text[start : i + 1])
                except (json.JSONDecodeError, ValueError):
                    break

    # Strategy 4: truncated JSON repair
    if truncated:
        repaired = _repair_truncated_json(text)
        if repaired is not None:
            return repaired

    raise SkillError("json_extraction", "Could not extract JSON from LLM response")


async def run_skill(
    skill: SkillDefinition,
    inputs: dict,
    config: SkillsConfig,
    pool: asyncpg.Pool | None = None,
    project_id: str | None = None,
) -> dict:
    """Execute a skill: render prompt, call LLM, extract JSON, optionally log."""
    import logging
    log = logging.getLogger(__name__)

    prompt = render_prompt(skill, inputs)

    model = skill.model or config.model
    temperature = skill.temperature if skill.temperature is not None else config.temperature
    max_tokens = skill.max_tokens or config.max_tokens

    try:
        client = AsyncAnthropic(api_key=config.api_key or None)
    except AuthenticationError as exc:
        raise SkillError(skill.name, f"Anthropic API key is missing or invalid. Set ANTHROPIC_API_KEY or LOOM_SKILLS_API_KEY. ({exc})") from exc
    start = time.monotonic()

    try:
        response = await client.messages.create(
            model=model,
            max_tokens=max_tokens,
            temperature=temperature,
            messages=[{"role": "user", "content": prompt}],
        )
    except (AuthenticationError, APIStatusError) as exc:
        if isinstance(exc, AuthenticationError) or (isinstance(exc, APIStatusError) and exc.status_code == 401):
            raise SkillError(skill.name, f"Anthropic API key is missing or invalid. Set ANTHROPIC_API_KEY or LOOM_SKILLS_API_KEY. ({exc})") from exc
        raise

    duration_ms = int((time.monotonic() - start) * 1000)
    response_text = response.content[0].text
    tokens_used = response.usage.input_tokens + response.usage.output_tokens
    truncated = response.stop_reason == "max_tokens"

    # If truncated, retry once with doubled max_tokens
    if truncated:
        retry_max = min(max_tokens * 2, 16384)
        if retry_max > max_tokens:
            log.warning(
                "Skill '%s' hit max_tokens (%d), retrying with %d",
                skill.name, max_tokens, retry_max,
            )
            retry_start = time.monotonic()
            response = await client.messages.create(
                model=model,
                max_tokens=retry_max,
                temperature=temperature,
                messages=[{"role": "user", "content": prompt}],
            )
            duration_ms += int((time.monotonic() - retry_start) * 1000)
            response_text = response.content[0].text
            tokens_used += response.usage.input_tokens + response.usage.output_tokens
            truncated = response.stop_reason == "max_tokens"

    output = _extract_json(response_text, truncated=truncated)

    # Log to skill_runs table if pool and project_id are available
    if pool and project_id:
        from loom.graph.store import record_skill_run
        await record_skill_run(
            pool, project_id, skill.name, inputs, output,
            model=model, tokens_used=tokens_used, duration_ms=duration_ms,
        )

    return output
