#!/usr/bin/env python3
# -*- coding: utf-8 -*-
from __future__ import annotations

import json
import os
import re
import time
import urllib.request
from collections import Counter
from http.server import BaseHTTPRequestHandler, ThreadingHTTPServer
from pathlib import Path
from typing import Any

ROOT = Path("/app")
OLLAMA_URL = os.environ.get("OLLAMA_URL", "http://ollama:11434").rstrip("/")
OLLAMA_MODEL = os.environ.get("OLLAMA_MODEL", "qwen3-8b-korean-coding-router:latest")
ROUTER_MODEL_ID = os.environ.get("ROUTER_MODEL_ID", "qwen3-8b-korean-coding-router")

GLOBAL = """당신은 한국어 개발/운영 보조 AI입니다.
최종 답변만 출력합니다. 내부 사고 과정과 메타 설명은 출력하지 않습니다.
반드시 한국어 존댓말로 답합니다.
Qwen3 사고 태그(<think>)와 중국어 한자 출력은 금지합니다.
위험한 작업은 확인 명령과 실행 명령을 분리합니다.
긴 명령은 .sh/.py/zip 방식으로 안내합니다.
필수 문자열은 동의어로 바꾸지 말고 그대로 포함합니다.
"""

LONGFORM_TEMPLATE = """주제: longform writing / 장문 생성 안정화
- 반복 루프 금지
- 같은 문장과 같은 구절 반복 금지
- 소제목 사용
- 코딩 주제라면 개념, 실무 흐름, 예시, 주의점, 학습 방향으로 나눕니다.
"""

ROUTE_KEYWORDS = {
    "linux": ["포트", "ss -lntp", "pid", "8010", "8000", "로그", "프로세스", "리스닝", "listen"],
    "docker": ["docker", "도커", "compose", "컨테이너", "open webui", "open-webui", "webui", "오픈웹", "ollama", "올라마", "router-api", "3001", "11435"],
    "style": ["긴 명령", "붙여넣기", "짧은 명령", "한 단계", "zip", "스크립트"],
    "lora": ["lora", "adapter", "어댑터"],
    "security": ["토큰", "api key", "hf_token", "보안"],
    "eval": ["raw", "cleaned", "평가", "benchmark"],
    "dataset": ["dataset", "데이터셋", "clean", "rejected", "bad_json"],
    "systemd": ["systemd", "systemctl", "service", "disable"],
    "vllm": ["vllm", "/v1/models", "gpu_memory_utilization"],
    "ops": ["자동 재시작", "다시 살아", "crontab", "nohup", "tmux"],
    "python": ["python", "파이썬", "json", "glob", "path"],
    "workflow": ["재개", "상태", "다음 단계", "stable"],
}

RULES = {
    "linux": {"must": [], "any": ["ss -lntp", "ps", "PID", "포트"]},
    "style": {"must": ["스크립트", "한 단계"], "any": ["zip", ".sh", ".py", "짧은 명령"]},
    "docker": {"must": [], "any": ["docker compose ps", "http://ollama:11434", "http://router-api:8011/v1"]},
    "lora": {"must": ["adapter", "benchmark"], "any": ["dataset", "stable", "rejected"]},
    "security": {"must": ["토큰", "로그"], "any": ["환경변수", "Git", "HF_TOKEN"]},
    "eval": {"must": ["raw", "cleaned"], "any": ["think", "meta", "cleanup"]},
    "dataset": {"must": ["clean", "rejected"], "any": ["report", "duplicate", "bad_json"]},
    "systemd": {"must": ["systemctl"], "any": ["status", "disable", "mask"]},
    "ops": {"must": ["systemd", "Docker"], "any": ["crontab", "ps", "ss"]},
    "workflow": {"must": [], "any": ["상태", "다음 단계", "stable"]},
    "longform": {"must": [], "any": ["소제목", "요약"]},
}


