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, aiter_stream_events
from fastgpt_client.exceptions import (
    APIError, AuthenticationError, RateLimitError, ValidationError
)
from ..core.fastgpt_client import get_fastgpt_client
from ..core.config import Config
from loguru import logger
import json
import re
import time

router = APIRouter()
FORM_EXTRACT_MODULE_NAME = "文本内容提取事故信息"
STATUS_CODE_MAP = {
    '0000': '结束通话',
    '0001': '转接人工',
    '0002': '语义无法识别转接人工',
    '0003': '有人伤转接人工',
    '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_state_and_content(data1: str) -> dict | None:
    """
    Extracts the state and content from a string in the format <state>STATE</state>content.

    Args:
        data1: The input string.

    Returns:
        A dictionary with 'state' and 'content' keys if a match is found,
        otherwise None.
    """
    data1 = data1.strip()
    regex = r"<state>(.*?)</state>(.*)"
    match = re.search(regex, data1, flags=re.DOTALL)
    
    if match:
        return {
            "state": match.group(1),
            "content": match.group(2),
        }
    return None


def parse_json_value(value):
    """Parse JSON string values when possible."""
    parsed = value
    for _ in range(3):
        if not isinstance(parsed, str):
            return parsed
        parsed = parsed.strip()
        if not parsed:
            return {}
        try:
            parsed = json.loads(parsed)
        except json.JSONDecodeError:
            return parsed
    return parsed


def extract_form_update_from_flow_nodes(nodes):
    """Extract form update data from the configured FastGPT content-extract node."""
    if not isinstance(nodes, list):
        return {}

    for node in nodes:
        if not isinstance(node, dict):
            continue
        if node.get("moduleName") != FORM_EXTRACT_MODULE_NAME:
            continue

        extract_result = node.get("extractResult", {})
        if not isinstance(extract_result, dict):
            return {}

        form_update = extract_result.get("formUpdate") or extract_result.get("form") or ""
        if not form_update:
            return {}
        return parse_json_value(form_update)

    return {}


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,
    client: AsyncChatClient = Depends(get_fastgpt_client)
):
    """Handle chat completion request."""
    json_data = request.model_dump()
    need_form_update = json_data.get('needFormUpdate', False)
    chat_variables = {'needFormUpdate': 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} 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 = {}
            try:
                # Use SDK's create_chat_completion with stream=True
                response = await client.create_chat_completion(
                    messages=[{"role": "user", "content": json_data['text']}],
                    chatId=json_data['sessionId'],
                    stream=True,
                    detail=True,
                    variables=chat_variables
                )
                logger.info(
                    "FastGPT stream response opened "
                    f"sessionId={json_data['sessionId']} "
                    f"open_latency_ms={(time.perf_counter() - stream_started_at) * 1000:.1f}"
                )
                
                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)
                
                async for event in aiter_stream_events(response):
                    try:
                        if not first_event_logged:
                            first_event_logged = True
                            logger.info(
                                "FastGPT stream first event "
                                f"sessionId={json_data['sessionId']} kind={event.kind} "
                                f"ttfb_ms={(time.perf_counter() - stream_started_at) * 1000:.1f}"
                            )

                        if event.kind == "flowResponses" and not module_form_sent:
                            form_update = extract_form_update_from_flow_nodes(event.data)
                            if form_update:
                                form_update_payload = form_update
                                logger.info(
                                    "FastGPT stream formUpdate extracted "
                                    f"sessionId={json_data['sessionId']} "
                                    f"type={type(form_update).__name__} "
                                    f"formUpdate={form_update!r}"
                                )
                                yield flush_form_update(form_update)
                                module_form_sent = True
                            continue

                        if event.kind not in {"answer", "fastAnswer", "data"}:
                            continue

                        data = event.data
                        if not isinstance(data, dict):
                            continue

                        try:
                            delta_content = data['choices'][0]['delta'].get('content', '')
                        except (KeyError, IndexError):
                            delta_content = ''
                        if not delta_content:
                            continue

                        buffer += delta_content

                        if not state_code_found:
                            # Check for <state>XXXX</state> pattern
                            match = re.search(r"<state>(.*?)</state>", buffer, flags=re.DOTALL)
                            if match:
                                state_code = match.group(1)

                                # 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

                                # Send remaining content as text_delta
                                remaining_content = buffer[match.end():]
                                if remaining_content:
                                    yield flush_text_delta(remaining_content)
                                buffer = "" # Clear buffer after extracting state
                        else:
                            yield flush_text_delta(delta_content)
                            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:
                    yield flush_text_delta(buffer)

                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"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:
        # Use SDK's create_chat_completion
        response = await client.create_chat_completion(
            messages=[{"role": "user", "content": json_data['text']}],
            chatId=json_data['sessionId'],
            stream=False,
            detail=True,
            variables=chat_variables
        )
        response.raise_for_status()
        data = response.json()
        logger.info(
            "FastGPT non-stream response received "
            f"sessionId={json_data['sessionId']} "
            f"latency_ms={(time.perf_counter() - request_started_at) * 1000:.1f}"
        )
        
    except AuthenticationError 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 RateLimitError 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 APIError 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:
        # Extract content from FastGPT response
        content = data['choices'][0]['message']['content']
        logger.info(f"FastGPT服务返回信息content: {content}")

        finish_reason = data['choices'][0]['finish_reason']

        # Extract state variables
        state = data.get('newVariables', {}).get('state', {})
        if isinstance(state, str):
            state = json.loads(state)

        transfer_to_human = state.get("transfer_to_human", False)
        ywrysw = state.get("ywrysw", False)
        ywfjdc = state.get("ywfjdc", False)
        ywmtc = state.get("ywmtc", False)
        jdcsl = state.get("jdcsl", 0)
        accident_info_complete = state.get("accident_info_complete", False)
        user_is_ready = state.get("user_is_ready", False)
        if isinstance(user_is_ready, str):
            user_is_ready = user_is_ready.lower() == 'true'
        driver_info_complete = state.get("driver_info_complete", False)
        drivers_info_complete = state.get("drivers_info_complete", False)
        driver_info_check = state.get("drivers_info_check", False)
        drivers_info_check = state.get("drivers_info_check", False)

        logger.debug(f"State variables: {data.get('newVariables', {})}")

        # Parse content - sometimes content is a string, sometimes it is a list
        content_stage_code = None
        if isinstance(content, list):
            logger.debug("content是一个list")
            content = content[0]['text']['content']
        
        if isinstance(content, str):
            logger.debug("content是一个str")
            state_and_content = extract_state_and_content(content)
            if state_and_content:
                logger.debug(f"解析后的state和content为: {state_and_content}")
                content_stage_code = state_and_content['state']
                content = state_and_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 data['newVariables']['status_code']
        nextStageCode = normalize_stage_code(nextStageCode)
        nextStage = STATUS_CODE_MAP.get(nextStageCode, '')
        form_update = extract_form_update_from_flow_nodes(data.get("responseData", []))
        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="删除会话失败"
        )