"""Simplified LLM routing using LangChain ReAct agents.

This is a drop-in replacement for `llm.py` — it exposes the same public surface
so `app_simple.py` (a copy of `app.py`) only needs to change the import line.

Why a simpler module?
- `llm.py` used litellm with OpenAI-style native tool calling. Most Ollama
  models (e.g. `qwen2.5:32b-instruct`) do NOT emit OpenAI-format tool_calls,
  so they end up returning prose and the agent breaks silently.
- example_app.py proved that LangChain `create_react_agent` with the ReAct
  text protocol works reliably for plain Ollama models.

This module uses one code path — ReAct via LangChain — for all providers
(Ollama, Anthropic, OpenAI).
"""

import json
import os
import queue
import re
import threading
import time as _time
import urllib.request
from typing import Callable, Generator

# ── Stop flag ─────────────────────────────────────────────────────────────────
_stop_event = threading.Event()


def request_stop() -> None:
    print("[llm_simple] request_stop() called", flush=True)
    _stop_event.set()


def _clear_stop() -> None:
    _stop_event.clear()


def _is_stopped() -> bool:
    return _stop_event.is_set()


# ── Constants ─────────────────────────────────────────────────────────────────
ANTHROPIC_MODELS = [
    "claude-opus-4-7",
    "claude-sonnet-4-6",
    "claude-haiku-4-5-20251001",
    "claude-3-5-sonnet-20241022",
    "claude-3-5-haiku-20241022",
]

OPENAI_MODELS = [
    "gpt-4o",
    "gpt-4o-mini",
    "gpt-4-turbo",
]

MAX_ITERATIONS = int(os.environ.get("AGENT_MAX_ITERATIONS", "5"))
MAX_RESULT_CHARS = 20_000


CHAT_INSTRUCTIONS = """\
You are a SQL query assistant for 3DCityDB v5 — a PostgreSQL-based 3D city model database \
following the CityGML standard.

Tool-use discipline:
- You have direct database access via the `run_query` tool. ALWAYS call it — never write SQL \
in your final answer expecting the user to run it.
- Always return the objectid of each matching feature in your final answer.

CRITICAL — No hallucination:
NEVER invent, fabricate, or guess building IDs, addresses, roof areas, or any other values.
Your Final Answer MUST be based ONLY on data returned by the run_query tool.
If all queries returned 0 rows or errors, you MUST tell the user "No results found" — do NOT
make up data. A wrong answer is far worse than an honest "no results found."

CRITICAL — Language rule:
ALWAYS respond in the EXACT language the user wrote in. If the user wrote in English, your entire
Final Answer MUST be in English — even if the database content (street names, values) is in German.
The language of database results NEVER affects the language of your reply.
Examples: user writes English → answer in English. User writes German → answer in German.
"""


# ── Provider discovery (copied from llm.py — small, stable helpers) ────────────

def get_ollama_models() -> list[str]:
    base = os.environ.get("OLLAMA_BASE_URL", "http://localhost:11434").rstrip("/")
    try:
        req = urllib.request.Request(f"{base}/api/tags", headers={"Accept": "application/json"})
        with urllib.request.urlopen(req, timeout=5) as resp:
            data = json.loads(resp.read())
        return [m["name"] for m in data.get("models", [])]
    except Exception:
        return []


def get_llamacpp_models() -> list[str]:
    """llama.cpp support is intentionally not implemented in the simple module.

    Returns an empty list so the UI dropdown is empty when llama.cpp is selected.
    """
    return []


def _ollama_reachable() -> bool:
    base = os.environ.get("OLLAMA_BASE_URL", "").rstrip("/")
    if not base:
        return False
    try:
        with urllib.request.urlopen(
            urllib.request.Request(f"{base}/api/tags"), timeout=3
        ) as resp:
            return resp.status == 200
    except Exception:
        return False


def _llamacpp_reachable() -> bool:
    return False