def load_templates() -> dict[str, Any]:
    marker = ROOT / "policies/STABLE_ROUTER_TEMPLATE.txt"
    if marker.exists():
        template_path = Path(marker.read_text(encoding="utf-8").strip())
        if not template_path.is_absolute():
            template_path = ROOT / template_path
    else:
        candidates = sorted((ROOT / "policies").glob("*templates*.json"))
        if not candidates:
            return {"templates": {}}
        template_path = candidates[-1]

    try:
        obj = json.loads(template_path.read_text(encoding="utf-8"))
    except Exception:
        obj = {"templates": {}}

    obj = dict(obj)
    obj["templates"] = dict(obj.get("templates", {}))
    obj["templates"]["longform"] = {"name": "longform_docker_router", "template": LONGFORM_TEMPLATE}
    return obj


TEMPLATES = load_templates()


def contains_any(text: str, terms: list[str]) -> bool:
    low = text.lower()
    return any(term.lower() in low for term in terms)


def is_stack_question(user: str) -> bool:
    low = user.lower()
    has_webui = ("open webui" in low) or ("open-webui" in low) or ("webui" in low) or ("오픈웹" in low)
    has_ollama_or_docker = (
        ("ollama" in low) or ("올라마" in low) or ("docker" in low) or ("도커" in low)
        or ("compose" in low) or ("모델" in low) or ("3001" in low) or ("router-api" in low)
    )
    return has_webui and has_ollama_or_docker


def is_port_or_process_question_text(user: str) -> bool:
    low = user.lower()
    strong = ["ss -lntp", "pid", "8010", "8000", "포트", "프로세스", "리스닝", "listen"]
    return any(item in low for item in strong)


def is_short_answer_request(user: str) -> bool:
    low = user.lower()
    return any(item in low for item in ["짧게", "간단히", "요약", "한 단계", "확인 순서", "명령어"])


def is_longform_context_only(user: str) -> bool:
    low = user.lower()
    context_patterns = [
        "방금 장문",
        "장문 이후",
        "장문 다음",
        "긴 글 이후",
        "4천자 이후",
        "4000자 이후",
        "이후에도 정상",
    ]
    return any(p in low for p in context_patterns)


def is_explicit_longform_request(user: str) -> bool:
    low = user.lower()
    if is_longform_context_only(user) or is_short_answer_request(user):
        return False

    length_signal = any(k in low for k in ["4천자", "4000", "3천자", "3000", "긴 글", "장문", "longform"])
    write_signal = any(k in low for k in ["써줘", "작성", "생성", "텍스트", "에세이", "글을", "설명해줘"])
    return length_signal and write_signal


def route_reason(user: str) -> dict[str, Any]:
    low = user.lower()
    signals: list[str] = []
    blocked: list[str] = []

    port_signal = is_port_or_process_question_text(user)
    stack_signal = is_stack_question(user)
    short_signal = is_short_answer_request(user)
    long_context_only = is_longform_context_only(user)
    explicit_longform = is_explicit_longform_request(user)

    if port_signal:
        signals.append("port_or_process_signal")
    if stack_signal:
        signals.append("docker_stack_signal")
    if short_signal:
        signals.append("short_answer_signal")
        blocked.append("longform")
    if long_context_only:
        signals.append("longform_context_only")
        blocked.append("longform")
    if explicit_longform:
        signals.append("explicit_longform_request")

    # Priority 1: safety/operation command questions.
    # A port/PID question should not become longform even if it mentions "장문 이후".
    if port_signal:
        return {"route": "linux", "signals": signals, "blocked_routes": sorted(set(blocked)), "priority": "safety_linux_first"}

    # Priority 2: current Docker/Open WebUI/Ollama stack.
    if stack_signal:
        return {"route": "docker", "signals": signals, "blocked_routes": sorted(set(blocked)), "priority": "docker_stack_second"}

    # Priority 3: script/style preference.
    if contains_any(user, ROUTE_KEYWORDS["style"]):
        signals.append("style_script_signal")
        return {"route": "style", "signals": signals, "blocked_routes": sorted(set(blocked)), "priority": "style_third"}

    # Priority 4: more specific engineering domains.
    for cat in ["systemd", "vllm", "ops", "dataset", "lora", "security", "eval", "python"]:
        if contains_any(user, ROUTE_KEYWORDS[cat]):
            signals.append(f"{cat}_signal")
            return {"route": cat, "signals": signals, "blocked_routes": sorted(set(blocked)), "priority": "domain_specific"}

    # Priority 5: longform only when explicit and not context-only.
    if explicit_longform:
        return {"route": "longform", "signals": signals, "blocked_routes": sorted(set(blocked)), "priority": "explicit_longform_late"}

    # Fallback scoring for remaining weak signals.
    hits: list[tuple[int, str]] = []
    for cat, kws in ROUTE_KEYWORDS.items():
        count = sum(1 for k in kws if k.lower() in low)
        if count:
            hits.append((count, cat))
    if hits:
        hits.sort(reverse=True)
        return {"route": hits[0][1], "signals": signals + [f"weak_keyword_{hits[0][1]}"], "blocked_routes": sorted(set(blocked)), "priority": "weak_keyword"}

    return {"route": "workflow", "signals": signals, "blocked_routes": sorted(set(blocked)), "priority": "fallback"}


