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

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

# --- 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 ---

async def _get_db():
    if not DB_PATH.exists():
        raise HTTPException(status_code=503, detail="Database not initialized. Run a proxy session first.")
    return await aiosqlite.connect(DB_PATH)


# --- Sessions API ---

@app.get("/sessions", response_model=List[SessionSummary])
async def list_sessions(
    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 = await _get_db()
    try:
        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
            ]
    finally:
        await db.close()


@app.get("/sessions/{session_id}", response_model=SessionDetail)
async def get_session(session_id: int):
    """특정 세션의 상세 정보를 반환합니다."""
    db = await _get_db()
    try:
        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]
            )
    finally:
        await db.close()


@app.patch("/sessions/{session_id}")
async def update_session(session_id: int, body: SessionUpdate):
    """세션의 이름이나 상태를 변경합니다."""
    db = await _get_db()
    try:
        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}
    finally:
        await db.close()


@app.delete("/sessions/{session_id}")
async def delete_session(session_id: int):
    """세션과 관련 로그, 명령을 전부 삭제합니다."""
    db = await _get_db()
    try:
        # 존재 여부 확인
        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}
    finally:
        await db.close()


@app.get("/sessions/{session_id}/logs", response_model=List[LogEntry])
async def get_session_logs(
    session_id: int,
    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 = await _get_db()
    try:
        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
            ]
    finally:
        await db.close()


# --- Timing Profiler API ---

@app.get("/sessions/{session_id}/timing", response_model=TimingStats)
async def get_timing_stats(session_id: int):
    """C2S→S2C 응답시간 통계를 반환합니다. (Thinking Time Profiler)"""
    db = await _get_db()
    try:
        # 세션 존재 확인
        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,
        )
    finally:
        await db.close()


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

@app.get("/logs", response_model=List[LogEntry])
async def get_logs(
    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 필터 지원."""
    if not DB_PATH.exists():
        return []

    async with aiosqlite.connect(DB_PATH) as db:
        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():
    """통계 요약을 반환합니다."""
    if not DB_PATH.exists():
        return {"total_requests": 0, "status": "No DB"}

    async with aiosqlite.connect(DB_PATH) as db:
        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",
        }
