from fastapi import APIRouter, HTTPException, Depends
from fastapi.responses import StreamingResponse
from ..schemas.models import ProcessRequest_chat, ProcessResponse_chat, ProcessRequest_get, ProcessResponse_get, ProcessRequest_set, ProcessResponse_set, ProcessResponse_delete_session, ProcessRequest_delete_session
from fastgpt_client import AsyncChatClient
from ..backends.chat import (
    ChatBackend,
    ChatBackendAPIError,
    ChatBackendAuthenticationError,
    ChatBackendRateLimitError,
    ChatInput,
    FormUpdate,
    TextDelta,
)
from ..core.fastgpt_client import get_chat_backend, get_fastgpt_client
from ..core.config import Config
from ..utils.text_chunker import SentenceTextChunker, SentenceTextChunkerConfig
from loguru import logger
import json
import re
import time

router = APIRouter()
STATE_TAG_PATTERN = re.compile(r"<state>\s*(\d+)\s*</state>", flags=re.DOTALL)
STATUS_CODE_MAP = {
    '0000': '结束通话',
    '0001': '转接人工',
    '0002': '语义无法识别转接人工',
    '0003': '有人伤转接人工',
    '0005': '拍照连续识别失败转人工',
    '1001': '未准备好通话',
    '1002': '通话中',
    '2000': '进入单车拍照',
    '2001': '请对准车辆碰撞部位拍摄照片',
    '2002': '请对准被撞物品拍摄照片',
    '2003': '请切换摄像头对准本人拍摄一张正面照片',
    '2004': '确认单车车牌',
    '2005': '请确认车损位置是在车辆前方、后方还是侧面',
    '2010': '进入双车拍照',
    '2011': '请对准第一辆车碰撞部位拍摄',
    '2012': '请对准第二辆车碰撞部位拍摄',
    '2013': '请对准第二方车辆侧后方，看清车牌拍摄',
    '2014': '请拍摄另一方驾驶人的正面照片',
    '2015': '请切换前置摄像头对准本人拍摄一张正面照片',
    '2016': '确认双车中的车牌'
}

def normalize_stage_code(stage_code: str) -> str:
    """Normalize FastGPT stage codes to external API stage codes."""
    if stage_code in ['3001', '3002', '1002']:
        return '1002'
    if stage_code == '2006':
        return '2004'
    if stage_code == '2017':
        return '2016'
    if stage_code == '2020':
        return '0002'
    return stage_code


def extract_first_state_and_clean_content(text: str) -> tuple[str | None, str]:
    """Return the first state code and content with all state tags removed."""
    match = STATE_TAG_PATTERN.search(text)
    if not match:
        return None, text
    return match.group(1), STATE_TAG_PATTERN.sub("", text)


def format_set_info_input(payload: dict, include_input_info: bool) -> str:
    """Build optional setInfo input for FastGPT helper calls."""
    if not include_input_info:
        return ""
    return f"<setInfo>{json.dumps(payload, ensure_ascii=False)}</setInfo>"

async def delete_last_two_chat_records(
    client: AsyncChatClient,
    session_id: str
) -> None:
    """Delete the last two chat records."""
    try:
        # Get chat records using SDK
        response = await client.get_chat_records(
            appId=Config.FASTGPT_APP_ID,
            chatId=session_id,
            offset=0,
            pageSize=10
        )
        response.raise_for_status()
        data = response.json()
        
        records = data.get('data', {}).get('list', [])
        if len(records) < 2:
            logger.warning(f"Less than 2 records found for session {session_id}")
            return
            
        last_two_data_ids = [record['dataId'] for record in records[-2:]]
        logger.info(f"last_two_data_ids: {last_two_data_ids}")
        
        # Delete records using SDK
        for content_id in last_two_data_ids:
            delete_response = await client.delete_chat_record(
                appId=Config.FASTGPT_APP_ID,
                chatId=session_id,
                contentId=content_id
            )
            delete_response.raise_for_status()
            
    except Exception as e:
        logger.error(f"Error deleting chat records: {e}")
        raise


