from logging import Logger
import os
from datetime import datetime
from typing import Dict, List, Optional
from dotenv import load_dotenv
from langchain_ollama import OllamaLLM
from langgraph.graph import END, StateGraph
from models import AgentState, DiagramType
from config import Config
from langchain_core.prompts import ChatPromptTemplate
from file_manager import FileManager
from agents.analysis import (
    EntitiesAnalysis,
    ActorsAnalysis,
    SequenceAnalysis,
    StateAnalysis,
    ArchitectureAnalysis,
    DeploymentAnalysis
)
from agents.code import (
    ClassDiagramGenerator,
    UseCaseDiagramGenerator,
    SequenceDiagramGenerator,
    ActivityDiagramGenerator,
    ComponentDiagramGenerator,
    DeploymentDiagramGenerator,
    StateDiagramGenerator,
    TimingDiagramGenerator,
    ObjectDiagramGenerator
)
from fastapi import FastAPI, HTTPException
from fastapi.responses import StreamingResponse
from fastapi.middleware.cors import CORSMiddleware
from pydantic import BaseModel
import asyncio
from langchain_core.callbacks import StreamingStdOutCallbackHandler
import json

# Load environment variables
load_dotenv()

# Initialize FastAPI app
app = FastAPI()

# Add CORS middleware configuration
app.add_middleware(
    CORSMiddleware,
    allow_origins=["*"],
    allow_credentials=True,
    allow_methods=["*"],
    allow_headers=["*"],
)


# Request model
class DiagramRequest(BaseModel):
    project_name: str
    project_description: str
    selected_diagrams: List[int]
    output_dir: Optional[str] = "output"




@app.post("/generate-diagrams")
async def generate_diagrams(request: DiagramRequest):
    try:
        initial_state = AgentState(
            project_name=request.project_name,
            project_description=request.project_description,
            output_dir=request.output_dir,
            selected_diagrams=request.selected_diagrams,
            entities_classes="",
            actors_use_cases="",
            sequence_interactions="",
            class_diagram="",
            use_case_diagram="",
            sequence_diagram="",
            activity_diagram="",
            component_diagram="",
            deployment_diagram="",
            state_diagram="",
            timing_diagram="",
            object_diagram=""
        )

    
        diagram_requirements = {
            DiagramType.CLASS: ["extract_entities"],
            DiagramType.USE_CASE: ["extract_actors"],
            DiagramType.SEQUENCE: ["extract_actors", "extract_sequence"],
            DiagramType.OBJECT: ["extract_entities"],
            DiagramType.ACTIVITY: ["extract_actors"],
            DiagramType.COMPONENT: ["extract_entities", "extract_architecture"],
            DiagramType.DEPLOYMENT: ["extract_architecture", "extract_deployment"],
            DiagramType.STATE: ["extract_entities", "extract_sequence", "extract_states"],
            DiagramType.TIMING: ["extract_sequence"]
        }

       
        analysis_functions = {
            "extract_entities": EntitiesAnalysis(),
            "extract_actors": ActorsAnalysis(),
            "extract_sequence": SequenceAnalysis(),
            "extract_states": StateAnalysis(),
            "extract_architecture": ArchitectureAnalysis(),
            "extract_deployment": DeploymentAnalysis()
        }

        generation_functions = {
            DiagramType.CLASS: ("generate_class", ClassDiagramGenerator()),
            DiagramType.USE_CASE: ("generate_use_case", UseCaseDiagramGenerator()),
            DiagramType.SEQUENCE: ("generate_sequence", SequenceDiagramGenerator()),
            DiagramType.ACTIVITY: ("generate_activity", ActivityDiagramGenerator()),
            DiagramType.COMPONENT: ("generate_component", ComponentDiagramGenerator()),
            DiagramType.DEPLOYMENT: ("generate_deployment", DeploymentDiagramGenerator()),
            DiagramType.STATE: ("generate_state", StateDiagramGenerator()),
            DiagramType.TIMING: ("generate_timing", TimingDiagramGenerator()),
            DiagramType.OBJECT: ("generate_object", ObjectDiagramGenerator())
        }

      
        selected_types = [DiagramType(d) for d in request.selected_diagrams]
        added_analysis_steps = set()
        steps = []

        for diagram_type in selected_types:
            required_analysis = diagram_requirements.get(diagram_type, [])
            for step in required_analysis:
                if step not in added_analysis_steps:
                    steps.append((step, analysis_functions[step]))
                    added_analysis_steps.add(step)

            if diagram_type in generation_functions:
                steps.append(generation_functions[diagram_type])

       
        async def event_stream():
            for step_name, step_func in steps:
                # Determine if it's an analysis step or generation step
                step_type = "analysis" if step_name in analysis_functions else "generation"
                # Determine file extension based on step type
                file_ext = "pu" if step_type == "generation" else "md"
                yield f'data: {{"type": "step", "step_type": "{step_type}", "file_name": "{step_name}.{file_ext}"}}\n\n'
                async for chunk in step_func(initial_state):
                    yield f"data: {chunk}\n\n"

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

    except Exception as e:
        raise HTTPException(status_code=500, detail=str(e))