import asyncio
import aiosqlite
import json
import logging
from pathlib import Path
from datetime import datetime
import os

from textual.app import App, ComposeResult
from textual.containers import Container, Horizontal, Vertical
from textual.widgets import Header, Footer, DataTable, Static, TextArea, Markdown, Input
from textual.screen import Screen

from app.core.resource import get_combined_metrics
from app.core.config import list_servers, list_providers, DB_PATH

class K9sHeader(Static):
    """K9s 스타일의 상단 헤더 — 컴팩트 1줄 레이아웃"""
    def compose(self) -> ComposeResult:
        user = os.getenv("USER", "local")
        with Horizontal(id="k9s_header"):
            yield Static("[bold green]MCP-LENS[/bold green]  [dim]v0.1.9[/dim]", id="banner_container")
            yield Static(f"[cyan]ctx:[/cyan] Local  [cyan]user:[/cyan] {user}", id="info_container")
            yield Label("[cyan]Cluster:[/cyan] [white]llm4s-v0.1.9[/white]", id="view_info")

class Label(Static):
    """단순 텍스트 레이블"""
    pass

class TopStatus(Static):
    """리소스 상태바"""
    model_name = "None"

    def on_mount(self) -> None:
        self.set_interval(2.0, self.update_stats)

    def update_stats(self) -> None:
        metrics = get_combined_metrics()
        sys = metrics.get("system", {})
        cpu = sys.get("cpu_usage", 0)
        ram = sys.get("ram_percent", 0)
        gpu = metrics.get("main_gpu")
        gpu_str = f"GPU: {gpu['load_percent']}%" if gpu else "GPU: -"
        model_display = metrics.get("hw_model", "None")
        
        status_text = f" CPU:{cpu}% RAM:{ram}% {gpu_str} - MODEL:{model_display} "
        self.update(status_text)

class ProviderListScreen(Screen):
    """Tier 1: Provider Dashboard"""
    def compose(self) -> ComposeResult:
        yield K9sHeader()
        yield TopStatus()
        yield Static(" PROVIDERS ", classes="view_title")
        yield DataTable(id="provider_table")
        yield Static("", id="status_bar", classes="status_bar")
        yield Footer()

    async def on_mount(self) -> None:
        self.query_one("#view_info", Label).update("[cyan]View:   [/cyan] [white]Providers[/white]")
        table = self.query_one("#provider_table", DataTable)
        table.cursor_type = "row"
        # 열 간격 확보 (각 열의 너비 지정)
        table.add_column("PROVIDER", width=20, key="provider")
        table.add_column("STATUS", width=12, key="status")
        table.add_column("SESSIONS", width=10, key="sessions")
        table.add_column("ENV_KEY", width=30, key="env_key")
        table.add_column("KEY_HINT", width=15, key="key_hint")
        # 즉시 로딩 후 주기적 갱신 (K9s 스타일)
        await self.refresh_providers()
        self.set_interval(2.0, self.refresh_providers)

    async def refresh_providers(self) -> None:
        table = self.query_one("#provider_table", DataTable)
        status_bar = self.query_one("#status_bar", Static)

        # 1. 설정 파일의 프로바이더 로드
        providers = list_providers()
        
        db_providers = []
        session_counts = {}
        active_counts = {}
        if DB_PATH.exists():
            try:
                async with aiosqlite.connect(DB_PATH) as db:
                    # 모든 프로바이더별 세션 수 및 실행 중인 세션 수 집계
                    async with db.execute("""
                        SELECT provider, 
                               COUNT(*), 
                               SUM(CASE WHEN status = 'RUNNING' THEN 1 ELSE 0 END) 
                        FROM sessions GROUP BY provider
                    """) as cursor:
                        async for row in cursor:
                            name, total, active = row
                            session_counts[name] = total
                            active_counts[name] = active
                        
                    # 설정에 없지만 DB에 있는 프로바이더 목록
                    for name in session_counts.keys():
                        if name and name not in [p["name"] for p in providers]:
                            db_providers.append({
                                "name": name,
                                "status": "DB-ONLY",
                                "env_key": "-",
                                "key_hint": "-"
                            })
            except Exception as e:
                import logging
                logging.getLogger("llm4s").warning(f"Failed to refresh providers: {e}")
                status_bar.update(f"[red]DB error: {e}[/red]")

        all_providers = providers + db_providers

        if not all_providers:
            table.display = False
            status_bar.update("[dim]No data yet — run: llm4s proxy -- <cmd>  |  ? for help[/dim]")
        else:
            table.display = True
            status_bar.update("")
            
            # 커서 위치 저장
            current_coordinate = table.cursor_coordinate
            
            # 상태 텍스트에 아이콘과 색상 입히기
            def get_status_rich(status):
                if status == "OK": return "[bold green]✓ OK[/bold green]"
                if status == "MISSING": return "[bold yellow]✗ MISS[/bold yellow]"
                if status == "DB-ONLY": return "[bold blue]◎ DB[/bold blue]"
                return status

            # 전체 행 내용 구성
            display_rows = []
            for p in all_providers:
                total = session_counts.get(p["name"], 0)
                active = active_counts.get(p["name"], 0)
                
                # 세션 표시: "합계(실행중)" 형식 (실행 중인 것이 있을 때만 괄호 표시)
                session_str = str(total)
                if active > 0:
                    session_str = f"{total}([bold green]{active}[/bold green])"
                
                status_rich = get_status_rich(p["status"])
                display_rows.append({
                    "key": p["name"],
                    "vals": [p["name"], status_rich, session_str, p["env_key"], p["key_hint"]],
                    "column_keys": ["provider", "status", "sessions", "env_key", "key_hint"]
                })

            if len(table.rows) != len(display_rows):
                table.clear()
                for row in display_rows:
                    table.add_row(*row["vals"], key=row["key"])
                try: table.cursor_coordinate = current_coordinate
                except Exception: pass # 테이블 갱신 중 커서 복구 실패는 치명적이지 않음
            else:
                for row in display_rows:
                    for i, val in enumerate(row["vals"]):
                        table.update_cell(row["key"], row["column_keys"][i], val)

    def on_data_table_row_selected(self, event: DataTable.RowSelected) -> None:
        provider_name = event.row_key.value
        self.app.push_screen(ModelListScreen(provider_name))

