# D:\GNPL-Private\TimeCraft\backend\app\services\llm_service.py

import os
import json
import httpx
from typing import List
from datetime import date
from dotenv import load_dotenv
from app.models.task import TaskCreateRequest, ScheduleResponse

load_dotenv()

LLM_PROVIDER = os.getenv("LLM_PROVIDER", "OLLAMA")
LLM_API_KEY = os.getenv("LLM_API_KEY")
LLM_API_URL = os.getenv("LLM_API_URL")
HF_MODEL_NAME = os.getenv("HF_MODEL_NAME", "google/gemma-2b-it")
OLLAMA_MODEL_NAME = os.getenv("OLLAMA_MODEL_NAME", "deepseek-r1:1.5b")
OLLAMA_API_URL = os.getenv("OLLAMA_API_URL", "http://localhost:11434/api/generate")

tokenizer = None
model = None

if LLM_PROVIDER == "HUGGINGFACE":
    try:
        from transformers import AutoTokenizer, AutoModelForCausalLM
        import torch
        tokenizer = AutoTokenizer.from_pretrained(HF_MODEL_NAME)
        model = AutoModelForCausalLM.from_pretrained(HF_MODEL_NAME, torch_dtype=torch.float16)
        if torch.cuda.is_available():
            model.to('cuda')
        print(f"[LLM] Hugging Face model '{HF_MODEL_NAME}' loaded.")
    except Exception as e:
        print(f"[ERROR] Failed to load Hugging Face model: {e}")
        raise RuntimeError("Hugging Face model loading failed.")

def infer_priority(task: TaskCreateRequest) -> str:
    desc = task.task_description.lower()
    if any(k in desc for k in ["non-negotiable", "critical", "mandatory", "must", "cannot be moved", "very high priority"]):
        return "high"
    if any(k in desc for k in ["optional", "nice to have", "if possible", "fun"]):
        return "low"
    return "medium"

async def get_scheduled_tasks(tasks: List[TaskCreateRequest]) -> ScheduleResponse:
    today_str = date.today().isoformat()
    prompt_tasks_data = [
        {
            **json.loads(task.model_dump_json()),
            "original_task_id": f"task{index + 1}",
            "priority": infer_priority(task)
        }
        for index, task in enumerate(tasks)
    ]

    system_message = (
        f"You are an AI-powered schedule and task planning assistant. Your job is to create a conflict-free, "
        f"intelligent 30-day schedule from today ({today_str}).\n\n"
        "Parse task details precisely: type, time_settings (start_time, end_time, duration), category, and priority.\n"
        "Avoid time conflicts. If two fixed-time tasks clash, prioritize the 'high' priority one. Try to shift "
        "'medium' priority tasks to nearby times. Skip low priority if no space.\n\n"
        "Respect fixed times first. For flexible tasks, use hints.\n"
        "Return strictly valid JSON:\n"
        "{\n"
        "  \"scheduled_tasks\": [\n"
        "    {\n"
        "      \"task_name\": \"string\",\n"
        "      \"date\": \"YYYY-MM-DD\",\n"
        "      \"start_time\": \"HH:MM\",\n"
        "      \"end_time\": \"HH:MM\",\n"
        "      \"category\": \"personal | work | errands | extra\",\n"
        "      \"original_task_id\": \"string\"\n"
        "    }\n"
        "  ]\n"
        "}"
    )

    user_message = f"Here are the tasks to schedule:\n{json.dumps(prompt_tasks_data, indent=2)}"
    full_prompt = f"{system_message}\n\n{user_message}"
    llm_output_content = ""

    try:
        if LLM_PROVIDER == "GENERIC_API":
            if not LLM_API_URL or not LLM_API_KEY:
                raise ValueError("Missing LLM_API_URL or LLM_API_KEY in .env")

            headers = {
                "Content-Type": "application/json",
                "Authorization": f"Bearer {LLM_API_KEY}"
            }

            payload = {
                "model": "your-preferred-llm-model-name",
                "messages": [
                    {"role": "system", "content": system_message},
                    {"role": "user", "content": user_message}
                ],
                "response_format": {"type": "json_object"}
            }

            async with httpx.AsyncClient(timeout=httpx.Timeout(600.0, connect=60.0)) as client:
                response = await client.post(LLM_API_URL, headers=headers, json=payload)
                response.raise_for_status()
                llm_output_content = response.json().get("choices", [{}])[0].get("message", {}).get("content", "")

        elif LLM_PROVIDER == "HUGGINGFACE":
            import torch
            input_ids = tokenizer(full_prompt, return_tensors="pt").input_ids
            if torch.cuda.is_available():
                input_ids = input_ids.to("cuda")
            outputs = model.generate(input_ids, max_new_tokens=1024, pad_token_id=tokenizer.eos_token_id)
            llm_output_content = tokenizer.decode(outputs[0], skip_special_tokens=True)

        elif LLM_PROVIDER == "OLLAMA":
            payload = {
                "model": OLLAMA_MODEL_NAME,
                "prompt": full_prompt,
                "stream": False,
                "format": "json"
            }

            async with httpx.AsyncClient(timeout=httpx.Timeout(600.0, connect=60.0)) as client:
                print("[LLM] Sending prompt to Ollama...")
                response = await client.post(OLLAMA_API_URL, json=payload)
                response.raise_for_status()
                llm_output_content = response.json().get("response", "")

                print("[LLM] Received response (first 500 chars):")
                print(llm_output_content[:500])

        else:
            raise ValueError(f"Unsupported LLM_PROVIDER: {LLM_PROVIDER}")

        json_start = llm_output_content.find('{')
        json_end = llm_output_content.rfind('}')
        if json_start == -1 or json_end == -1:
            raise ValueError("No valid JSON object found in LLM output.")
        json_string = llm_output_content[json_start:json_end + 1]
        parsed_data = json.loads(json_string)
        return ScheduleResponse(**parsed_data)

    except httpx.HTTPStatusError as e:
        print(f"[ERROR] Ollama returned HTTP {e.response.status_code}: {e.response.text}")
        raise RuntimeError(f"LLM API returned error {e.response.status_code}: {e.response.text}")
    except httpx.RequestError as e:
        print(f"[ERROR] Request to LLM service failed: {e}")
        raise RuntimeError(f"Failed to connect to LLM service. Details: {e}")
    except json.JSONDecodeError as e:
        print(f"[ERROR] JSON decoding failed: {e}")
        print(f"[DEBUG] Raw output: {llm_output_content}")
        raise ValueError("LLM did not return valid JSON.")
    except Exception as e:
        print(f"[ERROR] Unexpected exception: {e}")
        raise RuntimeError(f"Failed to schedule tasks due to an issue with the LLM service. Details: {e}")