"""
Tool: summarise_with_ai
Sends the PROCEDURE DIVISION of a COBOL program to an LLM and gets back
a plain-English business logic summary.

Supported providers:
  - anthropic  → Claude (claude-sonnet-4-20250514 default)
  - openai     → GPT-4o (gpt-4o default)
  - groq       → Llama3 (llama3-70b-8192 default)
  - ollama     → local models (llama3 default, no API key needed)
  - custom     → any OpenAI-compatible endpoint (pass base_url)
"""

import os
import logging
from pathlib import Path

logger = logging.getLogger(__name__)

_MAX_PROC_LINES = 400

_DEFAULT_MODELS = {
    "anthropic": "claude-sonnet-4-20250514",
    "openai":    "gpt-4o",
    "groq":      "llama3-70b-8192",
    "ollama":    "llama3",
    "custom":    "gpt-4o",
}

_BASE_URLS = {
    "openai": "https://api.openai.com/v1",
    "groq":   "https://api.groq.com/openai/v1",
    "ollama": "http://localhost:11434/v1",
}


def _safe_read(path: Path) -> str:
    for enc in ("utf-8", "cp1252", "latin-1"):
        try:
            return path.read_text(encoding=enc)
        except UnicodeDecodeError:
            continue
    return path.read_text(encoding="utf-8", errors="replace")


def _extract_procedure_division(content: str) -> str:
    upper = content.upper()
    start = upper.find("PROCEDURE DIVISION")
    if start == -1:
        return content[:3000]
    return content[start:]


def _truncate(text: str, max_lines: int) -> str:
    lines = text.splitlines()
    if len(lines) <= max_lines:
        return text
    half = max_lines // 2
    return (
        "\n".join(lines[:half])
        + f"\n\n... [{len(lines) - max_lines} lines omitted] ...\n\n"
        + "\n".join(lines[-half:])
    )


def _build_prompt(program_id: str, parsed: dict, proc_div: str) -> str:
    summary        = parsed.get("business_logic_summary", {})
    tables_read    = summary.get("db2_tables_read", [])
    tables_written = summary.get("db2_tables_written", [])
    screens        = summary.get("screens_used", [])
    programs       = summary.get("programs_called", [])

    return f"""You are a COBOL mainframe modernization expert. Analyse this COBOL program and provide a concise business logic summary.

Program: {program_id}
DB2 tables read: {', '.join(tables_read) or 'none'}
DB2 tables written: {', '.join(tables_written) or 'none'}
BMS screens used: {', '.join(screens) or 'none'}
Programs called: {', '.join(programs) or 'none'}

PROCEDURE DIVISION:
{proc_div}

Respond with a JSON object containing exactly these fields:
{{
  "purpose": "One sentence describing what this program does in business terms",
  "business_domain": "e.g. Business Entity Registration, Payroll Processing, etc.",
  "user_facing": true or false,
  "key_operations": ["list", "of", "3-5", "key", "things", "this", "program", "does"],
  "modernization_notes": "One sentence about anything a developer should know when modernizing this to a REST API"
}}

Respond ONLY with the JSON object. No markdown, no explanation."""


async def _call_anthropic(prompt: str, api_key: str, model: str) -> str:
    import anthropic
    client  = anthropic.Anthropic(api_key=api_key)
    message = client.messages.create(
        model=model,
        max_tokens=500,
        messages=[{"role": "user", "content": prompt}],
    )
    return message.content[0].text.strip()


async def _call_openai_compatible(prompt: str, api_key: str, model: str, base_url: str) -> str:
    try:
        from openai import OpenAI
    except ImportError:
        raise ImportError("openai package not installed — run: pip install openai")

    client   = OpenAI(api_key=api_key or "ollama", base_url=base_url)
    response = client.chat.completions.create(
        model=model,
        max_tokens=500,
        messages=[{"role": "user", "content": prompt}],
    )
    return response.choices[0].message.content.strip()


async def summarise_with_ai(
    file_path: str,
    parsed: dict,
    api_key: str | None = None,
    provider: str = "anthropic",
    model: str | None = None,
    base_url: str | None = None,
) -> str:
    """
    Call an LLM to summarise the business logic of a COBOL program.

    Args:
        file_path:  Path to the COBOL source file.
        parsed:     The dict returned by parse_cobol_file().
        api_key:    API key. Falls back to ANTHROPIC_API_KEY / OPENAI_API_KEY / GROQ_API_KEY.
                    Not needed for ollama (local).
        provider:   anthropic | openai | groq | ollama | custom. Default: anthropic.
        model:      Model override. Defaults: claude-sonnet-4-20250514 / gpt-4o / llama3-70b-8192 / llama3.
        base_url:   Custom base URL for 'custom' provider or self-hosted models.

    Returns:
        JSON string with ai_summary fields, or an error string.
    """
    path = Path(file_path)
    if not path.exists():
        return f"File not found: {file_path}"

    content        = _safe_read(path)
    proc_div       = _truncate(_extract_procedure_division(content), _MAX_PROC_LINES)
    program_id     = parsed.get("program", path.stem.upper())
    prompt         = _build_prompt(program_id, parsed, proc_div)
    resolved_model = model or _DEFAULT_MODELS.get(provider, "gpt-4o")

    try:
        if provider == "anthropic":
            key = api_key or os.environ.get("ANTHROPIC_API_KEY")
            if not key:
                return "ANTHROPIC_API_KEY not set"
            try:
                import anthropic  # noqa: F401
            except ImportError:
                return "anthropic package not installed — run: pip install anthropic"
            return await _call_anthropic(prompt, key, resolved_model)

        elif provider in ("openai", "groq", "ollama", "custom"):
            env_map = {"openai": "OPENAI_API_KEY", "groq": "GROQ_API_KEY", "ollama": None, "custom": "OPENAI_API_KEY"}
            env_var = env_map.get(provider)
            resolved_key = api_key or (os.environ.get(env_var, "ollama") if env_var else "ollama")
            resolved_url = base_url or _BASE_URLS.get(provider, "https://api.openai.com/v1")
            return await _call_openai_compatible(prompt, resolved_key, resolved_model, resolved_url)

        else:
            return f"Unknown provider '{provider}'. Choose from: anthropic, openai, groq, ollama, custom"

    except Exception as e:
        logger.error("AI summary failed for %s (%s): %s", program_id, provider, e)
        return f"AI summary unavailable ({provider}): {e}"
