#!/usr/bin/env python3
"""NL2SQL 전용 워커 프로세스

Usage: python nl2sql_worker.py <port>

담당 도구:
  run_sql, query_pv_history, get_tag_metadata, list_drawings, query_with_nl

특징:
  - PostgreSQL 직접 연결
  - LLM SQL 생성 + DB 실행 분리
  - 메모리: ~1GB (SQL 생성용 LLM)
  - 생명주기: 메인 서버 종료 시까지 유지
"""

from __future__ import annotations
import sys
import os

# mcp-server 디렉토리를 Python 경로에 추가
sys.path.insert(0, os.path.dirname(os.path.dirname(os.path.abspath(__file__))))

import logging
import asyncio
from functools import lru_cache

from fastapi import FastAPI, Request
import uvicorn
import httpx

# ── 설정 ─────────────────────────────────────────────────────────────────────

DB_CONNECTION_STRING = "postgresql://postgres:postgres@localhost:5432/iiot_platform"
DB_TIMEOUT = 10

VLLM_BASE_URL = "http://localhost:8000/v1"
VLLM_MODEL = "Qwen/Qwen3-Coder-Next-FP8"

logging.basicConfig(
    level=logging.INFO,
    stream=sys.stderr,
    format="%(asctime)s [nl2sql_worker] %(levelname)s %(message)s",
)

app = FastAPI()

# ── DB 연결 풀 ───────────────────────────────────────────────────────────────

def _get_db_connection():
    import psycopg
    return psycopg.connect(DB_CONNECTION_STRING, connect_timeout=DB_TIMEOUT)

# ── LLM 클라이언트 ───────────────────────────────────────────────────────────

@lru_cache(maxsize=1)
def _llm_client():
    from openai import AsyncOpenAI
    return AsyncOpenAI(base_url=VLLM_BASE_URL, api_key="dummy")

async def _generate_sql(natural_language: str) -> str:
    """자연어를 SQL로 변환."""
    client = _llm_client()
    
    prompt = f"""다음 자연어 질문을 PostgreSQL SQL 쿼리로 변환하세요.
데이터베이스는 TimescaleDB 기반의 IIoT 플랫폼입니다.

질문:
{natural_language}

SQL 쿼리 (SELECT 문만, 설명 없이):"""
    
    response = await client.chat.completions.create(
        model=VLLM_MODEL,
        messages=[
            {"role": "system", "content": "You are a PostgreSQL SQL expert. Generate only valid SQL queries."},
            {"role": "user", "content": prompt},
        ],
        max_tokens=1024,
        temperature=0.1,
    )
    return response.choices[0].message.content.strip()

# ── NL2SQL 도구 구현 ─────────────────────────────────────────────────────────

@app.get("/health")
async def health():
    """워커 헬스체크."""
    return {"status": "ok"}

@app.post("/execute")
async def execute(request: Request):
    """HTTP 요청을 MCP 도구 호출로 변환."""
    body = await request.json()
    tool = body["tool"]
    params = body["params"]
    
    try:
        if tool == "run_sql":
            result = await _run_sql(**params)
        elif tool == "query_pv_history":
            result = await _query_pv_history(**params)
        elif tool == "get_tag_metadata":
            result = await _get_tag_metadata(**params)
        elif tool == "list_drawings":
            result = await _list_drawings(**params)
        elif tool == "query_with_nl":
            result = await _query_with_nl(**params)
        else:
            return {"success": False, "error": f"Unknown tool: {tool}"}
        
        return result
    except Exception as e:
        logging.error(f"Error executing {tool}: {e}")
        return {"success": False, "error": str(e)}

async def _run_sql(sql: str) -> str:
    """SQL 실행."""
    conn = _get_db_connection()
    try:
        with conn.cursor() as cur:
            cur.execute(sql)
            if cur.description:
                columns = [desc[0] for desc in cur.description]
                rows = cur.fetchall()
                data = [dict(zip(columns, row)) for row in rows]
                return {
                    "success": True,
                    "columns": columns,
                    "count": len(data),
                    "data": data,
                }
            else:
                conn.commit()
                return {
                    "success": True,
                    "message": f"Query executed successfully. {cur.rowcount} rows affected.",
                }
    finally:
        conn.close()

