import json
import logging
import pytest
from llmschema.core import generate_response, _extract_json
from llmschema.schema_manager import SchemaManager
from llmschema.exceptions import LLMValidationError
from pydantic import BaseModel
from typing import List, Optional
import ollama

from pydantic import BaseModel, Field

# Configure logging for tests
logging.basicConfig(level=logging.DEBUG)

# Sample Schema for Testing
class ResponseSchema(BaseModel):
    text: str
    confidence: float
    alternatives: Optional[List[str]] = None

# Set the test schema
SchemaManager.set_schema(ResponseSchema)

def test_generate_response_valid_json():
    """Test if generate_response returns a properly structured JSON output using actual Ollama model."""
    result = generate_response("llama3.2:1b", "Say hello")
    
    assert isinstance(result, dict)
    assert "text" in result
    assert "confidence" in result

def test_generate_response_invalid_json_retries(caplog):
    """Test if generate_response retries when LLM returns invalid JSON."""
    caplog.set_level(logging.WARNING)

    try:
        result = generate_response("llama3.2:1b", "Say something weird", max_retries=1)
        assert isinstance(result, dict)  # Should still return a valid dict after retries
    except LLMValidationError:
        assert "Retrying..." in caplog.text

def test_extract_json_valid():
    """Test if _extract_json correctly parses valid JSON."""
    valid_json_str = '{"text": "Hello", "confidence": 0.95}'
    parsed_json = _extract_json(valid_json_str)
    
    assert parsed_json["text"] == "Hello"
    assert parsed_json["confidence"] == 0.95

def test_extract_json_handles_code_blocks():
    """Test if _extract_json properly extracts JSON from markdown code blocks."""
    json_with_code_block = "```json\n{\"text\": \"Hello\", \"confidence\": 0.95}\n```"
    parsed_json = _extract_json(json_with_code_block)
    
    assert parsed_json["text"] == "Hello"
    assert parsed_json["confidence"] == 0.95

def test_extract_json_invalid():
    """Test if _extract_json raises an error for malformed JSON."""
    invalid_json_str = "Hello, this is not JSON"
    with pytest.raises(json.JSONDecodeError):
        _extract_json(invalid_json_str)

def test_generate_response_logs_errors(caplog):
    """Ensure that errors are logged when the model returns bad JSON."""
    caplog.set_level(logging.ERROR)

    try:
        generate_response("llama3.2:1b", "Give me nonsense output")
    except LLMValidationError:
        assert "JSON decoding failed" in caplog.text

def test_generate_response_validates_schema():
    """Ensure the response strictly follows the user-defined schema."""
    result = generate_response("llama3.2:1b", "Provide a structured response")
    assert "text" in result
    assert isinstance(result["text"], str)
    assert "confidence" in result
    assert isinstance(result["confidence"], float)

def test_generate_response_handles_pydantic_and_json_schema():
    """Ensure generate_response can handle both Pydantic-based schemas and standard JSON output."""
    
    class CustomResponseSchema(BaseModel):
        tool: Optional[str] = None
        intent: Optional[str] = None
        receiver: str = Field(..., description="Who receives this response: 'user' or 'action_router'")
        parameters: str = Field(..., description="Parameters for the response (JSON string) - return empty string quoted obj if no params required.")
        response: str = Field(..., description="Final response generated by the LLM")
        actionType: str = Field(..., description="Specify one of: 'A' (Action), 'S' (Search), 'P' (Params Missing), 'R' (Response to User). Must always have a value, cannot be None.")
    
    # Set the new schema
    SchemaManager.set_schema(CustomResponseSchema)
    
    response = generate_response("llama3.2:1b", "Book a flight from Delhi to Bangalore for tomorrow.")
    
    assert isinstance(response, dict)
    assert "tool" in response
    assert "intent" in response
    assert "receiver" in response
    assert "parameters" in response
    assert "response" in response
    assert "actionType" in response
    
    # Ensure response adheres to schema
    validated_response = CustomResponseSchema(**response)
    assert validated_response.receiver in ["user", "action_router"]