def detect_default_provider() -> str | None:
    if os.environ.get("ANTHROPIC_API_KEY"):
        return "anthropic"
    if os.environ.get("OPENAI_API_KEY"):
        return "openai"
    if _ollama_reachable():
        return "ollama"
    return None


def models_for_provider(provider: str) -> list[str]:
    if provider == "anthropic":
        return ANTHROPIC_MODELS
    if provider == "openai":
        return OPENAI_MODELS
    if provider == "ollama":
        return get_ollama_models()
    if provider == "llamacpp":
        return []
    return []


def default_model(provider: str) -> str:
    models = models_for_provider(provider)
    return models[0] if models else ""


# ── Context-usage helpers ──────────────────────────────────────────────────────

_ollama_ctx_cache: dict[str, int] = {}


def _estimate_tokens(messages: list[dict]) -> int:
    return sum(len(str(m.get("content") or "")) for m in messages) // 3


def _get_ollama_model_ctx(model: str) -> int:
    if model in _ollama_ctx_cache:
        return _ollama_ctx_cache[model]
    base = os.environ.get("OLLAMA_BASE_URL", "http://localhost:11434").rstrip("/")
    try:
        req = urllib.request.Request(
            f"{base}/api/show",
            data=json.dumps({"model": model}).encode(),
            headers={"Content-Type": "application/json"},
            method="POST",
        )
        with urllib.request.urlopen(req, timeout=5) as resp:
            data = json.loads(resp.read())
        info = data.get("model_info", {})
        ctx = info.get("llama.context_length") or info.get("context_length")
        if ctx:
            _ollama_ctx_cache[model] = int(ctx)
            return int(ctx)
    except Exception:
        pass
    configured = int(os.environ.get("OLLAMA_NUM_CTX", "32768"))
    _ollama_ctx_cache[model] = configured
    return configured


def _get_provider_ctx_limit(provider: str, model: str) -> int:
    if provider == "ollama":
        return _get_ollama_model_ctx(model)
    if provider == "llamacpp":
        return int(os.environ.get("LLAMACPP_CTX_SIZE", "32768"))
    return 200_000


def _log_context_usage(provider: str, model: str, messages: list[dict]) -> tuple[int, int, bool]:
    total_tok = _estimate_tokens(messages)
    limit = _get_provider_ctx_limit(provider, model)
    pct = (total_tok / limit * 100) if limit > 0 else 0
    over = pct > 80
    print(f"[ctx] total={total_tok:,} / {limit:,} ({pct:.0f}%)", flush=True)
    return total_tok, limit, over


# ── Markdown helpers ───────────────────────────────────────────────────────────

def _rows_to_markdown_table(rows: list[dict], max_rows: int = 5) -> str:
    if not rows:
        return "*No rows returned.*"
    headers = list(rows[0].keys())
    header_line = "| " + " | ".join(str(h) for h in headers) + " |"
    sep_line = "| " + " | ".join("---" for _ in headers) + " |"
    body_lines = [
        "| " + " | ".join(str(row.get(h, "")) for h in headers) + " |"
        for row in rows[:max_rows]
    ]
    result = "\n".join([header_line, sep_line] + body_lines)
    if len(rows) > max_rows:
        result += f"\n*… and {len(rows) - max_rows} more rows.*"
    return result


def _truncate_result(raw: str) -> str:
    if len(raw) <= MAX_RESULT_CHARS:
        return raw
    return raw[:MAX_RESULT_CHARS] + f"\n[…truncated — result was {len(raw):,} chars total]"


def _parse_tool_result(raw: str) -> dict:
    try:
        data = json.loads(raw)
        rows = data.get("results", [])
        return {
            "row_count": data.get("row_count", len(rows)),
            "execution_time_ms": data.get("execution_time_ms", 0),
            "preview_rows": rows[:5],
            "all_rows": rows[:200],
            "error": data.get("error") or None,
        }
    except Exception:
        return {
            "row_count": 0,
            "execution_time_ms": 0,
            "preview_rows": [],
            "all_rows": [],
            "error": None,
        }


# ── Viz: highlight payload extraction (copied from llm.py) ────────────────────

