import asyncio
import sys
import os
import logging
import subprocess
import argparse
from typing import List

from app.core.logging import db_logger

logger = logging.getLogger(__name__)

class InterceptorProxy:
    """
    LLM (Claude Desktop 등)과 실제 MCP 서버 사이에서
    표준 입출력(stdio)을 통째로 중계하며 중간에서 메시지를 가로채어(Intercept) 기록하는 클래스입니다.
    """
    def __init__(self, target_command: List[str], provider: str = "Unknown", model: str = "Unknown"):
        """
        :param target_command: 실제 구동할 MCP 서버 명령어 (예: ["node", "index.js"])
        """
        self.target_command = target_command
        self.provider = provider
        self.model = model
        self.process: asyncio.subprocess.Process | None = None
        self.session_id: int | None = None

    async def start(self):
        """프록시 및 타겟 프로세스를 시작합니다."""
        # 1. 아까 만든 비동기 SQLite 로거 백그라운드 부팅
        await db_logger.initialize()

        # [세션 생성] K9s의 Pod과 같은 개념으로 세션 시작 기록
        # Claude Desktop inject 시 환경변수로 프로젝트 컨텍스트 전달됨
        project_name = os.environ.get("LLM4S_PROJECT_NAME")
        inject_source = "claude_desktop" if os.environ.get("LLM4S_PROJECT_NAME") else "manual"
        self.session_id = await db_logger.create_session(
            self.provider, self.model,
            project_name=project_name,
            inject_source=inject_source,
        )
        # CI/자동화 스크립트에서 세션 ID를 파싱할 수 있도록 stderr에 출력
        print(f"llm4s_session_id={self.session_id}", file=sys.stderr, flush=True)

        # [환경 변수 주입] 등록된 API 키가 있다면 주입
        from app.core.config import list_providers
        env = os.environ.copy()
        providers = list_providers()
        for p in providers:
            env_key = p.get("env_key")
            api_key = p.get("api_key")
            if env_key and api_key:
                env[env_key] = api_key

        # 2. 타겟 MCP 서버 자식 프로세스로 실행 (stdin, stdout 파이프 연결)
        try:
            self.process = await asyncio.create_subprocess_exec(
                *self.target_command,
                stdin=asyncio.subprocess.PIPE,
                stdout=asyncio.subprocess.PIPE,
                stderr=sys.stderr,
                env=env
            )
        except Exception as e:
            print(f"Error starting target process: {e}", file=sys.stderr)
            if self.session_id:
                await db_logger.update_session_status(self.session_id, "ERROR")
            return

        # 3. 양방향 데이터 중계를 위한 두 가지 비동기 태스크 시작
        #   Task 1: Claude(현재 sys.stdin) -> Target(process.stdin)
        #   Task 2: Target(process.stdout)  -> Claude(현재 sys.stdout)
        client_to_server_task = asyncio.create_task(self._forward_client_to_server())
        server_to_client_task = asyncio.create_task(self._forward_server_to_client())
        command_poll_task = asyncio.create_task(self._poll_commands())

        try:
            # 타겟 프로세스가 끝날때까지 (또는 이 프록시가 종료될 때까지) 대기
            await asyncio.gather(client_to_server_task, server_to_client_task, command_poll_task)
            await self.process.wait()
        except Exception as e:
            if self.session_id:
                await db_logger.update_session_status(self.session_id, "ERROR")
            raise e
        finally:
            if self.session_id:
                await db_logger.update_session_status(self.session_id, "COMPLETED")
            await db_logger.shutdown()

    async def _forward_client_to_server(self):
        """클라이언트(프롬프트) -> 실제 서버 방향 (C2S)"""
        # Python의 sys.stdin은 동기식이므로, asyncio 이벤트 루프 안에서 
        # 블로킹 없이 읽기 위해 StreamReader를 하나 래핑해서 생성합니다.
        loop = asyncio.get_running_loop()
        client_reader = asyncio.StreamReader()
        protocol = asyncio.StreamReaderProtocol(client_reader)
        await loop.connect_read_pipe(lambda: protocol, sys.stdin)

        while True:
            # MCP는 JSON-RPC 통신을 줄바꿈(\n) 단위로 주고받습니다.
            line = await client_reader.readline()
            if not line:
                break # 스트림 종료 (EOF)

            # [기록] DB 로거 큐에 즉시 삽입 (성능 저하 없음)
            await db_logger.log_message("C2S", line, session_id=self.session_id)

            # [개입 (향후 구현)] Phase 3에서는 여기서 UI의 Breakpoint 신호를 기다려서 내용을 수정할 수 이습니다.

            # [전송] 타겟 프로세스 (실제 MCP 서버)로 그대로 전달
            if self.process and self.process.stdin:
                try:
                    self.process.stdin.write(line)
                    await self.process.stdin.drain()
                except (ConnectionResetError, BrokenPipeError):
                    print(f"Target process '{self.target_command[0]}' disconnected.", file=sys.stderr)
                    break

    async def _forward_server_to_client(self):
        """실제 서버 -> 클라이언트(LLM) 방향 (S2C)"""
        # stdout 출력용
        loop = asyncio.get_running_loop()
        
        if not self.process or not self.process.stdout:
            return

        while True:
            # 타겟 프로세스에서 응답한 JSON-RPC 줄라인 읽기
            line = await self.process.stdout.readline()
            if not line:
                break # 타겟 프로세스 종료 (EOF)

            # [기록] DB 로거 큐에 즉시 삽입 (성능 저하 없음)
            await db_logger.log_message("S2C", line, session_id=self.session_id)

            # [전송] 최종 클라이언트 (sys.stdout)로 화면 출력
            sys.stdout.buffer.write(line)
            sys.stdout.buffer.flush()

    async def _poll_commands(self):
        """TUI로부터 주입된 명령(Replay 등)이 있는지 주기적으로 확인합니다."""
        if not self.session_id: return
        
        while True:
            await asyncio.sleep(0.5)
            
            try:
                rows = await db_logger.get_pending_commands(self.session_id)
                
                for cmd_id, payload in rows:
                    if self.process and self.process.stdin:
                        message = payload.strip() + "\n"
                        self.process.stdin.write(message.encode())
                        await self.process.stdin.drain()
                        
                        await db_logger.mark_command_sent(cmd_id)
                        
                        # 가로채 로그에도 기록 (C2S 방향)
                        await db_logger.log_message("C2S", f"[REPLAY] {message}".encode(), session_id=self.session_id)
            except Exception as e:
                logger.warning(f"Command polling error (SID:{self.session_id}): {e}")