class ModelListScreen(Screen):
    """Tier 2: Model List (Grouped by Provider)"""
    def __init__(self, provider_name: str):
        super().__init__()
        self.provider_name = provider_name

    def compose(self) -> ComposeResult:
        yield K9sHeader()
        yield TopStatus()
        yield Static(f" MODELS: {self.provider_name} ", classes="view_title")
        yield DataTable(id="model_table")
        yield Footer()

    async def on_mount(self) -> None:
        self.query_one("#view_info", Label).update(f"[cyan]View:   [/cyan] [white]Models({self.provider_name})[/white]")
        table = self.query_one("#model_table", DataTable)
        table.cursor_type = "row"
        table.add_column("MODEL", width=40, key="model")
        table.add_column("SESSION COUNT", width=15, key="count")
        table.add_column("LAST ACTIVE", width=25, key="last_active")
        
        # 즉시 로딩 후 주기적 갱신
        await self.refresh_models()
        self.set_interval(2.0, self.refresh_models)

    async def refresh_models(self) -> None:
        table = self.query_one("#model_table", DataTable)
        if not DB_PATH.exists(): return
        
        try:
            async with aiosqlite.connect(DB_PATH) as db:
                # 모델별로 세션 수와 최근 활동 시간 집계
                async with db.execute(
                    "SELECT model, COUNT(id), MAX(start_time) FROM sessions WHERE provider = ? GROUP BY model ORDER BY MAX(start_time) DESC",
                    (self.provider_name,)
                ) as cursor:
                    rows = await cursor.fetchall()
                    
                    # 커서 위치 저장
                    current_coordinate = table.cursor_coordinate

                    if len(table.rows) != len(rows):
                        table.clear()
                        for row in rows:
                            model_name, count, last_active = row
                            model_disp = model_name if model_name else "Unknown"
                            table.add_row(model_disp, str(count), str(last_active), key=model_disp)
                        try: table.cursor_coordinate = current_coordinate
                        except: pass
                    else:
                        for row in rows:
                            model_name, count, last_active = row
                            model_disp = model_name if model_name else "Unknown"
                            row_key = model_disp
                            vals = [model_disp, str(count), str(last_active)]
                            column_keys = ["model", "count", "last_active"]
                            for col_idx, val in enumerate(vals):
                                table.update_cell(row_key, column_keys[col_idx], val)
        except aiosqlite.OperationalError as e:
            self.notify(f"DB error: {e}", severity="error")

    def on_data_table_row_selected(self, event: DataTable.RowSelected) -> None:
        model_name = event.row_key.value
        self.app.push_screen(SessionListScreen(self.provider_name, model_name))