_VIZ_ID_KEYS = ["objectid", "gmlid", "GMLID", "gml_id", "building_id", "feature_id"]
_GMLID_REGEX = re.compile(r'\b([A-Z]{2,5}_[A-Z0-9_-]{5,})\b')


def extract_highlight_payload(events: list) -> dict:
    last_result = None
    for ev_type, ev_data in reversed(events):
        if ev_type == "tool_result" and not ev_data.get("error"):
            last_result = ev_data
            break
    if not last_result or not last_result.get("all_rows"):
        return {"buildings": [], "centroid": None}
    rows = last_result["all_rows"]
    found_key = next((k for k in _VIZ_ID_KEYS if k in rows[0]), None)
    if not found_key:
        return {"buildings": [], "centroid": None}
    gmlids = [str(row[found_key]) for row in rows if row.get(found_key)]
    gmlids = list(dict.fromkeys(gmlids))[:200]
    if not gmlids:
        return {"buildings": [], "centroid": None}
    centroid = _compute_centroid_wgs84(gmlids)
    return {
        "buildings": [{"gmlid": gid} for gid in gmlids],
        "centroid": [{"lat": centroid[0], "long": centroid[1]}] if centroid else None,
    }


def fallback_regex_extract(final_answer: str) -> list[str]:
    return list(dict.fromkeys(_GMLID_REGEX.findall(final_answer)))[:200]


def _compute_centroid_wgs84(gmlids: list[str]) -> tuple[float, float] | None:
    if not gmlids:
        return None
    try:
        import psycopg2
        from psycopg2.extras import RealDictCursor

        conn_params = {
            "host": os.environ.get("CITYDB_HOST", "localhost"),
            "port": int(os.environ.get("CITYDB_PORT", "5432")),
            "dbname": os.environ.get("CITYDB_NAME", "citydb"),
            "user": os.environ.get("CITYDB_USER", "citydb"),
            "password": os.environ.get("CITYDB_PASSWORD", "citydb"),
        }
        schema = os.environ.get("CITYDB_SCHEMA", "citydb")
        placeholders = ",".join(["%s"] * len(gmlids))
        sql = f"""
            SELECT
              ST_Y(ST_Centroid(ST_Transform(ST_Union(envelope), 4326))) AS lat,
              ST_X(ST_Centroid(ST_Transform(ST_Union(envelope), 4326))) AS lng
            FROM {schema}.feature
            WHERE objectid IN ({placeholders})
              AND envelope IS NOT NULL;
        """
        with psycopg2.connect(**conn_params) as conn:
            with conn.cursor(cursor_factory=RealDictCursor) as cur:
                cur.execute(sql, gmlids)
                row = cur.fetchone()
                if row and row["lat"] is not None:
                    return (float(row["lat"]), float(row["lng"]))
    except Exception as exc:
        print(f"[viz] centroid computation failed: {exc}", flush=True)
    return None


# ── SQL cleaning (Ollama models often wrap SQL in JSON or markdown) ────────────

def _clean_sql_input(raw: str) -> str:
    """Strip JSON wrappers, code fences, and stray braces from a model's tool input."""
    s = str(raw).strip()
    s = re.sub(r'^\s*json\s*', '', s, flags=re.IGNORECASE)
    s = s.replace("```json", "").replace("```sql", "").replace("```", "").strip()
    if s.startswith("{"):
        try:
            data = json.loads(s)
            s = data.get("sql") or data.get("query") or data.get("input_data") or s
        except Exception:
            pass
    s = s.strip().strip("{}").strip()
    return s


# ── ReAct prompt template ──────────────────────────────────────────────────────