# @app.exception_handler(RequestValidationError)
# async def validation_exception_handler(request: Request, exc: RequestValidationError):
#     logging.error(f"Validation Error: {exc.errors()}")
#     raise HTTPException(status_code=422, detail=exc.errors())

def create_sse_event(event: str, data: dict) -> str:
    """Format data as an SSE event."""
    return f"event: {event}\ndata: {json.dumps(data, ensure_ascii=False)}\n\n"

@router.post("/chat", response_model=ProcessResponse_chat)
async def chat(
    request: ProcessRequest_chat,
    stream: bool = False,
    backend: ChatBackend = Depends(get_chat_backend)
):
    """Handle chat completion request."""
    json_data = request.model_dump()
    need_form_update = json_data.get('needFormUpdate', False)
    use_text_chunk = json_data.get('useTextChunk', False)
    chat_input = ChatInput(
        session_id=json_data['sessionId'],
        text=json_data['text'],
        need_form_update=need_form_update,
    )
    request_started_at = time.perf_counter()
    logger.info(
        "Chat request received "
        f"sessionId={json_data['sessionId']} stream={stream} "
        f"needFormUpdate={need_form_update} useTextChunk={use_text_chunk} "
        f"text_len={len(json_data.get('text', ''))} "
        f"input={json_data.get('text', '')!r}"
    )
    
    if stream:
        async def event_generator():
            stream_started_at = time.perf_counter()
            first_event_logged = False
            first_text_delta_logged = False
            last_text_delta_at = None
            text_delta_count = 0
            output_chunks = []
            form_update_payload = {}
            text_chunker = (
                SentenceTextChunker(
                    SentenceTextChunkerConfig(
                        min_chars=1,
                        max_chars=0,
                        use_soft_breaks=False,
                    )
                )
                if use_text_chunk
                else None
            )
            try:
                logger.info(
                    "Chat backend stream opened "
                    f"sessionId={json_data['sessionId']} "
                    f"open_latency_ms={(time.perf_counter() - stream_started_at) * 1000:.1f}"
                )
                
                buffer = ""
                state_filter_buffer = ""
                state_code_found = False
                module_form_sent = False

                def flush_text_delta(text: str):
                    nonlocal first_text_delta_logged, last_text_delta_at, text_delta_count
                    now = time.perf_counter()
                    if not first_text_delta_logged:
                        first_text_delta_logged = True
                        logger.info(
                            "Chat stream first text_delta sent "
                            f"sessionId={json_data['sessionId']} "
                            f"text_delta_ttfb_ms={(now - stream_started_at) * 1000:.1f}"
                        )
                    last_text_delta_at = now
                    text_delta_count += 1
                    output_chunks.append(text)
                    return create_sse_event("text_delta", {"text": text})

                def flush_form_update(form_update):
                    return create_sse_event("formUpdate", form_update)

                def build_text_delta_events(text: str):
                    if not text:
                        return []
                    chunks = text_chunker.feed(text) if text_chunker else [text]
                    return [flush_text_delta(chunk) for chunk in chunks if chunk]

                def flush_text_chunker_events():
                    if not text_chunker:
                        return []
                    chunk = text_chunker.flush()
                    if not chunk:
                        return []
                    return [flush_text_delta(chunk)]

                def clean_later_state_tags(text: str, *, final: bool = False) -> str:
                    nonlocal state_filter_buffer
                    state_filter_buffer += text
                    state_filter_buffer = STATE_TAG_PATTERN.sub("", state_filter_buffer)

                    state_start_index = state_filter_buffer.find("<state")
                    if state_start_index >= 0:
                        cleaned = state_filter_buffer[:state_start_index]
                        state_filter_buffer = state_filter_buffer[state_start_index:]
                        return cleaned

                    if final:
                        if state_filter_buffer.startswith("<state"):
                            state_filter_buffer = ""
                            return ""
                        cleaned = state_filter_buffer
                        state_filter_buffer = ""
                        return cleaned

                    hold_length = 0
                    state_prefix = "<state>"
                    max_suffix_length = min(len(state_filter_buffer), len(state_prefix) - 1)
                    for suffix_length in range(1, max_suffix_length + 1):
                        if state_prefix.startswith(state_filter_buffer[-suffix_length:]):
                            hold_length = suffix_length

                    if hold_length:
                        cleaned = state_filter_buffer[:-hold_length]
                        state_filter_buffer = state_filter_buffer[-hold_length:]
                        return cleaned

                    cleaned = state_filter_buffer
                    state_filter_buffer = ""
                    return cleaned
                
                async for event in backend.stream(chat_input):
                    try:
                        if not first_event_logged:
                            first_event_logged = True
                            logger.info(
                                "Chat backend stream first event "
                                f"sessionId={json_data['sessionId']} "
                                f"kind={type(event).__name__} "
                                f"ttfb_ms={(time.perf_counter() - stream_started_at) * 1000:.1f}"
                            )

                        if isinstance(event, FormUpdate) and not module_form_sent:
                            if event.data:
                                form_update_payload = event.data
                                logger.info(
                                    "Chat backend stream formUpdate received "
                                    f"sessionId={json_data['sessionId']} "
                                    f"type={type(event.data).__name__} "
                                    f"formUpdate={event.data!r}"
                                )
                                yield flush_form_update(event.data)
                                module_form_sent = True
                            continue

                        if not isinstance(event, TextDelta):
                            continue

                        delta_content = event.text
                        if not delta_content:
                            continue

                        buffer += delta_content

                        if not state_code_found:
                            # Check for <state>XXXX</state> pattern
                            state_code, cleaned_content = extract_first_state_and_clean_content(buffer)
                            if state_code:

                                # Apply logic to map/adjust state code
                                nextStageCode = normalize_stage_code(state_code)
                                nextStage = STATUS_CODE_MAP.get(nextStageCode, '')
                                logger.info(
                                    "FastGPT stream stage_code parsed "
                                    f"sessionId={json_data['sessionId']} "
                                    f"nextStageCode={nextStageCode} nextStage={nextStage}"
                                )

                                # Send stage code event
                                yield create_sse_event("stage_code", {
                                    "nextStageCode": nextStageCode,
                                    "nextStage": nextStage
                                })

                                state_code_found = True

                                if cleaned_content:
                                    for text_event in build_text_delta_events(cleaned_content):
                                        yield text_event
                                buffer = "" # Clear buffer after extracting state
                        else:
                            cleaned_content = clean_later_state_tags(delta_content)
                            if cleaned_content:
                                for text_event in build_text_delta_events(cleaned_content):
                                    yield text_event
                            buffer = ""

                    except Exception as e:
                        logger.error(f"Error processing stream event: {e}")
                        continue

                # If stream ends and no state code found (unlikely if format is strict), 
                # we might want to send what we have
                if not state_code_found and buffer:
                    for text_event in build_text_delta_events(buffer):
                        yield text_event
                elif state_code_found:
                    cleaned_content = clean_later_state_tags("", final=True)
                    if cleaned_content:
                        for text_event in build_text_delta_events(cleaned_content):
                            yield text_event

                for text_event in flush_text_chunker_events():
                    yield text_event

                text_delta_end_ms = (
                    f"{(last_text_delta_at - stream_started_at) * 1000:.1f}"
                    if last_text_delta_at is not None
                    else "-"
                )
                logger.info(
                    "Chat stream completed "
                    f"sessionId={json_data['sessionId']} "
                    f"duration_ms={(time.perf_counter() - stream_started_at) * 1000:.1f} "
                    f"text_delta_end_ms={text_delta_end_ms} "
                    f"text_delta_count={text_delta_count} "
                    f"useTextChunk={use_text_chunk} "
                    f"stage_code_found={state_code_found} formUpdate_sent={module_form_sent} "
                    f"output={''.join(output_chunks)!r} "
                    f"formUpdate={form_update_payload!r}"
                )
                yield create_sse_event("done", {"status": "completed"})
                
            except Exception as e:
                logger.error(
                    "Chat stream failed "
                    f"sessionId={json_data['sessionId']} "
                    f"duration_ms={(time.perf_counter() - stream_started_at) * 1000:.1f} "
                    f"error={e}"
                )
                yield create_sse_event("error", {"msg": str(e), "code": "500"})

        return StreamingResponse(event_generator(), media_type="text/event-stream")

    try:
        result = await backend.complete(chat_input)
        logger.info(
            "Chat backend non-stream response received "
            f"sessionId={json_data['sessionId']} "
            f"latency_ms={(time.perf_counter() - request_started_at) * 1000:.1f}"
        )
        
    except ChatBackendAuthenticationError as e:
        logger.error(f"Authentication error: {e}")
        return ProcessResponse_chat(
            sessionId=json_data['sessionId'],
            timeStamp=json_data['timeStamp'],
            outputText="",
            nextStage="",
            nextStageCode="",
            code="401",
            msg="认证失败"
        )
    except ChatBackendRateLimitError as e:
        logger.error(f"Rate limit error: {e}")
        return ProcessResponse_chat(
            sessionId=json_data['sessionId'],
            timeStamp=json_data['timeStamp'],
            outputText="",
            nextStage="",
            nextStageCode="",
            code="429",
            msg="请求过于频繁，请稍后重试"
        )
    except ChatBackendAPIError as e:
        logger.error(f"API error: {e}")
        return ProcessResponse_chat(
            sessionId=json_data['sessionId'],
            timeStamp=json_data['timeStamp'],
            outputText="",
            nextStage="",
            nextStageCode="",
            code="500",
            msg="大模型服务器无响应"
        )
    except Exception as e:
        logger.error(f"Unexpected error: {e}")
        return ProcessResponse_chat(
            sessionId=json_data['sessionId'],
            timeStamp=json_data['timeStamp'],
            outputText="",
            nextStage="",
            nextStageCode="",
            code="500",
            msg="大模型服务器无响应"
        )
    
    try:
        content = result.content
        logger.info(f"Chat backend returned content: {content}")

        content_stage_code = None
        if isinstance(content, str):
            logger.debug("content是一个str")
            content_stage_code, content = extract_first_state_and_clean_content(content)
            if content_stage_code:
                logger.debug(
                    f"解析后的第一个state为: {content_stage_code}, "
                    f"移除state标签后的content为: {content}"
                )
            else:
                raise ValueError("大模型回复中的state解析失败")
        else:
            logger.error(f"content既不是list也不是str, type: {type(content)}")
            raise ValueError("大模型回复不是list也不是str")

        nextStageCode = content_stage_code or result.status_code
        if not nextStageCode:
            raise ValueError("大模型回复中缺少state")
        nextStageCode = normalize_stage_code(nextStageCode)
        nextStage = STATUS_CODE_MAP.get(nextStageCode, '')
        form_update = result.form_update
        logger.info(
            "Chat non-stream completed "
            f"sessionId={json_data['sessionId']} "
            f"duration_ms={(time.perf_counter() - request_started_at) * 1000:.1f} "
            f"nextStageCode={nextStageCode} nextStage={nextStage} "
            f"output_len={len(content)} formUpdate_type={type(form_update).__name__} "
            f"output={content!r} formUpdate={form_update!r}"
        )

        return ProcessResponse_chat(
            sessionId=json_data['sessionId'],
            timeStamp=json_data['timeStamp'],
            outputText=content,
            formUpdate=form_update,
            nextStage=nextStage,
            nextStageCode=nextStageCode,
            code="200",
            msg="",
        )
    except Exception as e:
        logger.error(f"解析信息发生错误: {e}")
        logger.error(f"content: {content}, type: {type(content)}")
        return ProcessResponse_chat(
            sessionId=json_data['sessionId'],
            timeStamp=json_data['timeStamp'],
            outputText="",
            nextStage="",
            nextStageCode="",
            code="500",
            msg="大模型服务返回消息不完整"
        )


