# model.py
import os
import json
from langchain_huggingface import HuggingFaceEndpoint, ChatHuggingFace
from langchain_openai import ChatOpenAI
from langchain_community.document_loaders import PyPDFLoader
from crewai import Agent, Task, Crew, Process

# --- 1. LLM Factory Engine ---
def initialize_llm(provider, selected_model, user_token=None):
    """
    Initializes and returns the appropriate LangChain chat wrapper 
    depending on the chosen provider (Hugging Face or OpenAI).
    """
    if provider == "Hugging Face Serverless (Default)":
        # Resolve token priority: Explicit user token -> Env variable secrets
        token = user_token if user_token else (os.getenv("HUGGINGFACEHUB_API_TOKEN") or os.getenv("HF_TOKEN"))
        if not token:
            raise ValueError("Hugging Face API Token missing! Please provide it in the sidebar or Space Secrets.")
            
        base_endpoint = HuggingFaceEndpoint(
            repo_id=selected_model,
            task="text-generation",
            huggingfacehub_api_token=token,
            max_new_tokens=1500,
            temperature=0.2
        )
        return ChatHuggingFace(llm=base_endpoint)
        
    elif provider == "OpenAI (Optional Override)":
        token = user_token if user_token else os.getenv("OPENAI_API_KEY")
        if not token:
            raise ValueError("OpenAI API Key missing! Please provide it in the sidebar input field.")
            
        return ChatOpenAI(model=selected_model, api_key=token, temperature=0.1)
    
    else:
        raise ValueError(f"Unknown provider type: {provider}")

# --- 2. Advanced Document Context Parser ---
def parse_uploaded_file_content(file_name, file_bytes):
    """
    Extracts raw text strings from multiple formats: .pdf, .ipynb, .py, .txt
    """
    # Write to a temporary file locally so loaders can reference it
    with open(file_name, "wb") as temp_file:
        temp_file.write(file_bytes)
        
    try:
        if file_name.endswith(".pdf"):
            loader = PyPDFLoader(file_name)
            pages = loader.load_and_split()
            extracted_text = "\n".join([page.page_content for page in pages])
            
        elif file_name.endswith(".ipynb"):
            notebook_data = json.loads(file_bytes.decode("utf-8"))
            code_lines = []
            for cell in notebook_data.get("cells", []):
                if cell.get("cell_type") == "code":
                    code_lines.append("".join(cell.get("source", [])))
            extracted_text = "\n# --- Notebook Cell Output ---\n".join(code_lines)
            
        else:
            # Fallback text decoder for standard code scripts (.py, .txt, .md)
            extracted_text = file_bytes.decode("utf-8")
            
        # Context cleaning/wrapping
        return f"\n\n--- REFERENCE DOCUMENT ATTACHED ({file_name}) ---\n{extracted_text}\n"
        
    except Exception as parse_error:
        return f"\n\n[Error parsing document data from {file_name}: {str(parse_error)}]\n"
    finally:
        # Guarantee local filesystem cleanup
        if os.path.exists(file_name):
            os.remove(file_name)

# --- 3. Specialist Agent Configurations Catalog ---
AGENT_PERSONAS = {
    "Personal Motivator": {
        "role": "High-Performance Life Coach",
        "backstory": "An inspiring professional who helps users crush creative blockages, break down bad habits, and build an unstoppable focus mindset.",
        "goal": "Inject intense enthusiasm, structured routines, and motivational insight into the user prompt."
    },
    "Code Reviewer": {
        "role": "Principal Software Engineer & Architect",
        "backstory": "An expert with decades of experience tracking down edge-cases, semantic errors, bad performance bottlenecks, and structural refactoring improvements.",
        "goal": "Thoroughly review structural files or algorithms to provide high-quality code changes."
    },
    "Technical Researcher": {
        "role": "Lead Data Analytics Scholar",
        "backstory": "A detail-oriented analyst who excels at parsing complex raw documents, compiling literature overviews, and summarizing trends accurately.",
        "goal": "Identify hidden insights from context and write clear technical documentation summaries."
    },
    "Physics Formula Explainer": {
        "role": "Theoretical Physics Professor",
        "backstory": "An educator who makes complex physical theories, mathematical proofs, and equations easy to understand by using intuitive analogies and practical real-world applications.",
        "goal": "Deconstruct complex physics models and equations step-by-step into clear, understandable explanations."
    },
    "General Purpose Assistant": {
        "role": "Contextual Problem Solver",
        "backstory": "A versatile technical assistant skilled at answering cross-functional development questions clearly.",
        "goal": "Fulfill the user request precisely while adjusting to the provided documentation constraints."
    }
}

# --- 4. Main Multi-Agent Execution Orchestrator ---
def run_agent_pipeline(provider, model_name, token, persona_key, user_prompt, context_text):
    """
    Coordinates the selected Agent Persona alongside a secondary Quality Assurance 
    Auditor to review and polish the output before displaying it.
    """
    # 1. Instantiate the chosen LLM runtime
    llm = initialize_llm(provider, model_name, user_token=token)
    
    # 2. Extract Persona Definitions
    persona = AGENT_PERSONAS[persona_key]
    
    # 3. Assemble Specialist Agent
    specialist_agent = Agent(
        role=persona["role"],
        backstory=persona["backstory"],
        goal=persona["goal"],
        llm=llm,
        verbose=True,
        allow_delegation=False
    )
    
    # 4. Assemble Independent QA Critic Agent
    auditor_agent = Agent(
        role="Senior QA & Compliance Auditor",
        backstory="An unyielding editor focused on correcting hallucinations, logic formatting errors, typos, and completeness issues.",
        goal="Critique the primary specialist's generated draft response and output a polished, final version.",
        llm=llm,
        verbose=True
    )
    
    # 5. Define Task Blueprints
    complete_task_description = f"{user_prompt}\n\n{context_text}"
    
    generation_task = Task(
        description=f"Address the user requirements thoroughly using your specific persona parameters.\nRequirements:\n{complete_task_description}",
        agent=specialist_agent,
        expected_output="A rich, highly comprehensive response matching your assigned persona specialty."
    )
    
    review_task = Task(
        description="Review the generated response for accuracy, clarity, and formatting. Correct any typos or logical errors, and format it cleanly in Markdown.",
        agent=auditor_agent,
        context=[generation_task],
        expected_output="The final, polished markdown response, safe for production delivery."
    )
    
    # 6. Execute via Sequential Flow Pipeline
    crew_system = Crew(
        agents=[specialist_agent, auditor_agent],
        tasks=[generation_task, review_task],
        process=Process.sequential,
        verbose=True
    )
    
    return crew_system.kickoff()