_REACT_TEMPLATE = """{system_prompt}

You have access to the following tools:
{tools}

Available tool names: {tool_names}

Use the following format EXACTLY:
Thought: what you are thinking about the next step
Action: the tool name (one of [{tool_names}])
Action Input: the raw SQL query string (NO JSON wrapper, NO code fences, NO quotes around it)
Observation: the result of the action
... (repeat Thought/Action/Action Input/Observation as needed)
Thought: I now know the final answer
Final Answer: the answer to the user, in the user's language. ONLY use data from Observations above — NEVER invent values. If all Observations returned no rows, say "No results found."

Begin!

Previous conversation:
{chat_history}

Question: {input}
Thought: {agent_scratchpad}"""


# ── Build LLM by provider ──────────────────────────────────────────────────────

def _build_llm(provider: str, model: str, temperature: float):
    if provider == "ollama":
        from langchain_ollama import ChatOllama
        return ChatOllama(
            base_url=os.environ.get("OLLAMA_BASE_URL", "http://localhost:11434"),
            model=model,
            temperature=temperature,
            num_ctx=int(os.environ.get("OLLAMA_NUM_CTX", "32768")),
        )
    if provider == "anthropic":
        from langchain_anthropic import ChatAnthropic
        return ChatAnthropic(
            model=model,
            temperature=temperature,
            max_tokens=4096,
        )
    if provider == "openai":
        from langchain_openai import ChatOpenAI
        return ChatOpenAI(
            model=model,
            temperature=temperature,
        )
    raise ValueError(f"Unsupported provider: {provider}")


def _format_chat_history(messages: list[dict]) -> str:
    """Render prior user/assistant turns as plain text for the ReAct prompt."""
    lines = []
    for m in messages:
        role = m["role"]
        if role == "user":
            lines.append(f"Human: {m['content']}")
        elif role == "assistant":
            lines.append(f"AI: {m['content']}")
    return "\n".join(lines) if lines else "(none)"


# ── Main agentic loop ──────────────────────────────────────────────────────────

