import logging
import socket
import subprocess
import psutil
import json
import asyncio
from functools import lru_cache
from ollama import AsyncClient
from modules.tools import available_functions
from modules.system_prompts.system_prompts import SYSTEM_PROMPTS

logging.basicConfig(level=logging.INFO)
logger = logging.getLogger(__name__)

@lru_cache(maxsize=1)
def list_models_cached():
    return list_models()

def list_models() -> list:
    try:
        import requests
        response = requests.get('http://localhost:11434/api/tags')
        response.raise_for_status()
        return [model['name'] for model in response.json().get('models', [])]
    except Exception as e:
        logger.error(f"Model listing failed: {str(e)}")
        return []

def ensure_ollama_running() -> None:
    def is_port_in_use(port):
        with socket.socket(socket.AF_INET, socket.SOCK_STREAM) as s:
            return s.connect_ex(('localhost', port)) == 0

    def free_port(port):
        for proc in psutil.process_iter(['pid', 'name']):
            for conn in proc.net_connections(kind='inet'):
                if conn.laddr.port == port:
                    try:
                        proc.kill()
                        return
                    except Exception as e:
                        logger.error(f"Failed to kill process on port {port}: {e}")

    try:
        import requests
        response = requests.get("http://localhost:11434")
        if response.status_code == 200:
            logger.info("Ollama is already running.")
            return
    except Exception:
        logger.info("Ollama not running. Attempting to start it...")
        if is_port_in_use(11434):
            logger.info("Port 11434 is in use, freeing it...")
            free_port(11434)

    try:
        subprocess.Popen(["ollama", "serve"], stdout=subprocess.DEVNULL, stderr=subprocess.DEVNULL)
        logger.info("Ollama server started.")
    except Exception as e:
        logger.error(f"Failed to start Ollama: {e}")

async def query_ollama(
        prompt: list,
        model: str,
        tools: list,
        context: dict,
        agent_type: str = 'core_llm'
    ) -> dict:
        system_prompt = SYSTEM_PROMPTS.get(agent_type.upper(), SYSTEM_PROMPTS['CORE_LLM'])
        
        formatted_messages = [
            {"role": "system", "content": system_prompt},
            *prompt,
            {"role": "context", "content": str(context)}
        ]

        try:
            client = AsyncClient()
            response = await client.chat(
                model=model,
                messages=formatted_messages,
                tools=tools,
                options={"temperature": 0.4}
            )
            msg = response.message
            assistant_text = msg.content
            logger.info("Assistant: %s", msg.tool_calls)
            if msg.tool_calls:
                tool_call = msg.tool_calls[0]
                fn_name = tool_call.function.name
                raw_args = tool_call.function.arguments

                args = raw_args if isinstance(raw_args, dict) else json.loads(raw_args)
                args = {k: int(v) if k == "max_results" and isinstance(v, str) and v.isdigit() else v for k, v in args.items()}

                if fn_name in available_functions:
                    result = await asyncio.to_thread(available_functions[fn_name], **args)
                    # messages.append({"role": "tool", "name": fn_name, "content": "\n\n".join(result)})
                    mini_sys = {"role": "system",
                                 "content": f"Based on this result:\n {result}, try to answer the original question by user: {prompt[-1]['content']}"}

                    followup = await client.chat(model=model, messages=[mini_sys]
                                                            , options={"temperature": 0.0})
                    return followup.message.content

            return assistant_text
        except Exception as e:
            logger.exception("Error in query_ollama:")
            return f"Ollama Tool Call Error: {str(e)}"