"""Transformation Agent 路由模組"""

import json
import logging
from typing import AsyncGenerator, List

from fastapi import APIRouter
from fastapi.responses import StreamingResponse

from ..models.schemas import (
    TransformationRequest,
    TransformationCategoryResult,
    TransformationDescription,
)
from ..prompts.transformation_prompt import (
    get_keyword_generation_prompt,
    get_batch_description_prompt,
)
from ..services.llm_service import ollama_provider, extract_json_from_response

logger = logging.getLogger(__name__)
router = APIRouter(prefix="/api/transformation", tags=["transformation"])


async def generate_transformation_events(
    request: TransformationRequest
) -> AsyncGenerator[str, None]:
    """Generate SSE events for transformation process"""
    try:
        temperature = request.temperature if request.temperature is not None else 0.7
        model = request.model

        # ========== Step 1: Generate new keywords ==========
        yield f"event: keyword_start\ndata: {json.dumps({'message': f'為「{request.category}」生成新關鍵字...'}, ensure_ascii=False)}\n\n"

        keyword_prompt = get_keyword_generation_prompt(
            category=request.category,
            attributes=request.attributes,
            keyword_count=request.keyword_count
        )
        logger.info(f"Keyword prompt: {keyword_prompt[:200]}")

        keyword_response = await ollama_provider.generate(
            keyword_prompt, model=model, temperature=temperature
        )
        logger.info(f"Keyword response: {keyword_response[:500]}")

        keyword_data = extract_json_from_response(keyword_response)
        new_keywords = keyword_data.get("keywords", [])

        yield f"event: keyword_complete\ndata: {json.dumps({'keywords': new_keywords}, ensure_ascii=False)}\n\n"

        if not new_keywords:
            yield f"event: error\ndata: {json.dumps({'error': '無法生成新關鍵字'}, ensure_ascii=False)}\n\n"
            return

        # ========== Step 2: Generate descriptions for each keyword ==========
        yield f"event: description_start\ndata: {json.dumps({'message': '生成創新應用描述...'}, ensure_ascii=False)}\n\n"

        # Use batch description prompt for efficiency
        desc_prompt = get_batch_description_prompt(
            query=request.query,
            category=request.category,
            keywords=new_keywords
        )
        logger.info(f"Description prompt: {desc_prompt[:300]}")

        desc_response = await ollama_provider.generate(
            desc_prompt, model=model, temperature=temperature
        )
        logger.info(f"Description response: {desc_response[:500]}")

        desc_data = extract_json_from_response(desc_response)
        descriptions_raw = desc_data.get("descriptions", [])

        # Convert to TransformationDescription objects
        descriptions: List[TransformationDescription] = []
        for desc in descriptions_raw:
            if isinstance(desc, dict) and "keyword" in desc and "description" in desc:
                descriptions.append(TransformationDescription(
                    keyword=desc["keyword"],
                    description=desc["description"]
                ))

        yield f"event: description_complete\ndata: {json.dumps({'count': len(descriptions)}, ensure_ascii=False)}\n\n"

        # ========== Build final result ==========
        result = TransformationCategoryResult(
            category=request.category,
            original_attributes=request.attributes,
            new_keywords=new_keywords,
            descriptions=descriptions
        )

        final_data = {
            "result": result.model_dump()
        }
        yield f"event: done\ndata: {json.dumps(final_data, ensure_ascii=False)}\n\n"

    except Exception as e:
        logger.error(f"Transformation error: {e}")
        yield f"event: error\ndata: {json.dumps({'error': str(e)}, ensure_ascii=False)}\n\n"


@router.post("/category")
async def transform_category(request: TransformationRequest):
    """處理單一類別的轉換"""
    return StreamingResponse(
        generate_transformation_events(request),
        media_type="text/event-stream",
        headers={
            "Cache-Control": "no-cache",
            "Connection": "keep-alive",
            "X-Accel-Buffering": "no",
        },
    )