def route(text: str) -> str:
    return str(route_reason(text)["route"])


def strip_think(text: str) -> str:
    text = re.sub(r"<think>.*?</think>", "", text, flags=re.I | re.S)
    return re.sub(r"</?think>", "", text, flags=re.I).strip()


def leak_flags(text: str) -> list[str]:
    flags: list[str] = []
    if re.search(r"<\s*/?\s*think\s*>|chain\s+of\s+thought|internal\s+monologue", text, re.I):
        flags.append("think")
    if re.search(r"[\u4e00-\u9fff]", text):
        flags.append("cjk")
    return flags


def missing_terms(text: str, terms: list[str]) -> list[str]:
    low = text.lower()
    return [x for x in terms if x.lower() not in low]


def repetition_flags(text: str) -> list[str]:
    compact = re.sub(r"\s+", " ", text)
    if re.search(r"(.{3,40}?)(?:\1){3,}", compact):
        return ["phrase_repeat"]

    chars = re.sub(r"\s+", "", text)
    if len(chars) > 500:
        grams = [chars[i:i + 8] for i in range(max(0, len(chars) - 7))]
        if grams and Counter(grams).most_common(1)[0][1] >= 16:
            return ["ngram_repeat"]
    return []


def select_rules(cat: str, user: str) -> dict[str, list[str]]:
    rule = RULES.get(cat, {"must": [], "any": []})
    must = list(rule.get("must", []))
    any_terms = list(rule.get("any", []))

    if is_port_or_process_question_text(user):
        must = sorted(set(must + ["ss -lntp", "ps"]))
        any_terms = sorted(set(any_terms + ["PID", "포트"]))

    if cat == "style" and ("긴 명령" in user or "짧은 명령" in user or "붙여넣" in user):
        must = sorted(set(must + ["스크립트", "한 단계"]))
        any_terms = sorted(set(any_terms + ["zip", ".sh", ".py", "짧은 명령"]))

    return {"must": must, "any": any_terms}


def safety_kill_block() -> str:
    return """[안전 기준]

- 확인 전에는 kill하지 않습니다.
- 먼저 ss -lntp로 포트와 PID를 확인합니다.
- 그 다음 ps로 실제 프로세스와 실행 인자를 확인합니다.
- systemd 서비스면 systemctl status를 먼저 확인하고, 종료가 필요하면 systemctl stop을 우선 검토합니다.
- Docker 서비스면 docker compose ps와 docker compose stop을 우선 검토합니다.
- kill -9는 데이터 손상 위험이 있으므로 마지막 수단입니다.
"""


def current_stack_block() -> str:
    return """[현재 Docker stack 기준]

```text
Open WebUI = http://127.0.0.1:3001
Router API = http://127.0.0.1:8011/v1
Ollama API = http://127.0.0.1:11435
Open WebUI -> Ollama = http://ollama:11434
Open WebUI -> Router API = http://router-api:8011/v1
Router API -> Ollama = http://ollama:11434
```

```bash
cd /home/saul9523/dgx_ai_factory/101_qwen_next_coding_1m_train/dist/docker_openwebui_router/qwen3-router-openwebui-stack
docker compose ps
curl -fsS http://127.0.0.1:8011/v1/models
curl -fsS http://127.0.0.1:11435/api/tags
```

주의: `http://127.0.0.1:3001`은 Open WebUI 화면 주소입니다. Ollama API 주소로 넣으면 안 됩니다.
"""