class SessionListScreen(Screen):
    """Tier 3: Session List"""
    BINDINGS = [
        ("d", "delete_session", "Delete Session"),
        ("n", "toggle_rename", "Rename"),
        ("C", "clear_all_sessions", "Clear All"),
        ("s", "focus_search", "Search"),
        ("/", "focus_search", "Search"),
    ]

    def __init__(self, provider_name: str, model_name: str):
        super().__init__()
        self.provider_name = provider_name
        self.model_name = model_name
        self.search_query = ""
        self.renaming_session_id = None

    def compose(self) -> ComposeResult:
        yield K9sHeader()
        yield TopStatus()
        yield Static(f" SESSIONS: {self.provider_name} > {self.model_name} ", classes="view_title")
        yield Input(placeholder="Search sessions (ID or Name)...", id="session_search")
        yield Input(placeholder="Enter new session name...", id="rename_input", classes="hidden")
        yield DataTable(id="session_table")
        yield Footer()

    async def on_mount(self) -> None:
        self.query_one("#view_info", Label).update(f"[cyan]View:   [/cyan] [white]Sessions({self.model_name})[/white]")
        table = self.query_one("#session_table", DataTable)
        table.cursor_type = "row"
        # K9s 스타일의 열 구성
        table.add_column("NAME", width=20, key="name")
        table.add_column("STATUS", width=12, key="status")
        table.add_column("AGE", width=10, key="age")
        table.add_column("ID", width=6, key="id")
        table.add_column("MSGS", width=8, key="msgs")
        table.add_column("ERR", width=6, key="err")
        table.add_column("TOKENS", width=10, key="tokens")
        table.add_column("AVG_MS", width=10, key="avgms")
        table.add_column("LATEST", width=20, key="latest")
        
        # 즉시 로딩 후 주기적 갱신
        await self.refresh_sessions()
        self.set_interval(2.0, self.refresh_sessions)
        
        # 테이블에 기본 포커스 설정 (검색창이 아닌 테이블을 먼저 탐색하도록)
        self.call_after_refresh(table.focus)

    def on_input_changed(self, event: Input.Changed) -> None:
        if event.input.id == "session_search":
            self.search_query = event.value.lower()
            # 검색어가 바뀌면 테이블 클리어 후 즉시 갱신 유도
            self.query_one("#session_table", DataTable).clear()
            asyncio.create_task(self.refresh_sessions())

    def on_input_submitted(self, event: Input.Submitted) -> None:
        if event.input.id == "session_search":
            self.query_one("#session_table", DataTable).focus()
        elif event.input.id == "rename_input":
            asyncio.create_task(self.submit_rename(event.value))

    async def submit_rename(self, new_name: str) -> None:
        if self.renaming_session_id and new_name:
            from app.core.logging import db_logger
            await db_logger.rename_session(self.renaming_session_id, new_name)
            self.notify(f"Session renamed to: {new_name}")
        
        self.renaming_session_id = None
        ri = self.query_one("#rename_input", Input)
        ri.add_class("hidden")
        ri.value = ""
        table = self.query_one("#session_table", DataTable)
        table.focus()
        await self.refresh_sessions()

    async def refresh_sessions(self) -> None:
        table = self.query_one("#session_table", DataTable)
        if not DB_PATH.exists(): return
        
        now = datetime.utcnow()

        try:
            async with aiosqlite.connect(DB_PATH) as db:
                query = f"""
                    SELECT s.id, s.name, s.status, s.start_time, 
                           (SELECT COUNT(*) FROM intercept_logs l WHERE l.session_id = s.id) as msg_count,
                           (SELECT SUM(tokens) FROM intercept_logs l WHERE l.session_id = s.id) as total_tokens,
                           (SELECT method FROM intercept_logs l WHERE l.session_id = s.id ORDER BY id DESC LIMIT 1) as last_method,
                           (SELECT COUNT(*) FROM intercept_logs l WHERE l.session_id = s.id AND l.is_error = 1) as error_count,
                           (SELECT AVG(response_time_ms) FROM intercept_logs l WHERE l.session_id = s.id AND l.response_time_ms IS NOT NULL) as avg_ms
                    FROM sessions s
                    WHERE s.provider = ? AND s.model = ?
                """
                params = [self.provider_name, self.model_name]
                
                if self.search_query:
                    query += " AND (LOWER(s.name) LIKE ? OR CAST(s.id AS TEXT) LIKE ?)"
                    params.extend([f"%{self.search_query}%", f"%{self.search_query}%"])
                
                query += " ORDER BY s.id DESC"

                async with db.execute(query, params) as cursor:
                    rows = await cursor.fetchall()
                    
                    current_coordinate = table.cursor_coordinate
                    
                    if len(table.rows) != len(rows):
                        table.clear()
                        for row in rows:
                            sid, sname, status, start_str, msg_count, total_tokens, last_method, error_count, avg_ms = row
                            age_str = self._calc_age(start_str, now)
                            disp_name = sname if sname else f"session-{sid}"
                            latest = last_method if last_method else "-"
                            tokens_str = f"[bold cyan]{total_tokens or 0}[/bold cyan]"
                            err_count = error_count or 0
                            err_str = f"[bold red]🔴 {err_count}[/bold red]" if err_count > 0 else "0"
                            # 응답시간 시각화: 🟢 <1s, 🟡 1~3s, 🔴 >3s
                            if avg_ms is not None:
                                if avg_ms < 1000:
                                    ms_str = f"[green]🟢{avg_ms:.0f}[/green]"
                                elif avg_ms < 3000:
                                    ms_str = f"[yellow]🟡{avg_ms:.0f}[/yellow]"
                                else:
                                    ms_str = f"[red]🔴{avg_ms:.0f}[/red]"
                            else:
                                ms_str = "-"
                            table.add_row(disp_name, status, age_str, str(sid), str(msg_count), err_str, tokens_str, ms_str, latest, key=str(sid))
                        try: table.cursor_coordinate = current_coordinate
                        except: pass
                    else:
                        for i, row in enumerate(rows):
                            sid, sname, status, start_str, msg_count, total_tokens, last_method, error_count, avg_ms = row
                            age_str = self._calc_age(start_str, now)
                            disp_name = sname if sname else f"session-{sid}"
                            latest = last_method if last_method else "-"
                            tokens_str = f"[bold cyan]{total_tokens or 0}[/bold cyan]"
                            err_count = error_count or 0
                            err_str = f"[bold red]🔴 {err_count}[/bold red]" if err_count > 0 else "0"
                            if avg_ms is not None:
                                if avg_ms < 1000:
                                    ms_str = f"[green]🟢{avg_ms:.0f}[/green]"
                                elif avg_ms < 3000:
                                    ms_str = f"[yellow]🟡{avg_ms:.0f}[/yellow]"
                                else:
                                    ms_str = f"[red]🔴{avg_ms:.0f}[/red]"
                            else:
                                ms_str = "-"
                            
                            row_key = str(sid)
                            vals = [disp_name, status, age_str, str(sid), str(msg_count), err_str, tokens_str, ms_str, latest]
                            column_keys = ["name", "status", "age", "id", "msgs", "err", "tokens", "avgms", "latest"]
                            
                            for col_idx, val in enumerate(vals):
                                table.update_cell(row_key, column_keys[col_idx], val)
        except aiosqlite.OperationalError as e:
            self.notify(f"DB error: {e}", severity="error")

    async def action_delete_session(self) -> None:
        table = self.query_one("#session_table", DataTable)
        try:
            row_key, _ = table.coordinate_to_cell_key(table.cursor_coordinate)
            session_id = int(row_key.value)
            
            from app.core.logging import db_logger
            await db_logger.delete_session(session_id)
            self.notify(f"Session {session_id} deleted.")
            await self.refresh_sessions()
        except:
            self.notify("No session selected.", severity="warning")

    async def action_clear_all_sessions(self) -> None:
        from app.core.logging import db_logger
        await db_logger.clear_sessions(self.provider_name, self.model_name)
        self.notify(f"All sessions for {self.model_name} cleared.")
        await self.refresh_sessions()

    def action_focus_search(self) -> None:
        self.query_one("#session_search", Input).focus()

    def action_toggle_rename(self) -> None:
        table = self.query_one("#session_table", DataTable)
        try:
            row_key, _ = table.coordinate_to_cell_key(table.cursor_coordinate)
            self.renaming_session_id = int(row_key.value)
            
            ri = self.query_one("#rename_input", Input)
            ri.remove_class("hidden")
            ri.placeholder = f"Rename Session {self.renaming_session_id}..."
            ri.focus()
        except:
            self.notify("No session selected.", severity="warning")

    def _calc_age(self, start_str: str, now: datetime) -> str:
        try:
            start_dt = datetime.strptime(start_str, "%Y-%m-%d %H:%M:%S")
            diff = now - start_dt
            seconds = int(diff.total_seconds())
            if seconds < 60: return f"{seconds}s"
            elif seconds < 3600: return f"{seconds//60}m{seconds%60}s"
            else: return f"{seconds//3600}h{(seconds%3600)//60}m"
        except: return "-"

    def on_data_table_row_selected(self, event: DataTable.RowSelected) -> None:
        session_id = event.row_key.value
        self.app.push_screen(LogDetailScreen(session_id))

