from fastapi import APIRouter, Response
from src.models.analysis_models import InputModel
from src.workflows.analysis_workflow import MLAnalysisWorkflow, MLImplementationPlanner
from datetime import datetime
from phi.storage.workflow.sqlite import SqlWorkflowStorage
from phi.utils.pprint import pprint_run_response
from phi.utils.log import logger
from typing import Iterator

router = APIRouter()

@router.get("/")
async def read_root():
    return Response("Ml-Analysis workflow from user problem is Up!")

@router.post("/analyze-problem")
async def analyze_problem(data: InputModel):

    analysis_workflow = MLAnalysisWorkflow(
        session_id=f"ml-analysis-{datetime.now().strftime('%Y%m%d_%H%M%S')}",
        storage=SqlWorkflowStorage(
            table_name="ml_analysis_workflows",
            db_file="storage/workflows.db"
        )
    )

    analysis_response: Iterator[RunResponse] = analysis_workflow.run(data.problem_statement)
    
    pprint_run_response(analysis_response, markdown=True)

    requirements_result = analysis_workflow.requirements_analyst.run_response.content if analysis_workflow.requirements_analyst.run_response else None
    research_result = analysis_workflow.technical_researcher.run_response.content if analysis_workflow.technical_researcher.run_response else None

    if requirements_result:
        logger.info("===Planning Phase===")
        planning_workflow = MLImplementationPlanner(
            session_id=f"ml-planning-{datetime.now().strftime('%Y%m%d_%H%M%S')}",
            storage=SqlWorkflowStorage(
                table_name="ml_planning_workflows",
                db_file="storage/workflows.db"
            )
        )

        # run and print planning workflow
        planning_response_stream: Iterator[RunResponse] = planning_workflow.run(requirements_result, research_result)
        
        pprint_run_response(planning_response_stream, markdown=True)

        return {"Response": planning_workflow.writer.run_response.content}

    else:
        return {"Error": "Requirements analysis did not complete successfully."}