async def _query_pv_history(tag_names: list[str], time_from: str, time_to: str, limit: int = 100) -> str:
    """과거 값(PV) 히스토리 조회."""
    if not tag_names:
        return {"success": False, "error": "tag_names is required"}
    
    conn = _get_db_connection()
    try:
        with conn.cursor() as cur:
            # TimescaleDB의 time_bucket 함수 사용
            cur.execute(
                """
                SELECT time_bucket('1 min', ts) AS time, tag_name, value
                FROM realtime_table
                WHERE tag_name = ANY(%s)
                  AND ts >= %s
                  AND ts <= %s
                ORDER BY time DESC
                LIMIT %s
                """,
                (tag_names, time_from, time_to, limit),
            )
            columns = ["time", "tag_name", "value"]
            rows = cur.fetchall()
            data = [dict(zip(columns, row)) for row in rows]
            return {
                "success": True,
                "tag_names": tag_names,
                "time_range": {"from": time_from, "to": time_to},
                "limit": limit,
                "count": len(data),
                "data": data,
            }
    finally:
        conn.close()

async def _get_tag_metadata(query: str, limit: int = 10) -> str:
    """태그 메타데이터 검색."""
    conn = _get_db_connection()
    try:
        with conn.cursor() as cur:
            cur.execute(
                """
                SELECT DISTINCT tag_name, unit, description
                FROM realtime_table
                WHERE tag_name ILIKE %s
                ORDER BY tag_name
                LIMIT %s
                """,
                (f"%{query}%", limit),
            )
            columns = ["tag_name", "unit", "description"]
            rows = cur.fetchall()
            data = [dict(zip(columns, row)) for row in rows]
            return {
                "success": True,
                "query": query,
                "count": len(data),
                "tags": data,
            }
    finally:
        conn.close()

async def _list_drawings(unit_no: str = None) -> str:
    """단위별 도면 목록 조회."""
    conn = _get_db_connection()
    try:
        with conn.cursor() as cur:
            if unit_no:
                cur.execute(
                    """
                    SELECT DISTINCT name
                    FROM node_map_master
                    WHERE name LIKE %s
                    ORDER BY name
                    """,
                    (f"{unit_no}%",),
                )
            else:
                cur.execute(
                    """
                    SELECT DISTINCT name
                    FROM node_map_master
                    ORDER BY name
                    """
                )
            columns = ["name"]
            rows = cur.fetchall()
            data = [dict(zip(columns, row[0])) for row in rows]
            return {
                "success": True,
                "unit_no": unit_no,
                "count": len(data),
                "names": [d["name"] for d in data],
            }
    finally:
        conn.close()

async def _query_with_nl(question: str) -> str:
    """자연어로 SQL 쿼리 실행."""
    sql = await _generate_sql(question)
    
    conn = _get_db_connection()
    try:
        with conn.cursor() as cur:
            cur.execute(sql)
            if cur.description:
                columns = [desc[0] for desc in cur.description]
                rows = cur.fetchall()
                data = [dict(zip(columns, row)) for row in rows]
                return {
                    "success": True,
                    "sql": sql,
                    "columns": columns,
                    "count": len(data),
                    "data": data,
                }
            else:
                conn.commit()
                return {
                    "success": True,
                    "sql": sql,
                    "message": f"Query executed successfully. {cur.rowcount} rows affected.",
                }
    except Exception as db_error:
        return {
            "success": False,
            "sql": sql,
            "error": str(db_error),
        }
    finally:
        conn.close()

# ── 메인 ─────────────────────────────────────────────────────────────────────

if __name__ == "__main__":
    port = int(sys.argv[1]) if len(sys.argv) > 1 else 5003
    logging.info(f"Starting NL2SQL worker on port {port}")
    uvicorn.run(app, host="0.0.0.0", port=port)