@router.post("/set_info", response_model=ProcessResponse_set)
async def set_info(
    request: ProcessRequest_set,
    client: AsyncChatClient = Depends(get_fastgpt_client)
):
    """Set information in chat state."""
    json_data = request.model_dump()
    set_info_payload = {'key': json_data['key'], 'value': json_data['value']}
    set_info_input = format_set_info_input(
        set_info_payload,
        json_data.get('includeInputInfo', False)
    )
    
    try:
        # Get current state
        response = await client.create_chat_completion(
            messages=[{"role": "user", "content": set_info_input}],
            chatId=json_data['sessionId'],
            stream=False,
            detail=True
        )
        response.raise_for_status()
        data = response.json()
        
        current_state = data['newVariables']['state']
        if isinstance(current_state, str):
            current_state = json.loads(current_state)
        logger.debug(f"Current state: {current_state}")
            
    except Exception as e:
        logger.error(f"Error getting current state: {e}")
        return ProcessResponse_set(
            sessionId=json_data['sessionId'],
            timeStamp=json_data['timeStamp'],
            code="500",
            msg="大模型服务器无响应"
        )
    
    try:
        await delete_last_two_chat_records(client, json_data['sessionId'])
    except Exception as e:
        logger.error(f"Error deleting chat records: {e}")
        return ProcessResponse_set(
            sessionId=json_data['sessionId'],
            timeStamp=json_data['timeStamp'],
            code="500",
            msg="大模型后台无响应"
        )
    
    try:
        # Update state
        key = json_data['key']
        value = json_data['value']
        current_state[key] = value
        logger.info(f'即将设置 {key} 为 {value}')
        logger.info(f'即将上传 {current_state}')
        
        # Update state using SDK
        response = await client.create_chat_completion(
            messages=[{"role": "user", "content": set_info_input}],
            chatId=json_data['sessionId'],
            stream=False,
            detail=True,
            variables={'state': current_state}
        )
        response.raise_for_status()
        
        # Delete records again after update
        await delete_last_two_chat_records(client, json_data['sessionId'])
        
        return ProcessResponse_set(
            sessionId=json_data['sessionId'],
            timeStamp=json_data['timeStamp'],
            code="200",
            msg=""
        )
        
    except Exception as e:
        logger.error(f"Error setting info: {e}")
        return ProcessResponse_set(
            sessionId=json_data['sessionId'],
            timeStamp=json_data['timeStamp'],
            code="500",
            msg="大模型后台无响应"
        )
    