class LogDetailScreen(Screen):
    """Tier 3: Detailed Logs"""
    BINDINGS = [
        ("c", "copy_payload", "Copy Payload"),
        ("e", "toggle_edit", "Edit Payload"),
        ("r", "replay", "Replay (Inject)"),
        ("i", "copy_session_id", "Copy SID"),
        ("l", "copy_log_id", "Copy Log ID"),
        ("x", "export_logs", "Export MD"),
        ("/", "focus_search", "Search"),
        ("s", "focus_search", "Search"),
    ]
    
    def __init__(self, session_id: str):
        super().__init__()
        self.session_id = session_id

    def compose(self) -> ComposeResult:
        yield K9sHeader()
        yield TopStatus()
        yield Static(f" LOGS: Session {self.session_id} ", classes="view_title")
        with Horizontal(id="log_main_split"):
            with Vertical(id="log_left_pane"):
                yield Input(placeholder="Search logs (method or payload)...", id="log_search")
                yield DataTable(id="log_table", classes="log_left")
            with Vertical(id="log_right"):
                with Vertical(id="thought_container"):
                    yield Static(" THINKING / CONTEXT ", classes="panel_title")
                    yield Markdown(id="thought_markdown")
                with Vertical(id="payload_panel"):
                    yield Static(" PAYLOAD DETAIL ", id="editor_title")
                    yield TextArea(id="payload_editor", language="json", show_line_numbers=True)
        yield Footer()

    def on_mount(self) -> None:
        self.query_one("#view_info", Label).update(f"[cyan]View:   [/cyan] [white]Logs(SID:{self.session_id})[/white]")
        table = self.query_one("#log_table", DataTable)
        table.cursor_type = "row"
        table.add_columns("ID", "DIR", "TOKENS", "MS", "METHOD", "PREVIEW")
        self.call_after_refresh(table.focus)
        
        # [신규] TextArea 기본 설정
        editor = self.query_one("#payload_editor", TextArea)
        editor.read_only = True # 기본은 읽기 전용
        
        self.set_interval(1.0, self.refresh_logs)
        self.search_query = ""

    def on_input_changed(self, event: Input.Changed) -> None:
        if event.input.id == "log_search":
            self.search_query = event.value.lower()
            # 검색어가 바뀌면 테이블 클리어 후 즉시 갱신 유도
            self.query_one("#log_table", DataTable).clear()
            asyncio.create_task(self.refresh_logs())

    def on_input_submitted(self, event: Input.Submitted) -> None:
        """검색창에서 엔터를 치면 테이블로 포커스를 이동합니다."""
        if event.input.id == "log_search":
            self.query_one("#log_table", DataTable).focus()

    async def refresh_logs(self) -> None:
        table = self.query_one("#log_table", DataTable)
        if not DB_PATH.exists(): return
        
        query = "SELECT id, direction, method, payload, tokens, is_error, error_type, response_time_ms FROM intercept_logs WHERE session_id = ?"
        params = [self.session_id]
        
        if self.search_query:
            # [고급 필터링] "key:value" 또는 "key>value" 형식 지원
            if ":" in self.search_query or ">" in self.search_query or "<" in self.search_query:
                parts = self.search_query.split()
                for part in parts:
                    if ":" in part:
                        k, v = part.split(":", 1)
                        if k == "method":
                            query += " AND LOWER(method) LIKE ?"
                            params.append(f"%{v}%")
                        elif k == "dir":
                            query += " AND LOWER(direction) = ?"
                            params.append(v.upper())
                        elif k == "error" and v.lower() in ("true", "1", "yes"):
                            query += " AND is_error = 1"
                    elif ">" in part:
                        k, v = part.split(">", 1)
                        if k == "tokens" and v.isdigit():
                            query += " AND tokens > ?"
                            params.append(int(v))
                        elif k == "ms" and v.isdigit():
                            query += " AND response_time_ms > ?"
                            params.append(int(v))
                    elif "<" in part:
                        k, v = part.split("<", 1)
                        if k == "tokens" and v.isdigit():
                            query += " AND tokens < ?"
                            params.append(int(v))
                        elif k == "ms" and v.isdigit():
                            query += " AND response_time_ms < ?"
                            params.append(int(v))
            else:
                query += " AND (LOWER(method) LIKE ? OR LOWER(payload) LIKE ?)"
                params.extend([f"%{self.search_query}%", f"%{self.search_query}%"])
            
        query += " ORDER BY id DESC LIMIT 100"

        async with aiosqlite.connect(DB_PATH) as db:
            async with db.execute(query, params) as cursor:
                rows = await cursor.fetchall()
                existing_ids = {row.value for row in table.rows}
                for row in reversed(rows):
                    log_id, direction, method, payload, tokens, is_error, error_type, response_time_ms = row
                    if str(log_id) not in existing_ids:
                        # 에러 로그 시각적 강조
                        if is_error:
                            arrow = "[bold red]🔴[/bold red]"
                        else:
                            arrow = "[cyan]→[/cyan]" if direction == "C2S" else "[green]←[/green]"
                        
                        # 메서드 강조
                        disp_method = method
                        if is_error:
                            disp_method = f"[bold red]{method}[/bold red]"
                        elif method == "tools/call":
                            disp_method = "[bold magenta]TOOLS/CALL[/bold magenta]"
                        
                        token_str = f"[bold cyan]{tokens}[/bold cyan]" if tokens else "-"
                        
                        # 응답시간 표시
                        if response_time_ms is not None:
                            if response_time_ms < 1000:
                                ms_str = f"[green]{response_time_ms:.0f}[/green]"
                            elif response_time_ms < 3000:
                                ms_str = f"[yellow]{response_time_ms:.0f}[/yellow]"
                            else:
                                ms_str = f"[red]{response_time_ms:.0f}[/red]"
                        else:
                            ms_str = "-"
                        
                        preview = self._format_preview(payload)
                        table.add_row(str(log_id), arrow, token_str, ms_str, disp_method, preview, key=str(log_id))
                        table.scroll_end(animate=False)

    async def on_data_table_row_highlighted(self, event: DataTable.RowHighlighted) -> None:
        """커서가 옮겨질 때마다 우측 상세 패널을 업데이트합니다 (실시간 보기)"""
        if event.row_key is None:
            return
        log_id = event.row_key.value
        await self.update_detail_pane(log_id)

    async def on_data_table_row_selected(self, event: DataTable.RowSelected) -> None:
        """엔터를 쳐도 상세 정보를 업데이트합니다 (기존 하위 호환)"""
        if event.row_key is None:
            return
        log_id = event.row_key.value
        await self.update_detail_pane(log_id)

    async def update_detail_pane(self, log_id: str) -> None:
        if not DB_PATH.exists(): return
        async with aiosqlite.connect(DB_PATH) as db:
            async with db.execute("SELECT payload FROM intercept_logs WHERE id = ?", (log_id,)) as cursor:
                row = await cursor.fetchone()
                if row:
                    payload_str = row[0]
                    editor = self.query_one("#payload_editor", TextArea)
                    try:
                        parsed = json.loads(payload_str)
                        editor.load_text(json.dumps(parsed, indent=2, ensure_ascii=False))
                        editor.refresh()
                        thought = self.extract_thought_process(parsed)

                        md_panel = self.query_one("#thought_markdown", Markdown)
                        if thought:
                            md_panel.update(f"### Reasoning Process\n\n{thought}")
                        else:
                            method = parsed.get("method", parsed.get("result", {}).get("method", "N/A"))
                            params = parsed.get("params", parsed.get("result", {}))
                            params_keys = ", ".join(params.keys()) if isinstance(params, dict) else "-"

                            info = f"### Message Info (ID: {log_id})\n\n"
                            info += f"- **Method**: `{method}`\n"
                            info += f"- **Keys**: `{params_keys}`\n"
                            md_panel.update(info)
                    except Exception as e:
                        editor.load_text(payload_str)
                        editor.refresh()
                        self.query_one("#thought_markdown", Markdown).update(f"Error parsing JSON: {e}")

    def extract_thought_process(self, data: dict) -> str:
        try:
            params = data.get("params", {})
            if "messages" in params:
                for msg in params["messages"]:
                    content = msg.get("content", [])
                    if isinstance(content, list):
                        for item in content:
                            if item.get("type") in ["thought", "thinking"]:
                                return item.get("text", "")
        except: pass
        for key in ["thought", "reasoning", "thinking", "cot"]:
            val = data.get(key) or data.get("params", {}).get(key)
            if val and isinstance(val, str): return val
        return ""

    def action_copy_payload(self) -> None:
        """현재 선택된 페이로드를 클립보드에 복사합니다."""
        editor = self.query_one("#payload_editor", TextArea)
        payload = editor.text
        if payload:
            self.app.copy_to_clipboard(payload)
            self.notify("Payload copied to clipboard!", severity="information")

    def action_focus_search(self) -> None:
        """검색창으로 포커스를 이동합니다."""
        self.query_one("#log_search", Input).focus()

    def _format_preview(self, payload_str: str) -> str:
        """JSON 페이로드에서 핵심 정보만 추출하여 한 줄 프리뷰로 만듭니다."""
        try:
            data = json.loads(payload_str)
            # 1. 에러 응답인 경우
            if "error" in data:
                err = data["error"]
                msg = err.get("message", str(err))
                return f"[bold red]ERR:[/bold red] {msg}"
            
            # 2. 결과 응답인 경우
            if "result" in data:
                res = data["result"]
                if not res: return "{}"
                # 결과 요약 (예: 도구 실행 결과 등)
                res_str = json.dumps(res, ensure_ascii=False)
                return res_str[:60] + "..." if len(res_str) > 60 else res_str

            # 3. 요청/알림인 경우 (params 추출) - 데이터가 비어있어도 보여줌
            if "params" in data:
                params = data["params"]
                if not params and "method" in data:
                    return f"{data['method']} ()"
                p_str = json.dumps(params, ensure_ascii=False)
                return p_str[:60] + "..." if len(p_str) > 60 else p_str
            
            # 4. 기타 JSON (method 등)
            summary = {k:v for k,v in data.items() if k not in ["jsonrpc", "id"]}
            s_str = json.dumps(summary, ensure_ascii=False)
            return s_str[:60] + "..." if len(s_str) > 60 else s_str
        except:
            return payload_str[:50] + "..." if len(payload_str) > 50 else payload_str

    async def action_export_logs(self) -> None:
        """현재 세션의 로그를 Markdown으로 내보냅니다."""
        if not DB_PATH.exists():
            self.notify("Database not found.", severity="error")
            return

        async with aiosqlite.connect(DB_PATH) as db:
            async with db.execute(
                "SELECT timestamp, direction, method, payload, tokens FROM intercept_logs WHERE session_id = ? ORDER BY id ASC",
                (self.session_id,)
            ) as cursor:
                logs = await cursor.fetchall()
        
        if not logs:
            self.notify(f"No logs found for Session {self.session_id}", severity="warning")
            return
            
        export_dir = Path.home() / "Documents" / "llm4s_exports"
        export_dir.mkdir(parents=True, exist_ok=True)
        
        filename = f"session_{self.session_id}_{datetime.now().strftime('%Y%m%d_%H%M%S')}.md"
        filepath = export_dir / filename
        
        content = f"# LLM4s Session Export (ID: {self.session_id})\n\n"
        content += f"- Exported at: {datetime.now().strftime('%Y-%m-%d %H:%M:%S')}\n"
        content += f"- Session ID: {self.session_id}\n\n"
        content += "---\n"
        
        for timestamp, direction, method, payload, tokens in logs:
            arrow = "→ (C2S)" if direction == "C2S" else "← (S2C)"
            token_info = f" ({tokens} tokens)" if tokens else ""
            content += f"### [{timestamp}] {arrow} {method}{token_info}\n"
            try:
                formatted_json = json.dumps(json.loads(payload), indent=2, ensure_ascii=False)
                content += f"```json\n{formatted_json}\n```\n\n"
            except:
                content += f"```\n{payload}\n```\n\n"
        
        with open(filepath, "w", encoding="utf-8") as f:
            f.write(content)
            
        self.notify(f"Exported to: {filepath}", severity="information")

    def action_toggle_edit(self) -> None:
        """페이로드 편집 모드를 토글합니다."""
        editor = self.query_one("#payload_editor", TextArea)
        editor.read_only = not editor.read_only
        
        status = "EDITING" if not editor.read_only else "READ-ONLY"
        self.notify(f"Editor mode: {status}")
        
        if not editor.read_only:
            editor.focus()
            self.query_one("#editor_title", Static).update(" [bold yellow]PAYLOAD DETAIL (EDITING...)[/bold yellow] ")
        else:
            self.query_one("#log_table", DataTable).focus()
            self.query_one("#editor_title", Static).update(" PAYLOAD DETAIL ")

    async def action_replay(self) -> None:
        """현재 편집기 내용을 실제 타겟 프로세스(MCP 서버)로 재전송합니다."""
        editor = self.query_one("#payload_editor", TextArea)
        payload = editor.text.strip()
        
        if not payload:
            self.notify("Cannot replay empty payload.", severity="warning")
            return
            
        try:
            data = json.loads(payload)
        except json.JSONDecodeError:
            self.notify("Invalid JSON payload.", severity="error")
            return

        # JSON-RPC 메서드 화이트리스트 검증
        _ALLOWED_METHODS = {
            "initialize", "ping",
            "tools/list", "tools/call",
            "resources/list", "resources/read", "resources/subscribe",
            "prompts/list", "prompts/get",
            "logging/setLevel",
        }
        method = data.get("method") if isinstance(data, dict) else None
        if method and method not in _ALLOWED_METHODS:
            self.notify(f"Replay blocked: method '{method}' is not allowed.", severity="error")
            return

        from app.core.logging import db_logger
        await db_logger.replay_message(int(self.session_id), payload)
        self.notify("Message queued for replay!", severity="information")
        
        # 편집 모드 종료
        editor.read_only = True
        self.query_one("#editor_title", Static).update(" PAYLOAD DETAIL ")
        self.query_one("#log_table", DataTable).focus()

    def action_copy_session_id(self) -> None:
        self.app.copy_to_clipboard(str(self.session_id))
        self.notify(f"Session ID {self.session_id} copied.")

    def action_copy_log_id(self) -> None:
        table = self.query_one("#log_table", DataTable)
        try:
            row_key, _ = table.coordinate_to_cell_key(table.cursor_coordinate)
            log_id = row_key.value
            self.app.copy_to_clipboard(str(log_id))
            self.notify(f"Log ID {log_id} copied.")
        except:
            self.notify("No log selected.", severity="warning")

