from contextlib import asynccontextmanager
from fastapi import FastAPI, HTTPException, Query, Request
from pydantic import BaseModel
from typing import List, Optional
import aiosqlite
from pathlib import Path
from app.core.config import DB_PATH


@asynccontextmanager
async def lifespan(app: FastAPI):
    if DB_PATH.exists():
        db = await aiosqlite.connect(DB_PATH)
        await db.execute("PRAGMA journal_mode=WAL")
        await db.execute("PRAGMA busy_timeout=5000")
        app.state.db = db
    else:
        app.state.db = None
    yield
    if app.state.db:
        await app.state.db.close()


app = FastAPI(
    title="LLM4s (MCP-Lens) API",
    description="REST API for accessing locally intercepted MCP JSON-RPC logs.",
    version="0.2.0",
    lifespan=lifespan,
)

# --- Pydantic Models ---

class LogEntry(BaseModel):
    id: int
    timestamp: str
    direction: str
    method: str
    payload: str
    session_id: Optional[int] = None
    tokens: Optional[int] = None
    size_bytes: Optional[int] = None
    is_error: Optional[bool] = False
    error_type: Optional[str] = None
    response_time_ms: Optional[float] = None
    jsonrpc_id: Optional[str] = None

class SessionSummary(BaseModel):
    id: int
    name: Optional[str] = None
    provider: str
    model: str
    status: str
    start_time: str
    end_time: Optional[str] = None
    log_count: int = 0
    error_count: int = 0
    avg_response_ms: Optional[float] = None

class SessionDetail(SessionSummary):
    total_tokens: int = 0
    total_size_bytes: int = 0

class SessionUpdate(BaseModel):
    name: Optional[str] = None
    status: Optional[str] = None

class TimingEntry(BaseModel):
    method: str
    count: int
    avg_ms: float
    max_ms: float
    min_ms: float

class TimingStats(BaseModel):
    session_id: int
    total_calls: int
    avg_response_ms: float
    max_response_ms: float
    per_method: List[TimingEntry]


# --- Helper ---

def _get_db(request: Request) -> aiosqlite.Connection:
    db = request.app.state.db
    if db is None:
        raise HTTPException(status_code=503, detail="Database not initialized. Run a proxy session first.")
    return db


# --- Sessions API ---

@app.get("/sessions", response_model=List[SessionSummary])
async def list_sessions(
    request: Request,
    provider: Optional[str] = None,
    status: Optional[str] = None,
    limit: int = Query(100, ge=1, le=1000),
    offset: int = Query(0, ge=0),
):
    """세션 목록을 반환합니다. provider, status 필터 지원."""
    db = _get_db(request)
    query = """
        SELECT s.id, s.name, s.provider, s.model, s.status,
               s.start_time, s.end_time,
               COUNT(l.id) as log_count,
               SUM(CASE WHEN l.is_error = 1 THEN 1 ELSE 0 END) as error_count,
               AVG(CASE WHEN l.response_time_ms IS NOT NULL THEN l.response_time_ms END) as avg_response_ms
        FROM sessions s
        LEFT JOIN intercept_logs l ON s.id = l.session_id
        WHERE 1=1
    """
    params = []

    if provider:
        query += " AND s.provider = ?"
        params.append(provider)
    if status:
        query += " AND s.status = ?"
        params.append(status)

    query += " GROUP BY s.id ORDER BY s.start_time DESC LIMIT ? OFFSET ?"
    params.extend([limit, offset])

    async with db.execute(query, params) as cursor:
        rows = await cursor.fetchall()
        return [
            SessionSummary(
                id=r[0], name=r[1], provider=r[2], model=r[3], status=r[4],
                start_time=r[5], end_time=r[6], log_count=r[7],
                error_count=r[8] or 0,
                avg_response_ms=round(r[9], 1) if r[9] else None
            ) for r in rows
        ]


