from __future__ import annotations

import re


# Provider import patterns: (old_import_re, bootstrap_replacement_name)
# Each pattern handles both single-line and multi-line (parenthesized) imports:
#   from openai import OpenAI
#   from openai import (\n    OpenAI,\n    AsyncOpenAI,\n)
def _provider_re(module: str) -> re.Pattern[str]:
    escaped = re.escape(module)
    return re.compile(
        rf"^(?:import {escaped}|from {escaped} import \([^)]*\)|from {escaped} import .+)$",
        re.MULTILINE,
    )

_PROVIDER_IMPORTS: list[tuple[re.Pattern[str], str]] = [
    (_provider_re("openai"), "openai"),
    (_provider_re("anthropic"), "anthropic"),
    (_provider_re("google.generativeai"), "gemini"),
    (_provider_re("google.genai"), "genai"),
    (_provider_re("mistralai"), "mistral"),
    (_provider_re("cohere"), "cohere"),
]

_ROUTE_HANDLER_RE = re.compile(
    r"^([ \t]*)@(?:app|router)\.\s*(?:get|post|put|patch|delete|route)\s*\("
    r"|^([ \t]*)@api_view\s*\(",
    re.MULTILINE,
)

_USER_INPUT_PATTERNS = re.compile(
    r"(?:request\.json|request\.body|request\.get_json\(\)"
    r"|await request\.json\(\)"
    r"|data\s*=\s*await\s+request\.json\(\)"
    r"|body\s*=\s*await\s+request\.body\(\)"
    r"|request\.data)",
    re.MULTILINE,
)

_DEF_RE = re.compile(
    r"^([ \t]*)(async\s+)?def\s+(\w+)\s*\([^)]*\)\s*.*?:\s*$",
    re.MULTILINE,
)


def _add_bootstrap_import(source: str, bootstrap_import: str) -> str:
    """Replace provider imports with the bootstrap import, or append after existing imports."""
    result = source
    replaced = False
    for pattern, _name in _PROVIDER_IMPORTS:
        match = pattern.search(result)
        if match:
            if not replaced:
                result = result[:match.start()] + bootstrap_import + result[match.end():]
                replaced = True
            else:
                result = result[:match.start()] + result[match.end():]

    if not replaced:
        last_import = -1
        for m in re.finditer(r"^(?:import |from \S+ import )", result, re.MULTILINE):
            end = result.index("\n", m.start()) if "\n" in result[m.start():] else len(result)
            if end > last_import:
                last_import = end
        if last_import >= 0:
            return result[:last_import] + "\n" + bootstrap_import + result[last_import:]
        return bootstrap_import + "\n" + result

    return result


def _find_handler_functions(source: str) -> list[re.Match[str]]:
    """Find route handler decorated function definitions."""
    matches: list[re.Match[str]] = []
    for route_match in _ROUTE_HANDLER_RE.finditer(source):
        search_start = route_match.end()
        def_match = _DEF_RE.search(source, search_start)
        if def_match and def_match.start() - search_start < 200:
            matches.append(def_match)
    return matches


def _wrap_handler_with_session(
    source: str,
    handler_match: re.Match[str],
    agent_id: str,
    description: str | None = None,
) -> str:
    """Wrap a handler function body with agent.session context (advanced tier)."""
    indent = handler_match.group(1)
    is_async = bool(handler_match.group(2))
    func_name = handler_match.group(3)
    body_start = handler_match.end()

    next_newline = source.find("\n", body_start)
    if next_newline == -1:
        return source

    body_indent = indent + "    "
    session_indent = body_indent + "    "

    remaining = source[next_newline + 1:]
    body_lines: list[str] = []
    for line in remaining.split("\n"):
        if line.strip() == "" or line.startswith(body_indent) or line.startswith(indent + "\t\t"):
            body_lines.append(line)
        else:
            break

    body_end = next_newline + 1 + sum(len(line) + 1 for line in body_lines)
    original_body = source[next_newline + 1:body_end]

    re_indented_body = ""
    for line in original_body.split("\n"):
        if line.strip():
            if line.startswith(body_indent):
                re_indented_body += session_indent + line[len(body_indent):] + "\n"
            else:
                re_indented_body += "    " + line + "\n"
        else:
            re_indented_body += "\n"

    ctx_keyword = "async with" if is_async else "with"

    track_line = ""
    if _USER_INPUT_PATTERNS.search(original_body):
        track_line = f'{session_indent}s.track_user_message("user input")  # TODO: extract actual user message\n'

    agent_args = f'"{agent_id}"'
    if description:
        safe_desc = description.replace('"', '\\"')
        agent_args += f', description="{safe_desc}"'

    session_block = (
        f"{body_indent}# TODO(required): Replace with the real user ID from your auth/request context.\n"
        f"{body_indent}# Without a real user_id, per-user funnels, retention, and cohorts won't work.\n"
        f'{body_indent}_user_id = "anonymous"  # e.g. request.user.id, token.sub, session.user_id\n'
        f'{body_indent}agent = ai.agent({agent_args})\n'
        f"{body_indent}{ctx_keyword} agent.session("
        f"user_id=_user_id"
        f") as s:\n"
        f"{track_line}"
        f"{re_indented_body}"
        f"{body_indent}ai.flush()\n"
    )

    return source[:next_newline + 1] + session_block + source[body_end:]


def instrument_file(
    source: str,
    file_path: str,
    tier: str,
    bootstrap_import: str,
    agent_id: str,
    description: str | None = None,
    providers: list[str] | None = None,
) -> str:
    """Apply regex-based instrumentation transforms to Python source code.

    Tiers:
        quick_start -- returns source unchanged.
        standard -- adds bootstrap import; notes LLM call sites.
        advanced -- standard + session context wrapping for route handlers + ai.flush().
    """
    if tier == "quick_start":
        return source

    result = _add_bootstrap_import(source, bootstrap_import)

    if tier == "advanced":
        handler_matches = _find_handler_functions(result)
        offset = 0
        for handler_match in handler_matches:
            adjusted_start = handler_match.start() + offset
            adjusted_source = result
            shifted_match = _DEF_RE.search(adjusted_source, adjusted_start)
            if shifted_match is None:
                continue
            new_source = _wrap_handler_with_session(result, shifted_match, agent_id, description)
            offset += len(new_source) - len(result)
            result = new_source

    return result