class WelcomeScreen(Screen):
    """데이터가 없을 때 표시되는 시작 화면"""

    SPINNERS = ["⠋", "⠙", "⠹", "⠸", "⠼", "⠴", "⠦", "⠧", "⠇", "⠏"]
    _spinner_idx = 0

    def compose(self) -> ComposeResult:
        yield Static(id="welcome_content")
        yield TextArea(
            "llm4s proxy -- node /path/to/your/server.js",
            id="proxy_cmd",
            read_only=True,
        )
        yield Footer()

    def on_mount(self) -> None:
        self._render_welcome()
        self.set_interval(0.5, self._tick_spinner)
        self.set_interval(2.0, self._check_for_data)

    def _render_welcome(self) -> None:
        sp = self.SPINNERS[self._spinner_idx]
        content = (
            "\n"
            " [bold green]MCP-LENS[/bold green]  [dim]v0.1.5 — K9s for MCP[/dim]\n\n"
            " Watching for MCP traffic...  " + sp + "\n\n"
            " ─────────────────────────────────────────────────\n\n"
            " [bold]Quick Start[/bold]\n\n"
            "  [cyan]Step 1[/cyan]  Configure API key:\n"
            "          [dim]llm4s setup[/dim]  or  [dim]llm4s config --add gemini[/dim]\n\n"
            "  [cyan]Step 2[/cyan]  Start the proxy in another terminal\n"
            "          [dim](edit the path below, then copy & run)[/dim]\n\n"
            "  [cyan]Step 3[/cyan]  Data appears here automatically.\n\n"
            " ─────────────────────────────────────────────────\n\n"
            " [dim]Press [bold]?[/bold] for keybindings  ·  [bold]q[/bold] to quit[/dim]\n"
        )
        self.query_one("#welcome_content", Static).update(content)

    def _tick_spinner(self) -> None:
        self._spinner_idx = (self._spinner_idx + 1) % len(self.SPINNERS)
        self._render_welcome()

    async def _check_for_data(self) -> None:
        """세션이 생기면 메인 화면으로 자동 이동"""
        if not DB_PATH.exists():
            return
        try:
            async with aiosqlite.connect(DB_PATH) as db:
                async with db.execute("SELECT COUNT(*) FROM sessions") as cur:
                    row = await cur.fetchone()
                    if row and row[0] > 0:
                        self.app.pop_screen()
                        self.app.push_screen(ProviderListScreen())
        except Exception:
            pass