@app.get("/sessions/{session_id}", response_model=SessionDetail)
async def get_session(session_id: int, request: Request):
    """특정 세션의 상세 정보를 반환합니다."""
    db = _get_db(request)
    async with db.execute("""
        SELECT s.id, s.name, s.provider, s.model, s.status,
               s.start_time, s.end_time,
               COUNT(l.id) as log_count,
               SUM(CASE WHEN l.is_error = 1 THEN 1 ELSE 0 END) as error_count,
               AVG(CASE WHEN l.response_time_ms IS NOT NULL THEN l.response_time_ms END) as avg_response_ms,
               COALESCE(SUM(l.tokens), 0) as total_tokens,
               COALESCE(SUM(l.size_bytes), 0) as total_size_bytes
        FROM sessions s
        LEFT JOIN intercept_logs l ON s.id = l.session_id
        WHERE s.id = ?
        GROUP BY s.id
    """, (session_id,)) as cursor:
        r = await cursor.fetchone()
        if not r:
            raise HTTPException(status_code=404, detail=f"Session {session_id} not found")
        return SessionDetail(
            id=r[0], name=r[1], provider=r[2], model=r[3], status=r[4],
            start_time=r[5], end_time=r[6], log_count=r[7],
            error_count=r[8] or 0,
            avg_response_ms=round(r[9], 1) if r[9] else None,
            total_tokens=r[10], total_size_bytes=r[11]
        )


@app.patch("/sessions/{session_id}")
async def update_session(session_id: int, body: SessionUpdate, request: Request):
    """세션의 이름이나 상태를 변경합니다."""
    db = _get_db(request)
    updates = []
    params = []
    if body.name is not None:
        updates.append("name = ?")
        params.append(body.name)
    if body.status is not None:
        valid_statuses = {"RUNNING", "IDLE", "COMPLETED", "ERROR"}
        if body.status not in valid_statuses:
            raise HTTPException(status_code=400, detail=f"Invalid status. Must be one of: {valid_statuses}")
        updates.append("status = ?")
        params.append(body.status)
        if body.status in ("COMPLETED", "ERROR"):
            updates.append("end_time = CURRENT_TIMESTAMP")

    if not updates:
        raise HTTPException(status_code=400, detail="No fields to update. Provide 'name' or 'status'.")

    query = f"UPDATE sessions SET {', '.join(updates)} WHERE id = ?"
    params.append(session_id)

    result = await db.execute(query, params)
    await db.commit()

    if result.rowcount == 0:
        raise HTTPException(status_code=404, detail=f"Session {session_id} not found")

    return {"ok": True, "session_id": session_id}


@app.delete("/sessions/{session_id}")
async def delete_session(session_id: int, request: Request):
    """세션과 관련 로그, 명령을 전부 삭제합니다."""
    db = _get_db(request)
    async with db.execute("SELECT id FROM sessions WHERE id = ?", (session_id,)) as cursor:
        if not await cursor.fetchone():
            raise HTTPException(status_code=404, detail=f"Session {session_id} not found")

    await db.execute("DELETE FROM intercept_logs WHERE session_id = ?", (session_id,))
    await db.execute("DELETE FROM commands WHERE session_id = ?", (session_id,))
    await db.execute("DELETE FROM sessions WHERE id = ?", (session_id,))
    await db.commit()
    return {"ok": True, "deleted_session_id": session_id}


@app.get("/sessions/{session_id}/logs", response_model=List[LogEntry])
async def get_session_logs(
    session_id: int,
    request: Request,
    direction: Optional[str] = None,
    method: Optional[str] = None,
    errors_only: bool = False,
    limit: int = Query(100, ge=1, le=1000),
    offset: int = Query(0, ge=0),
):
    """특정 세션의 로그를 반환합니다. direction, method, errors_only 필터 지원."""
    db = _get_db(request)
    query = """
        SELECT id, timestamp, direction, method, payload,
               session_id, tokens, size_bytes, is_error, error_type,
               response_time_ms, jsonrpc_id
        FROM intercept_logs WHERE session_id = ?
    """
    params: list = [session_id]

    if direction:
        query += " AND direction = ?"
        params.append(direction)
    if method:
        query += " AND method = ?"
        params.append(method)
    if errors_only:
        query += " AND is_error = 1"

    query += " ORDER BY id DESC LIMIT ? OFFSET ?"
    params.extend([limit, offset])

    async with db.execute(query, params) as cursor:
        rows = await cursor.fetchall()
        return [
            LogEntry(
                id=r[0], timestamp=r[1], direction=r[2], method=r[3], payload=r[4],
                session_id=r[5], tokens=r[6], size_bytes=r[7],
                is_error=bool(r[8]), error_type=r[9],
                response_time_ms=r[10], jsonrpc_id=r[11]
            ) for r in rows
        ]


# --- Timing Profiler API ---