def remove_unsafe_kill_lines(answer: str) -> str:
    safe_words = ("금지", "하지", "마지막", "위험", "권장하지", "확인 전", "수단")
    kept: list[str] = []
    for line in answer.splitlines():
        normalized = line.replace("`kill`", "kill").replace("`kill -9`", "kill -9")
        has_kill = re.search(r"\bkill\b", normalized, re.I) is not None
        if has_kill and not any(word in normalized for word in safe_words):
            continue
        kept.append(normalized)
    return "\n".join(kept).strip()


def remove_wrong_stack_url_lines(answer: str) -> str:
    wrong = [
        "host.docker.internal:11434",
        "172.17.0.1:11434",
        "http://127.0.0.1:3000",
        "http://127.0.0.1:3001/api",
    ]
    kept: list[str] = []
    for line in answer.splitlines():
        if any(item in line for item in wrong):
            continue
        kept.append(line)
    return "\n".join(kept).strip()


def apply_safety_patch(answer: str, cat: str, user: str) -> str:
    patched = answer.strip()
    port_question = is_port_or_process_question_text(user)
    stack_question = is_stack_question(user)

    # Category-independent safety guard:
    # Even if routing ever misclassifies a port/PID prompt, this guard still applies.
    if port_question:
        patched = remove_unsafe_kill_lines(patched)
        patched = patched.replace("`kill`", "kill").replace("`kill -9`", "kill -9")
        if "확인 전에는 kill하지 않습니다" not in patched:
            patched = safety_kill_block() + "\n\n" + patched

    if stack_question:
        patched = remove_wrong_stack_url_lines(patched)
        if "http://ollama:11434" not in patched or "http://router-api:8011/v1" not in patched or "http://127.0.0.1:3001" not in patched:
            patched = current_stack_block() + "\n\n" + patched

    if re.search(r"\bkill\s+-9\b", patched, re.I) and "마지막 수단" not in patched:
        patched = safety_kill_block() + "\n\n" + remove_unsafe_kill_lines(patched)

    return patched.strip()


def system_prompt(cat: str, user: str = "", repair: bool = False, missing_must: list[str] | None = None, missing_any: list[str] | None = None) -> str:
    template = TEMPLATES.get("templates", {}).get(cat) or TEMPLATES.get("templates", {}).get("workflow") or {"template": ""}
    text = GLOBAL + "\n\n선택된 주제 템플릿:\n" + template.get("template", "")

    rr = route_reason(user) if user else {"signals": [], "blocked_routes": []}
    text += "\n\n라우팅 판정 힌트:\n"
    text += "route=" + cat + "\n"
    text += "signals=" + ", ".join(rr.get("signals", [])) + "\n"
    text += "blocked_routes=" + ", ".join(rr.get("blocked_routes", [])) + "\n"

    if cat in {"linux", "ops", "systemd"} or is_port_or_process_question_text(user):
        text += "\n\n운영 안전 규칙:\n" + safety_kill_block()

    if cat == "docker" or is_stack_question(user):
        text += "\n\nDocker/Open WebUI stack 주소 규칙:\n" + current_stack_block()

    if repair:
        text += "\n\nPOST_CHECK_REWRITE=ON\n"
        text += "누락 must 문자열: " + (", ".join(missing_must or []) or "(없음)") + "\n"
        text += "누락 any 문자열: " + (", ".join(missing_any or []) or "(없음)") + "\n"
        text += "위 문자열을 그대로 포함하고 최종 답변만 출력합니다.\n"

    return text


def ollama_chat(messages: list[dict[str, str]], longform: bool = False, num_predict: int = 900) -> str:
    payload = {
        "model": OLLAMA_MODEL,
        "messages": messages,
        "stream": False,
        "options": {
            "temperature": 0.15 if longform else 0,
            "top_p": 0.9,
            "repeat_penalty": 1.18 if longform else 1.1,
            "num_predict": num_predict,
        },
    }
    req = urllib.request.Request(
        OLLAMA_URL + "/api/chat",
        data=json.dumps(payload, ensure_ascii=False).encode("utf-8"),
        headers={"Content-Type": "application/json"},
    )
    with urllib.request.urlopen(req, timeout=300) as resp:
        obj = json.loads(resp.read().decode("utf-8"))
    return strip_think(obj.get("message", {}).get("content", ""))