class HelpScreen(Screen):
    """? 키로 열리는 키바인딩 오버레이"""

    BINDINGS = [
        ("escape", "dismiss_help", "Close"),
        ("q", "dismiss_help", "Close"),
        ("?", "dismiss_help", "Close"),
    ]

    def compose(self) -> ComposeResult:
        yield Static(id="help_content")
        yield Footer()

    def on_mount(self) -> None:
        help_text = (
            "\n"
            " [bold cyan]MCP-LENS Keybindings[/bold cyan]\n\n"
            " [bold]Navigation[/bold]\n"
            "   [cyan]Enter[/cyan]    Drill down (Provider → Model → Session → Logs)\n"
            "   [cyan]Esc[/cyan]      Go back\n"
            "   [cyan]F5  / R[/cyan]  Refresh current view\n"
            "   [cyan]?[/cyan]        Show this help\n"
            "   [cyan]q[/cyan]        Quit\n\n"
            " [bold]Search[/bold]\n"
            "   [cyan]/ or s[/cyan]   Focus search bar\n"
            "   Advanced: [dim]method:tools/call  tokens>500  dir:S2C  error:true[/dim]\n\n"
            " [bold]Log Detail (inside a session)[/bold]\n"
            "   [cyan]e[/cyan]        Edit payload (toggle)\n"
            "   [cyan]r[/cyan]        Replay edited message\n"
            "   [cyan]c[/cyan]        Copy payload to clipboard\n"
            "   [cyan]x[/cyan]        Export session to Markdown\n"
            "   [cyan]i[/cyan]        Copy session ID\n"
            "   [cyan]l[/cyan]        Copy log ID\n\n"
            " [bold]Session Management[/bold]\n"
            "   [cyan]n[/cyan]        Rename session\n"
            "   [cyan]d[/cyan]        Delete session\n"
            "   [cyan]C[/cyan]        Clear all sessions\n\n"
            " [bold]Filter Syntax[/bold]\n"
            "   [dim]method:tools/call[/dim]  Filter by method\n"
            "   [dim]dir:C2S[/dim]            Filter by direction (C2S / S2C)\n"
            "   [dim]tokens>500[/dim]         Token threshold\n"
            "   [dim]error:true[/dim]         Errors only\n"
            "   [dim]ms>3000[/dim]            Slow responses (>3s)\n\n"
            " [dim]Press Esc, q, or ? to close[/dim]\n"
        )
        self.query_one("#help_content", Static).update(help_text)

    def action_dismiss_help(self) -> None:
        self.app.pop_screen()