def agentic_stream(
    provider: str,
    model: str,
    temperature: float,
    messages: list[dict],
    tool_executor: Callable[[str], str],
    *,
    enable_thinking: bool = True,
) -> Generator[tuple, None, None]:
    """ReAct agent that yields events compatible with app.py's chat_stream.

    Events:
        ("status", str)
        ("ping", "")
        ("context_update", {"used": int, "limit": int})
        ("thinking", str)
        ("tool_call", {"tool": str, "args": dict, "iteration": int})
        ("tool_result", {"row_count": int, "execution_time_ms": int,
                         "preview_rows": list, "all_rows": list, "error": str|None, "iteration": int})
        ("final", str)
        ("stopped", "")
    """
    from langchain.agents import AgentExecutor, create_react_agent
    from langchain.tools import Tool
    from langchain_core.prompts import PromptTemplate

    # Separate system prompt, history, and the current user question.
    system_content = ""
    history_msgs: list[dict] = []
    user_input = ""
    for m in messages:
        if m["role"] == "system":
            system_content = m["content"]
        else:
            history_msgs.append(m)
    if history_msgs and history_msgs[-1]["role"] == "user":
        user_input = history_msgs[-1]["content"]
        history_msgs = history_msgs[:-1]

    used_tok = _estimate_tokens(messages)
    ctx_limit = _get_provider_ctx_limit(provider, model)
    yield ("context_update", {"used": used_tok, "limit": ctx_limit})
    yield ("status", "Thinking…" if enable_thinking else "Formulating…")

    # Build the run_query tool — wrap tool_executor with input cleaning.
    def _run_sql(sql_input: str) -> str:
        cleaned = _clean_sql_input(sql_input)
        print(f"[llm_simple] run_query → {cleaned[:120]}", flush=True)
        return tool_executor(cleaned)

    tools = [
        Tool(
            name="run_query",
            func=_run_sql,
            description=(
                "Execute a read-only SQL SELECT query against the 3DCityDB database "
                "and return the rows as JSON. Input MUST be a raw SQL string."
            ),
        )
    ]

    prompt = PromptTemplate.from_template(_REACT_TEMPLATE).partial(
        system_prompt=system_content
    )

    llm = _build_llm(provider, model, temperature)
    agent = create_react_agent(llm, tools, prompt)
    executor = AgentExecutor(
        agent=agent,
        tools=tools,
        verbose=False,
        handle_parsing_errors=True,
        max_iterations=MAX_ITERATIONS,
        return_intermediate_steps=True,
    )

    chat_history_str = _format_chat_history(history_msgs)

    # Run executor.stream(...) in a worker thread so the Gradio generator can
    # still poll the stop flag and emit pings while the LLM is generating.
    q: queue.Queue = queue.Queue()
    SENTINEL = object()

    def _worker():
        try:
            for chunk in executor.stream(
                {"input": user_input, "chat_history": chat_history_str}
            ):
                q.put(("chunk", chunk))
        except Exception as exc:
            q.put(("error", exc))
        finally:
            q.put((SENTINEL, None))

    t = threading.Thread(target=_worker, daemon=True)
    t.start()

    iteration = 0
    seen_actions: set[int] = set()
    seen_steps: set[int] = set()
    final_output = ""

    while True:
        if _is_stopped():
            print("[llm_simple] stop flag set — exiting agentic_stream", flush=True)
            yield ("stopped", "")
            return
        try:
            kind, payload = q.get(timeout=0.25)
        except queue.Empty:
            yield ("ping", "")
            continue

        if kind is SENTINEL:
            break

        if kind == "error":
            exc = payload
            print(f"[llm_simple] executor error: {type(exc).__name__}: {exc}", flush=True)
            raise exc

        # kind == "chunk"
        chunk = payload

        # Actions emitted by this step.
        actions = chunk.get("actions") or []
        for action in actions:
            key = id(action)
            if key in seen_actions:
                continue
            seen_actions.add(key)
            log_text = getattr(action, "log", "") or ""
            thought = _extract_thought_from_log(log_text)
            if thought:
                yield ("thinking", thought)
            tool_input = action.tool_input
            if isinstance(tool_input, dict):
                sql = tool_input.get("sql") or tool_input.get("query") or json.dumps(tool_input)
            else:
                sql = str(tool_input)
            sql = _clean_sql_input(sql)
            yield (
                "tool_call",
                {"tool": "run_query", "args": {"sql": sql}, "iteration": iteration},
            )

        # Steps (action + observation pairs).
        steps = chunk.get("steps") or []
        for step in steps:
            key = id(step)
            if key in seen_steps:
                continue
            seen_steps.add(key)
            observation = getattr(step, "observation", "")
            raw = str(observation) if observation is not None else ""
            result_data = _parse_tool_result(raw)
            yield ("tool_result", {**result_data, "iteration": iteration})
            iteration += 1

        # Final output.
        if "output" in chunk:
            final_output = str(chunk["output"])

    # Drain finished — emit the final answer.
    if final_output:
        # Strip any stray "Final Answer:" prefix the parser might have left.
        cleaned = re.sub(r"^\s*Final Answer:\s*", "", final_output).strip()
        yield ("final", cleaned)
    else:
        yield ("final", "_The model did not produce a final answer. Try rephrasing._")


# ── Aliases & stubs for app.py compatibility ──────────────────────────────────

# app.py falls back to prompt_agentic_stream when agentic_stream produces no
# content. With this module the agentic_stream IS already prompt-based, so the
# fallback would be redundant — alias it to the same function.
prompt_agentic_stream = agentic_stream


def stream_response(
    provider: str,
    model: str,
    temperature: float,
    system_prompt: str,
    messages: list[dict],
) -> Generator[str, None, None]:
    """Non-agentic streaming. Kept as a stub for app.py imports — not used."""
    yield ""


# ── ReAct log parsing ──────────────────────────────────────────────────────────

_THOUGHT_RE = re.compile(r"Thought:\s*(.+?)(?=\n(?:Action|Final Answer)|\Z)", re.DOTALL)


def _extract_thought_from_log(log_text: str) -> str:
    """Pull the 'Thought: ...' segment out of LangChain's action log."""
    m = _THOUGHT_RE.search(log_text)
    if m:
        return m.group(1).strip()
    return ""