def gen_normal(cat: str, user: str, history: list[dict[str, str]]) -> str:
    rule = select_rules(cat, user)
    messages = [{"role": "system", "content": system_prompt(cat, user)}] + history[-6:] + [{"role": "user", "content": user}]
    first = ollama_chat(messages)

    mm = missing_terms(first, rule["must"])
    ma = missing_terms(first, rule["any"])
    leaks = leak_flags(first)
    reps = repetition_flags(first)
    needs_repair = bool(mm or leaks or reps or (cat in {"linux", "style", "docker", "lora"} and len(ma) >= 2))

    if not needs_repair:
        return apply_safety_patch(first, cat, user)

    repair_user = (
        f"원래 질문:\n{user}\n\n"
        f"기존 답변:\n{first}\n\n"
        f"누락 must 문자열: {', '.join(mm) if mm else '(없음)'}\n"
        f"누락 any 문자열: {', '.join(ma) if ma else '(없음)'}\n"
        "위 문자열을 그대로 포함하고 최종 답변만 출력해 주세요."
    )
    second = ollama_chat([
        {"role": "system", "content": system_prompt(cat, user, True, mm, ma)},
        {"role": "user", "content": repair_user},
    ])

    mm2 = missing_terms(second, rule["must"])
    ma2 = missing_terms(second, rule["any"])
    leaks2 = leak_flags(second)
    reps2 = repetition_flags(second)

    if len(mm2) < len(mm) or (not leaks2 and leaks) or (not reps2 and reps) or (len(mm2) == len(mm) and len(ma2) < len(ma)):
        return apply_safety_patch(second, cat, user)

    return apply_safety_patch(first, cat, user)


def target_chars(user: str) -> int:
    if "4천" in user or "4000" in user:
        return 4000
    if "3천" in user or "3000" in user:
        return 3000
    return 1800


def gen_long(user: str) -> str:
    target = target_chars(user)
    sections = 5 if target >= 3800 else 4
    titles = ["문제 정의와 코딩의 의미", "실무 개발 흐름", "디버깅과 테스트", "협업과 운영", "학습 로드맵과 요약"]
    parts: list[str] = []
    for idx in range(sections):
        title = titles[idx] if idx < len(titles) else "추가"
        prompt = (
            f"원래 요청:\n{user}\n\n"
            f"총 목표 분량: 최소 {target}자.\n"
            f"현재 section: {idx + 1}/{sections}\n"
            f"이번 section 제목: {title}\n"
            "이번 section만 작성하고 반복 루프를 만들지 마세요."
        )
        text = ollama_chat([
            {"role": "system", "content": system_prompt("longform", user)},
            {"role": "user", "content": prompt},
        ], True, 720)
        if repetition_flags(text):
            text = ollama_chat([
                {"role": "system", "content": system_prompt("longform", user)},
                {"role": "user", "content": prompt + "\n반복 없이 새 문장으로 다시 작성하세요."},
            ], True, 720)
        parts.append(text.strip())
    return "\n\n".join(parts)


def generate(messages: list[dict[str, str]]) -> str:
    user_messages = [m for m in messages if m.get("role") == "user"]
    user = str(user_messages[-1].get("content", "")) if user_messages else ""

    history: list[dict[str, str]] = []
    for m in messages[-8:-1]:
        if m.get("role") in {"user", "assistant"}:
            content = str(m.get("content", ""))
            # Long previous assistant answers are summarized as a placeholder to avoid context spillover.
            if m.get("role") == "assistant" and len(content) > 1200:
                content = "[이전 답변은 장문이므로 히스토리에서는 생략됨]"
            else:
                content = content[:1000]
            history.append({"role": str(m["role"]), "content": content})

    rr = route_reason(user)
    cat = str(rr["route"])

    if cat == "longform":
        answer = gen_long(user)
    else:
        answer = gen_normal(cat, user, history)

    return apply_safety_patch(answer, cat, user)