class MCPLensApp(App):
    """K9s 스타일에 최적화된 LLM4s TUI"""
    CSS_PATH = Path(__file__).parent / "styles.tcss"
    BINDINGS = [
        ("q", "quit", "Quit"),
        ("escape", "safe_pop_screen", "Back"),
        ("R", "refresh_current", "Refresh"),
        ("f5", "refresh_current", "Refresh"),
        ("?", "show_help", "Help"),
    ]

    def action_safe_pop_screen(self) -> None:
        if isinstance(self.focused, Input):
            if isinstance(self.screen, SessionListScreen):
                self.screen.query_one("#session_table", DataTable).focus()
            elif isinstance(self.screen, LogDetailScreen):
                self.screen.query_one("#log_table", DataTable).focus()
            return

        if isinstance(self.screen, (ProviderListScreen, WelcomeScreen)):
            return

        if len(self.screen_stack) > 1:
            self.pop_screen()

    def action_refresh_current(self) -> None:
        screen = self.screen
        if hasattr(screen, "refresh_providers"): asyncio.create_task(screen.refresh_providers())
        elif hasattr(screen, "refresh_models"): asyncio.create_task(screen.refresh_models())
        elif hasattr(screen, "refresh_sessions"): asyncio.create_task(screen.refresh_sessions())
        elif hasattr(screen, "refresh_logs"): asyncio.create_task(screen.refresh_logs())

    def action_show_help(self) -> None:
        if not isinstance(self.screen, HelpScreen):
            self.push_screen(HelpScreen())

    async def on_mount(self) -> None:
        await self._ensure_db_schema()
        # DB에 세션이 있으면 바로 Provider 화면, 없으면 Welcome 화면
        has_data = False
        if DB_PATH.exists():
            try:
                async with aiosqlite.connect(DB_PATH) as db:
                    async with db.execute("SELECT COUNT(*) FROM sessions") as cur:
                        row = await cur.fetchone()
                        has_data = bool(row and row[0] > 0)
            except Exception:
                pass
        if has_data:
            self.push_screen(ProviderListScreen())
        else:
            self.push_screen(WelcomeScreen())

    async def _ensure_db_schema(self) -> None:
        """TUI 시작 시 DB 스키마가 최신인지 확인하고 마이그레이션을 실행합니다."""
        if not DB_PATH.exists():
            return
        try:
            async with aiosqlite.connect(DB_PATH) as db:
                migrations = [
                    "ALTER TABLE intercept_logs ADD COLUMN is_error BOOLEAN DEFAULT 0",
                    "ALTER TABLE intercept_logs ADD COLUMN error_type TEXT",
                    "ALTER TABLE intercept_logs ADD COLUMN response_time_ms FLOAT",
                    "ALTER TABLE intercept_logs ADD COLUMN jsonrpc_id TEXT",
                ]
                for sql in migrations:
                    try:
                        await db.execute(sql)
                    except Exception:
                        pass  # 이미 존재하면 무시
                await db.commit()
        except Exception as e:
            import logging
            logging.getLogger("llm4s").warning(f"Failed to ensure DB schema: {e}")

def run_tui():
    app = MCPLensApp()
    app.run()

if __name__ == "__main__":
    run_tui()