@router.post("/get_info", response_model=ProcessResponse_get)
async def get_info(
    request: ProcessRequest_get,
    client: AsyncChatClient = Depends(get_fastgpt_client)
):
    """Get information from chat state."""
    json_data = request.model_dump()
    get_info_payload = {'key': json_data['key']}
    get_info_input = format_set_info_input(
        get_info_payload,
        json_data.get('includeInputInfo', False)
    )
    
    try:
        # Get current state
        response = await client.create_chat_completion(
            messages=[{"role": "user", "content": get_info_input}],
            chatId=json_data['sessionId'],
            stream=False,
            detail=True
        )
        response.raise_for_status()
        data = response.json()
        
        current_state = data['newVariables']['state']
        if isinstance(current_state, str):
            current_state = json.loads(current_state)
        logger.debug(f"Current state: {current_state}")
            
    except Exception as e:
        logger.error(f"Error getting state: {e}")
        return ProcessResponse_get(
            sessionId=json_data['sessionId'],
            timeStamp=json_data['timeStamp'],
            value="",
            code="500",
            msg="大模型服务器无响应"
        )
    
    try:
        await delete_last_two_chat_records(client, json_data['sessionId'])
    except Exception as e:
        logger.error(f"Error deleting records: {e}")
        return ProcessResponse_get(
            sessionId=json_data['sessionId'],
            timeStamp=json_data['timeStamp'],
            value="",
            code="500",
            msg="大模型后台无响应"
        )
    
    try:
        key = json_data['key']
        acd_keys = ['ywrysw', 'ywfjdc', 'ywmtc', 'bjrjs', 'sgfssj', 'sfsgxc', 'jdcsl', 'sgyy']
        human1_keys = ['xm1', 'hpzl1', 'hphm1', 'sfzmhm1', 'sfzmwh1', 'sjhm1', 'sjwh1', 'csbw1']
        human2_keys = ['xm2', 'hpzl2', 'hphm2', 'sfzmhm2', 'sfzmwh2', 'sjhm2', 'sjwh2', 'csbw2']
        
        def bool_to_str(v):
            """Convert boolean to string representation."""
            return '1' if v is True else '0' if v is False else v
        
        if key == 'all':
            acd_values = {k: bool_to_str(current_state.get(k, '')) for k in acd_keys}
            human1_values = {k: bool_to_str(current_state.get(k, '')) for k in human1_keys}
            human2_values = {k: bool_to_str(current_state.get(k, '')) for k in human2_keys}
            value = json.dumps({
                'acdinfo': acd_values,
                'acdhuman1': human1_values,
                'acdhuman2': human2_values
            })
        elif key == "acdinfo":
            acd_values = {k: bool_to_str(current_state.get(k, '')) for k in acd_keys}
            value = json.dumps(acd_values)
        elif key == 'acdhuman1':
            human1_values = {k: bool_to_str(current_state.get(k, '')) for k in human1_keys}
            value = json.dumps(human1_values)
        elif key == 'acdhuman2':
            human2_values = {k: bool_to_str(current_state.get(k, '')) for k in human2_keys}
            value = json.dumps(human2_values)
        else:
            value = json.dumps(current_state.get(key, ''))
        
        logger.debug(f"Returning value for key '{key}': {value}")
        return ProcessResponse_get(
            sessionId=json_data['sessionId'],
            timeStamp=json_data['timeStamp'],
            value=value,
            code="200",
            msg=""
        )
        
    except Exception as e:
        logger.error(f"Error getting info: {e}")
        return ProcessResponse_get(
            sessionId=json_data['sessionId'],
            timeStamp=json_data['timeStamp'],
            value="",
            code="500",
            msg="入参不规范"
        )
    

@router.delete("/delete_session", response_model=ProcessResponse_delete_session)
async def delete_session(
    request: ProcessRequest_delete_session,
    client: AsyncChatClient = Depends(get_fastgpt_client)
):
    """Delete a chat session."""
    json_data = request.model_dump()
    chat_id = json_data.get('sessionId')
    
    if not chat_id:
        return ProcessResponse_delete_session(
            sessionId="",
            timeStamp=json_data['timeStamp'],
            code="400",
            msg="sessionId is required"
        )
    
    try:
        # Use SDK's delete_chat_history
        response = await client.delete_chat_history(
            appId=Config.FASTGPT_APP_ID,
            chatId=chat_id
        )
        response.raise_for_status()
        data = response.json()
        
        if data.get('code') == 200:
            return ProcessResponse_delete_session(
                sessionId=json_data['sessionId'],
                timeStamp=json_data['timeStamp'],
                code="200",
                msg=""
            )
        else:
            raise ValueError("删除会话失败")
            
    except Exception as e:
        logger.error(f"Error deleting session: {e}")
        return ProcessResponse_delete_session(
            sessionId=json_data['sessionId'],
            timeStamp=json_data['timeStamp'],
            code="500",
            msg="删除会话失败"
        )