import json
from fastapi import FastAPI, HTTPException
from pydantic import BaseModel
import openai
import asyncio
import os
from typing import List

app = FastAPI()

# Pydantic model for input
class StudyInput(BaseModel):
    overall_study_pattern: str
    memorization_study_pattern: str
    problem_solving_study_pattern: str
    visualization_study_pattern: str
    obstacle_study_pattern: str
    new_topic_approach: str
    old_topic_approach: str
    topic_ratio: str
    hours_of_study: str
    hours_of_study_weekends: str
    revision_days: str
    test_days: str
    physicsStartIndex: int
    chemistryStartIndex: int
    mathematicsStartIndex: int
    completed_phy_chapters: List[str]
    completed_chem_chapters: List[str]
    completed_maths_chapters: List[str]
    currentDate: str

# Utility function to remove completed chapters
def remove_completed_chapters(subject_data, completed_chapters):
    subject_data["chapters"] = [
        chapter for chapter in subject_data["chapters"] if chapter["chapter"] not in completed_chapters
    ]
    return subject_data

# Utility function to get chapter at a specific index
def get_data_at_index(json_data, index):
    if 0 <= index < len(json_data['chapters']):
        return json_data['chapters'][index]
    return {}

# Agent to generate a roadmap for a subject
async def generate_subject_roadmap(subject_name, subject_data, study_input):
    user_persona = f"""
    You are generating a JEE roadmap for {subject_name}.
    Student Preferences:
    - Study Pattern: {study_input.overall_study_pattern}
    - Memorization: {study_input.memorization_study_pattern}
    - Problem-Solving: {study_input.problem_solving_study_pattern}
    - Visualization: {study_input.visualization_study_pattern}
    - New Topics: {study_input.new_topic_approach}
    - Old Topics: {study_input.old_topic_approach}
    - Study Hours (Weekdays): {study_input.hours_of_study}
    - Study Hours (Weekends): {study_input.hours_of_study_weekends}
    - Revision Days: {study_input.revision_days}
    - Test Days: {study_input.test_days}
    """

    output_structure = """{
      "schedule": [
        {
          "dayNumber": int,
          "date": YYYY-MM-DD,
          "subjects": [
            {
              "name": "string",
              "tasks": [
                {
                  "ChapterName": "string",
                  "type": "string",
                  "topic": "string",
                  "time": "string"
                }
              ]
            }
          ]
        }
      ]
    }"""

    system_prompt = f"""
    Generate a structured roadmap for {subject_name} using the following data: {subject_data}.
    The roadmap must include Concept Learning, Question Practice, Revision, and Tests.
    Stick to the time allocations and ensure the JSON format follows:
    {output_structure}
    """

    openai.api_key = os.getenv("KEY")
    response = await openai.ChatCompletion.acreate(
        model="gpt-4o-mini",
        messages=[
            {"role": "system", "content": system_prompt},
            {"role": "user", "content": user_persona}
        ]
    )
    return json.loads(response["choices"][0]["message"]["content"])

# API endpoint for roadmap generation
@app.post("/generate_roadmap")
async def generate_roadmap(study_input: StudyInput):
    try:
        # Load JSON data for each subject
        with open('Physics.json', 'r', encoding='utf-8') as file:
            phy = json.load(file)
        with open('Chemistry.json', 'r', encoding='utf-8') as file:
            chem = json.load(file)
        with open('Maths.json', 'r', encoding='utf-8') as file:
            maths = json.load(file)

        # Remove completed chapters
        phy = remove_completed_chapters(phy, study_input.completed_phy_chapters)
        chem = remove_completed_chapters(chem, study_input.completed_chem_chapters)
        maths = remove_completed_chapters(maths, study_input.completed_maths_chapters)

        # Get the chapters at the given index
        phy = get_data_at_index(phy, study_input.physicsStartIndex)
        chem = get_data_at_index(chem, study_input.chemistryStartIndex)
        maths = get_data_at_index(maths, study_input.mathematicsStartIndex)

        # Run agents in parallel
        phy_task = asyncio.create_task(generate_subject_roadmap("Physics", phy, study_input))
        chem_task = asyncio.create_task(generate_subject_roadmap("Chemistry", chem, study_input))
        maths_task = asyncio.create_task(generate_subject_roadmap("Maths", maths, study_input))

        # Collect results
        physics_roadmap, chemistry_roadmap, maths_roadmap = await asyncio.gather(phy_task, chem_task, maths_task)

        # Combine results
        final_roadmap = {
            "Physics": physics_roadmap,
            "Chemistry": chemistry_roadmap,
            "Maths": maths_roadmap
        }

        return json.dumps(final_roadmap, indent=4)
    except Exception as e:
        raise HTTPException(status_code=500, detail=str(e))

# Run FastAPI
if __name__ == "__main__":
    import uvicorn
    uvicorn.run(app, host="0.0.0.0", port=8000)