@app.get("/sessions/{session_id}/timing", response_model=TimingStats)
async def get_timing_stats(session_id: int, request: Request):
    """C2S→S2C 응답시간 통계를 반환합니다. (Thinking Time Profiler)"""
    db = _get_db(request)
    async with db.execute("SELECT id FROM sessions WHERE id = ?", (session_id,)) as cursor:
        if not await cursor.fetchone():
            raise HTTPException(status_code=404, detail=f"Session {session_id} not found")

    async with db.execute("""
        SELECT COUNT(*), AVG(response_time_ms), MAX(response_time_ms)
        FROM intercept_logs
        WHERE session_id = ? AND response_time_ms IS NOT NULL
    """, (session_id,)) as cursor:
        total = await cursor.fetchone()
        total_calls = total[0] or 0
        avg_ms = round(total[1] or 0, 1)
        max_ms = round(total[2] or 0, 1)

    async with db.execute("""
        SELECT method,
               COUNT(*) as cnt,
               AVG(response_time_ms) as avg_ms,
               MAX(response_time_ms) as max_ms,
               MIN(response_time_ms) as min_ms
        FROM intercept_logs
        WHERE session_id = ? AND response_time_ms IS NOT NULL AND direction = 'S2C'
        GROUP BY method
        ORDER BY avg_ms DESC
    """, (session_id,)) as cursor:
        method_rows = await cursor.fetchall()

    per_method = [
        TimingEntry(
            method=r[0] or "unknown",
            count=r[1],
            avg_ms=round(r[2], 1),
            max_ms=round(r[3], 1),
            min_ms=round(r[4], 1)
        ) for r in method_rows
    ]

    return TimingStats(
        session_id=session_id,
        total_calls=total_calls,
        avg_response_ms=avg_ms,
        max_response_ms=max_ms,
        per_method=per_method,
    )


# --- Existing Endpoints (Enhanced) ---

@app.get("/logs", response_model=List[LogEntry])
async def get_logs(
    request: Request,
    limit: int = Query(100, ge=1, le=1000),
    offset: int = Query(0, ge=0),
    direction: Optional[str] = None,
    session_id: Optional[int] = None,
    method: Optional[str] = None,
    errors_only: bool = False,
):
    """모든 로그를 반환합니다. session_id, direction, method, errors_only 필터 지원."""
    db = _get_db(request)
    query = """
        SELECT id, timestamp, direction, method, payload,
               session_id, tokens, size_bytes, is_error, error_type,
               response_time_ms, jsonrpc_id
        FROM intercept_logs WHERE 1=1
    """
    params: list = []

    if direction:
        query += " AND direction = ?"
        params.append(direction)
    if session_id is not None:
        query += " AND session_id = ?"
        params.append(session_id)
    if method:
        query += " AND method = ?"
        params.append(method)
    if errors_only:
        query += " AND is_error = 1"

    query += " ORDER BY id DESC LIMIT ? OFFSET ?"
    params.extend([limit, offset])

    async with db.execute(query, params) as cursor:
        rows = await cursor.fetchall()
        return [
            LogEntry(
                id=r[0], timestamp=r[1], direction=r[2], method=r[3], payload=r[4],
                session_id=r[5], tokens=r[6], size_bytes=r[7],
                is_error=bool(r[8]) if r[8] is not None else False,
                error_type=r[9], response_time_ms=r[10], jsonrpc_id=r[11]
            ) for r in rows
        ]


@app.get("/stats")
async def get_stats(request: Request):
    """통계 요약을 반환합니다."""
    if request.app.state.db is None:
        return {"total_requests": 0, "status": "No DB", "dropped_log_count": 0}

    db = _get_db(request)
    async with db.execute("SELECT COUNT(*) FROM intercept_logs") as cursor:
        total = (await cursor.fetchone())[0]

    async with db.execute("SELECT COUNT(*) FROM intercept_logs WHERE is_error = 1") as cursor:
        errors = (await cursor.fetchone())[0]

    async with db.execute("SELECT COUNT(*) FROM sessions") as cursor:
        sessions = (await cursor.fetchone())[0]

    async with db.execute("SELECT AVG(response_time_ms) FROM intercept_logs WHERE response_time_ms IS NOT NULL") as cursor:
        avg_time = (await cursor.fetchone())[0]

    return {
        "total_intercepted_messages": total,
        "total_errors": errors,
        "total_sessions": sessions,
        "avg_response_ms": round(avg_time, 1) if avg_time else None,
        "status": "RUNNING",
    }