class Handler(BaseHTTPRequestHandler):
    server_version = "Qwen3RouterAPI/4.1"

    def send_json(self, obj: Any, status: int = 200) -> None:
        body = json.dumps(obj, ensure_ascii=False).encode("utf-8")
        self.send_response(status)
        self.send_header("Content-Type", "application/json; charset=utf-8")
        self.send_header("Content-Length", str(len(body)))
        self.end_headers()
        self.wfile.write(body)

    def do_GET(self) -> None:
        if self.path.rstrip("/") == "/v1/models":
            self.send_json({
                "object": "list",
                "data": [{
                    "id": ROUTER_MODEL_ID,
                    "object": "model",
                    "created": int(time.time()),
                    "owned_by": "local",
                }],
            })
        elif self.path == "/health":
            self.send_json({
                "status": "ok",
                "server": self.server_version,
                "ollama_url": OLLAMA_URL,
                "ollama_model": OLLAMA_MODEL,
                "routing": "priority_intent_router_v4_1",
            })
        else:
            self.send_json({"error": "not found"}, 404)

    def do_POST(self) -> None:
        if self.path.rstrip("/") == "/v1/route/debug":
            try:
                length = int(self.headers.get("Content-Length", "0"))
                req = json.loads(self.rfile.read(length).decode("utf-8"))
                text = str(req.get("text") or req.get("prompt") or "")
                self.send_json(route_reason(text))
            except Exception as exc:
                self.send_json({"error": f"{type(exc).__name__}: {exc}"}, 500)
            return

        if self.path.rstrip("/") != "/v1/chat/completions":
            self.send_json({"error": "not found"}, 404)
            return

        try:
            length = int(self.headers.get("Content-Length", "0"))
            req = json.loads(self.rfile.read(length).decode("utf-8"))
            answer = generate(req.get("messages", []))

            if req.get("stream"):
                self.send_response(200)
                self.send_header("Content-Type", "text/event-stream; charset=utf-8")
                self.send_header("Cache-Control", "no-cache")
                self.end_headers()
                chunk = {
                    "id": "chatcmpl-router",
                    "object": "chat.completion.chunk",
                    "created": int(time.time()),
                    "model": ROUTER_MODEL_ID,
                    "choices": [{"index": 0, "delta": {"role": "assistant", "content": answer}, "finish_reason": None}],
                }
                self.wfile.write(("data: " + json.dumps(chunk, ensure_ascii=False) + "\n\n").encode("utf-8"))
                done = {
                    "id": "chatcmpl-router",
                    "object": "chat.completion.chunk",
                    "created": int(time.time()),
                    "model": ROUTER_MODEL_ID,
                    "choices": [{"index": 0, "delta": {}, "finish_reason": "stop"}],
                }
                self.wfile.write(("data: " + json.dumps(done, ensure_ascii=False) + "\n\n").encode("utf-8"))
                self.wfile.write(b"data: [DONE]\n\n")
                return

            self.send_json({
                "id": "chatcmpl-router",
                "object": "chat.completion",
                "created": int(time.time()),
                "model": ROUTER_MODEL_ID,
                "choices": [{"index": 0, "message": {"role": "assistant", "content": answer}, "finish_reason": "stop"}],
                "usage": {"prompt_tokens": 0, "completion_tokens": 0, "total_tokens": 0},
            })
        except Exception as exc:
            self.send_json({"error": {"message": f"{type(exc).__name__}: {exc}", "type": "router_error"}}, 500)

    def log_message(self, fmt: str, *args: Any) -> None:
        print("%s - %s" % (self.address_string(), fmt % args), flush=True)


if __name__ == "__main__":
    print(
        f"router api serving on 0.0.0.0:8011, server=Qwen3RouterAPI/4.1, model={ROUTER_MODEL_ID}, ollama={OLLAMA_URL}/{OLLAMA_MODEL}",
        flush=True,
    )
    ThreadingHTTPServer(("0.0.0.0", 8011), Handler).serve_forever()