# Created by aaronkueh on 9/26/2025
# aom/aom_api.py
from __future__ import annotations

from typing import Any, Dict, TypedDict, Optional, Union
from fastapi import FastAPI, HTTPException, Request
from fastapi.middleware.cors import CORSMiddleware
from pydantic import BaseModel, Field
from aom.utils.logger import init_logging, get_logger
from aom.utils.report_helpers import performance_pdf
from aom.agents.orchestrator import Orchestrator
from aom.utils.agent_helpers import ui_sanitize
from datetime import datetime


class OrchestratorState(TypedDict, total=False):
    user_prompt: str
    session_id: str  # Added for Qdrant storage
    intent: str
    isa_params: Dict[str, Any]
    result: Dict[str, Any]
    error: str

class OrchestrateRequest(BaseModel):
    """
    Request body for orchestration with free-form prompt only.
    """
    prompt: str = Field(..., description='Free-form prompt text')
    mode: Optional[str] = 'general'
    session_id: Optional[str] = None

class OrchestrateResponse(BaseModel):
    """
    Response envelope for orchestration result.
    """
    result: Union[str, Dict[str, Any]]
    session_id: Optional[str] = None

class ReportRequest(BaseModel):
    report: str = Field(..., description='Report text content to be converted into PDF')


# App initialization
init_logging(app_name='aom-api')
log = get_logger('api_server')

app = FastAPI(title='AOM Orchestrator API', version='1.0.0', docs_url='/docs', redoc_url='/redoc')

# CORS (adjust origins as needed)
app.add_middleware(
    CORSMiddleware,
    allow_origins=['*'],  # Narrow this in production
    allow_credentials=True,
    allow_methods=['*'],
    allow_headers=['*'],
)

# Create orchestrator (uses default SQL profile)
_orch = Orchestrator(db_profile=None)

@app.get('/health')
def health() -> Dict[str, Any]:
    """Health check endpoint"""
    return {
        'status': 'ok',
        'checkpoint_enabled': _orch.checkpointer is not None
    }


async def _default_orchestrate(req: OrchestrateRequest, request: Request) -> OrchestrateResponse:
    """
    The original orchestrator logic, extracted into a helper.
    """
    try:
        # Extract mode and session_id from request
        mode = req.mode if hasattr(req, 'mode') and req.mode else 'general'
        sim_session = f'sess_0000000000000_testing_01'
        session_id = req.session_id if hasattr(req, 'session_id') and req.session_id else sim_session

        # Run orchestrator in a thread pool to avoid blocking event loop
        # LangGraph's sync checkpointer operations can block FastAPI's async loop
        import asyncio
        from concurrent.futures import ThreadPoolExecutor

        loop = asyncio.get_event_loop()
        with ThreadPoolExecutor() as pool:
            response_data = await loop.run_in_executor(
                pool,
                lambda: _orch.run(req.prompt, mode=mode, session_id=session_id)
            )

        result = response_data.get('result', {})
        print('Response received: \n', result)

        if isinstance(result, dict):
            # if result.get("summary") == 'table':
            if 'table' in (result.keys()):
                table = result.get('table')
                if table is None:
                    summary = result.get('summary', str(result))
                else:
                    summary = {'table': table}
            elif 'report' in result.keys():
                report = result.get('report', str(result)).replace('<END>', "")
                file_link = None
                if report:
                    try:
                        timestamp = datetime.now().strftime("%Y%m%d_%H%M%S")
                        filename = f'report_{timestamp}.pdf'
                        file_link = performance_pdf(report, file_name=filename)
                        log.info(f'Report PDF generated at {file_link}')
                    except Exception as e:
                        log.exception('Failed to write report PDF: %s', e)

                # Return the normal summary to the UI
                s = result.get('summary', str(result))
                s = ui_sanitize(s, strip_end_token=True)

                if file_link is not None:
                    summary = {
                        'summary': s,
                        'report_path': str(file_link)
                    }
                else:
                    summary = s
            else:
                summary = result.get('summary', str(result))
                summary = ui_sanitize(summary, strip_end_token=True)
        else:
            summary = str(result)
            summary = ui_sanitize(summary, strip_end_token=True)

        returned_session_id = response_data.get('session_id')

        return OrchestrateResponse(result=summary, session_id=returned_session_id)
    except Exception as e:
        log.exception('orchestrate failed')
        raise HTTPException(status_code=500, detail=str(e))


@app.post('/orchestrate', response_model=OrchestrateResponse)
async def orchestrate(
    req: OrchestrateRequest,
    request: Request,
):
    return await _default_orchestrate(req, request)


if __name__ == '__main__':
    # Local run helper: uvicorn aom.api_server:app --reload
    import uvicorn

    uvicorn.run('aom.aom_api:app', host='127.0.0.1', port=8080, reload=False)